import streamlit as st
import plotly.express as px
import pandas as pd

def display_dashboard(df):
    st.subheader("📊 System Summary")
    col1, col2, col3 = st.columns(3)
    col1.metric("Total Poles", df.shape[0])
    col2.metric("🚨 Red Alerts", df[df['AlertLevel'] == "Red"].shape[0])
    col3.metric("⚡ Power Issues", df[df['PowerSufficient'] == "No"].shape[0])

def display_charts(df):
    st.subheader("⚙️ Energy Generation Trends")
    st.bar_chart(df.set_index("PoleID")[["SolarGen(kWh)", "WindGen(kWh)"]])
    st.subheader("📉 Tilt vs Vibration")
    st.scatter_chart(df.rename(columns={"Tilt(°)": "Tilt", "Vibration(g)": "Vibration"}).set_index("PoleID")[["Tilt", "Vibration"]])

def display_heatmap(df):
    # Map AlertLevel to numeric values for heatmap intensity
    alert_map = {"Green": 0, "Yellow": 1, "Red": 2}
    df["AlertValue"] = df["AlertLevel"].map(alert_map)
    
    # Create a pivot table for heatmap (single row for all poles)
    pivot_df = df[["PoleID", "AlertValue"]].set_index("PoleID").T
    
    # Create heatmap using Plotly
    fig = px.imshow(
        pivot_df,
        color_continuous_scale=["green", "yellow", "red"],
        zmin=0,
        zmax=2,
        labels=dict(color="Alert Level"),
        title="Pole Alert Heatmap",
        height=300
    )
    fig.update_layout(
        xaxis_title="Pole ID",
        yaxis_title="",
        yaxis_showticklabels=False,
        coloraxis_colorbar=dict(
            tickvals=[0, 1, 2],
            ticktext=["Green", "Yellow", "Red"]
        )
    )
    st.plotly_chart(fig, use_container_width=True)