consider incident type as string and apply colors

src/streamlit_app.py

@@ -25,7 +25,7 @@ INCIDENT_COLOR_MAP = {
     "2": ("#eab308", "Rupt/Exp"),
     "3": ("#3b82f6", "EMS"),
     "4": ("#22c55e", "Hazardous"),
-    "S": ("#a855f7", "Public Assist"),
+    "5": ("#a855f7", "Public Assist"),
     "6": ("#06b6d4", "Good Intent"),
     "7": ("#9ca3af", "False Alarms"),
     "8": ("#84cc16", "Weather"),
@@ -60,10 +60,17 @@ INCIDENT_NAME_COLOR_MAP = {
     "special type": ("#c0c0c0", "Special Type"),
 }
 
+INCIDENT_COL_NAMES = {"incidenttype", "incident_type", "incidentclassification", "incident_classification", "incident_category", "incidentcategory"}
+
+def _detect_incident_col(df: pd.DataFrame) -> str | None:
+    for col in df.columns:
+        if col.strip().lower() in INCIDENT_COL_NAMES:
+            return col
+    return None
 
 def _incident_label_and_color(value) -> tuple[str, str]:
     """Return (display_label, hex_color) for a raw incidenttype value."""
-    if value is None or (isinstance(value, float) and pd.isna(value)):
+    if value is None or value == "nan" or (isinstance(value, float) and pd.isna(value)):
         return INCIDENT_NULL_COLOR[1], INCIDENT_NULL_COLOR[0]
     s = str(value).strip()
     if s == "":
@@ -90,13 +97,15 @@ def _incident_label_and_color(value) -> tuple[str, str]:
 
 
 def _add_incident_category(df: pd.DataFrame, col: str) -> pd.DataFrame:
-    """Add _incident_label and _incident_color helper columns (non-destructive)."""
     df = df.copy()
     if col not in df.columns or df.empty:
         return df
-    labels, colors = zip(*df[col].map(_incident_label_and_color))
-    df["_incident_label"] = labels
-    df["_incident_color"] = colors
+    df = df.reset_index(drop=True)
+    # Force plain string series — eliminates Categorical/MultiIndex issues
+    plain = df[col].astype(str).where(df[col].notna(), other=None)
+    mapped = plain.map(_incident_label_and_color)
+    df["_incident_label"] = mapped.apply(lambda x: x[0])
+    df["_incident_color"] = mapped.apply(lambda x: x[1])
     return df
 
 
@@ -423,25 +432,33 @@ def _guess_chart_type(df: pd.DataFrame) -> str:
 
 def render_chart(
     df: pd.DataFrame,
-    incident_col: str | None = None,
+    incident_cols: list[str] | None = None,
     chart_type_raw: str | None = None,
     key_prefix: str = "chart",
 ):
     if df.empty:
         st.info("No data to chart.")
         return
+    
+    df = df.copy()
+    for c in df.columns:
+        if hasattr(df[c], "cat"):
+            df[c] = df[c].astype(str).replace("nan", None)
+
 
     chart_type = _normalise_chart_type(chart_type_raw) or _guess_chart_type(df)
 
-    df_plot = _add_incident_category(df, incident_col) if incident_col else df.copy()
+    incident_cols = incident_cols or []
+    active_incident_col: str | None = None
+    df_plot = df.copy()
 
     cols = list(df.columns)
     numeric = df.select_dtypes(include="number").columns.tolist()
     cat = df.select_dtypes(exclude="number").columns.tolist()
 
     default_x = (
-        incident_col
-        if incident_col and incident_col in cols
+        incident_cols[0]
+        if incident_cols and incident_cols[0] in cols
         else (cat[0] if cat else cols[0])
     )
     default_y = numeric[0] if numeric else (cols[1] if len(cols) > 1 else cols[0])
@@ -508,16 +525,11 @@ def render_chart(
 
     # ── Incident color mapping ─────────────────────────
     incident_color_map = None
-    use_incident_x = False
-    if (
-        incident_col
-        and "_incident_label" in df_plot.columns
-        and "_incident_color" in df_plot.columns
-    ):
+    if active_incident_col and "_incident_label" in df_plot.columns:
         incident_color_map = dict(
             zip(df_plot["_incident_label"], df_plot["_incident_color"])
         )
-        use_incident_x = x_col == incident_col
+
 
     # ── Build chart ────────────────────────────────────
     fig = None
@@ -526,6 +538,19 @@ def render_chart(
         # ───── BAR ─────
         if chart_type == "bar":
 
+            active_incident_col = x_col if x_col in incident_cols else None
+            df_plot = (
+                _add_incident_category(df, active_incident_col)
+                if active_incident_col
+                else df.copy()
+            )
+
+            incident_color_map = (
+                dict(zip(df_plot["_incident_label"].tolist(), df_plot["_incident_color"].tolist()))
+                if active_incident_col and "_incident_label" in df_plot.columns
+                else None
+            )
+
             if color_col:
                 bar_kwargs = dict(
                     x=x_col,
@@ -534,13 +559,13 @@ def render_chart(
                     barmode="group",
                     template=PLOTLY_TEMPLATE,
                 )
-                if incident_color_map and color_col == incident_col:
-                    bar_kwargs["x"] = "_incident_label" if use_incident_x else x_col
+                if incident_color_map and color_col in incident_cols:
+                    bar_kwargs["x"] = "_incident_label" if active_incident_col is not None else x_col
                     bar_kwargs["color"] = "_incident_label"
                     bar_kwargs["color_discrete_map"] = incident_color_map
                 fig = px.bar(df_plot, **bar_kwargs)
 
-            elif incident_color_map and use_incident_x:
+            elif incident_color_map and active_incident_col is not None:
                 fig = px.bar(
                     df_plot,
                     x="_incident_label",
@@ -561,6 +586,19 @@ def render_chart(
         # ───── LINE ─────
         elif chart_type == "line":
 
+            active_incident_col = x_col if x_col in incident_cols else None
+            df_plot = (
+                _add_incident_category(df, active_incident_col)
+                if active_incident_col
+                else df.copy()
+            )
+
+            incident_color_map = (
+                dict(zip(df_plot["_incident_label"].tolist(), df_plot["_incident_color"].tolist()))
+                if active_incident_col and "_incident_label" in df_plot.columns
+                else None
+            )
+
             if color_col:
                 line_kwargs = dict(
                     x=x_col,
@@ -569,13 +607,13 @@ def render_chart(
                     markers=True,
                     template=PLOTLY_TEMPLATE,
                 )
-                if incident_color_map and color_col == incident_col:
-                    line_kwargs["x"] = "_incident_label" if use_incident_x else x_col
+                if incident_color_map and color_col in incident_cols:
+                    line_kwargs["x"] = "_incident_label" if active_incident_col is not None else x_col
                     line_kwargs["color"] = "_incident_label"
                     line_kwargs["color_discrete_map"] = incident_color_map
                 fig = px.line(df_plot, **line_kwargs)
 
-            elif incident_color_map and use_incident_x:
+            elif incident_color_map and active_incident_col is not None:
                 fig = px.line(
                     df_plot,
                     x="_incident_label",
@@ -598,7 +636,21 @@ def render_chart(
         # ───── PIE ─────
         elif chart_type == "pie":
 
-            if incident_color_map and use_incident_x:
+            active_incident_col = x_col if x_col in incident_cols else None
+            df_plot = (
+                _add_incident_category(df, active_incident_col)
+                if active_incident_col
+                else df.copy()
+            )
+
+            incident_color_map = (
+                dict(zip(df_plot["_incident_label"].tolist(), df_plot["_incident_color"].tolist()))
+                if active_incident_col and "_incident_label" in df_plot.columns
+                else None
+            )
+
+
+            if incident_color_map and active_incident_col is not None:
                 fig = px.pie(
                     df_plot,
                     names="_incident_label",
@@ -784,7 +836,7 @@ def render_message(msg):
     # ── Multi-view data panel ──────────────────────────────────────────────
     data = msg.get("data", [])
     chart_hint = msg.get("best_suitable_chart")  # e.g. "bar", "line", "pie"
-    incident_col = msg.get("incident_col")
+
 
     if data and len(data) > 0 and msg.get("status") == "success":
         df = pd.DataFrame(data)
@@ -823,7 +875,17 @@ def render_message(msg):
             # key_prefix is unique per message (uses timestamp) so each
             # chart's column selectors have independent Streamlit widget keys.
             chart_key = f"chart_{msg.get('ts', 'x').replace(':', '_')}"
-            render_chart(df, incident_col, chart_hint, key_prefix=chart_key)
+            incident_cols = []
+            if msg.get("incident_col"):
+                incident_cols.append(msg["incident_col"])
+            detected = _detect_incident_col(df)
+            if detected and detected not in incident_cols:
+                incident_cols.append(detected)
+
+            for col in incident_cols:
+                if col in df.columns:
+                    df[col] = df[col].astype("category")
+            render_chart(df, incident_cols, chart_hint, key_prefix=chart_key)
 
     elif msg.get("status") == "success" and not data:
         st.markdown(