Update modules/visuals.py

modules/visuals.py

@@ -1,4 +1,6 @@
 import streamlit as st
+import matplotlib.pyplot as plt
+import seaborn as sns
 
 def display_dashboard(df):
     st.subheader("📊 System Summary")
@@ -10,5 +12,20 @@ def display_dashboard(df):
 def display_charts(df):
     st.subheader("⚙️ Energy Generation Trends")
     st.bar_chart(df.set_index("Pole ID")[["Solar Gen (kWh)", "Wind Gen (kWh)"]])
+    
     st.subheader("📉 Tilt vs Vibration")
-    st.scatter_chart(df.rename(columns={"Tilt (°)": "Tilt", "Vibration (g)": "Vibration"}).set_index("Pole ID")[["Tilt", "Vibration"]])
+    scatter_df = df.rename(columns={"Tilt (°)": "Tilt", "Vibration (g)": "Vibration"})
+    st.scatter_chart(scatter_df.set_index("Pole ID")[["Tilt", "Vibration"]])
+
+def display_heatmap(df):
+    st.subheader("🌡️ Site-wise Pole Alert Heatmap")
+    
+    alert_map = {"Green": 0, "Yellow": 1, "Red": 2}
+    df["Alert Code"] = df["Alert Level"].map(alert_map)
+
+    pivot = df.pivot(index="Site", columns="Pole ID", values="Alert Code")
+
+    fig, ax = plt.subplots(figsize=(14, 4))
+    sns.heatmap(pivot, annot=True, cmap="YlOrRd", linewidths=0.5, cbar_kws={'label': 'Alert Level (0=Green, 2=Red)'}, ax=ax)
+    st.pyplot(fig)
+