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

def display_dashboard(df):
    st.subheader("📊 System Summary")
    col1, col2, col3, col4 = st.columns(4)
    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])
    col4.metric("📍 Locations", len(df['Location'].unique()))

def display_charts(df):
    st.subheader("⚙️ Energy Generation Trends")
    st.bar_chart(df.groupby("Location")[["SolarGen(kWh)", "WindGen(kWh)"]].sum())
    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 = df.copy()  # Avoid modifying the original DataFrame
    df["AlertValue"] = df["AlertLevel"].map(alert_map)
    
    # Pivot table for heatmap values (AlertValue)
    pivot_df = df.pivot_table(index="Location", columns="PoleID", values="AlertValue", fill_value=0)
    
    # Create hover text DataFrame
    hover_data = df[["Location", "PoleID", "AlertLevel", "Anomalies"]].copy()
    hover_text = pivot_df.copy().astype(str)
    for loc in pivot_df.index:
        for pole in pivot_df.columns:
            # Find matching row in hover_data
            match = hover_data[(hover_data["Location"] == loc) & (hover_data["PoleID"] == pole)]
            if not match.empty:
                alert = match["AlertLevel"].iloc[0]
                anomalies = match["Anomalies"].iloc[0]
                hover_text.loc[loc, pole] = f"Pole: {pole}<br>Alert: {alert}<br>Anomalies: {anomalies}"
            else:
                hover_text.loc[loc, pole] = f"Pole: {pole}<br>Alert: None<br>Anomalies: None"
    
    # Create heatmap using Plotly
    fig = px.imshow(
        pivot_df,
        color_continuous_scale=["green", "yellow", "red"],
        zmin=0,
        zmax=2,
        title="Pole Alert Heatmap by Location",
        text_auto=False,
        height=500
    )
    fig.update_traces(hovertemplate="%{customdata}<br>%{x}<br>%{y}", customdata=hover_text)
    fig.update_layout(
        xaxis_title="Pole ID",
        yaxis_title="Location",
        coloraxis_colorbar=dict(
            tickvals=[0, 1, 2],
            ticktext=["Green", "Yellow", "Red"]
        )
    )
    st.plotly_chart(fig, use_container_width=True)