utils/visualizations.py

"""Plotly visualization helpers for the PFAS-SBEAD dashboard."""

from __future__ import annotations

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots


COLOR_SCHEME = px.colors.qualitative.Set2
TEMPLATE = "plotly_white"


def degradation_scatter(df: pd.DataFrame) -> go.Figure:
    """Scatter: PFAS degradation vs voltage, colored by current density."""
    fig = px.scatter(
        df,
        x="voltage_V",
        y="PFAS_degradation_pct",
        color="current_density_A_m2",
        size="HRT_days",
        hover_data=["OLR_kg_m3_d", "pH", "AI_score"],
        title="PFAS Degradation vs Applied Voltage",
        labels={
            "voltage_V": "Voltage (V)",
            "PFAS_degradation_pct": "PFAS Degradation (%)",
            "current_density_A_m2": "Current Density (A/m²)",
            "HRT_days": "HRT (days)",
        },
        color_continuous_scale="Viridis",
        template=TEMPLATE,
    )
    fig.update_layout(height=450)
    return fig


def ai_score_distribution(df: pd.DataFrame) -> go.Figure:
    """Histogram of AI optimization scores."""
    fig = px.histogram(
        df,
        x="AI_score",
        nbins=25,
        title="AI Optimization Score Distribution",
        labels={"AI_score": "AI Score"},
        color_discrete_sequence=[COLOR_SCHEME[0]],
        template=TEMPLATE,
    )
    fig.update_layout(height=350)
    return fig


def mass_balance_sunburst(mb_df: pd.DataFrame) -> go.Figure:
    """Average mass balance breakdown as a pie chart."""
    avg = mb_df[
        ["remaining_in_water_ug_L", "adsorbed_sludge_ug_L",
         "adsorbed_electrode_ug_L", "short_chain_products_ug_L", "mineralized_PFAS_ug_L"]
    ].mean()
    labels = ["Remaining in Water", "Adsorbed on Sludge", "Adsorbed on Electrode",
              "Short-Chain Products", "Mineralized (Degraded)"]
    fig = go.Figure(data=[go.Pie(
        labels=labels,
        values=avg.values,
        hole=0.4,
        marker_colors=COLOR_SCHEME[:5],
    )])
    fig.update_layout(
        title="Average PFAS Mass Balance Distribution",
        template=TEMPLATE,
        height=400,
    )
    return fig


def feature_importance_bar(imp_df: pd.DataFrame) -> go.Figure:
    """Horizontal bar chart of feature importances."""
    fig = px.bar(
        imp_df.head(10),
        x="importance",
        y="feature",
        orientation="h",
        title="Top Feature Importances (SHAP-proxy)",
        labels={"importance": "Importance", "feature": ""},
        color="importance",
        color_continuous_scale="Blues",
        template=TEMPLATE,
    )
    fig.update_layout(height=400, showlegend=False)
    return fig


def degradation_heatmap(df: pd.DataFrame) -> go.Figure:
    """Heatmap: voltage vs OLR bins, showing mean degradation."""
    df_copy = df.copy()
    df_copy["OLR_bin"] = pd.cut(df_copy["OLR_kg_m3_d"], bins=6)
    df_copy["voltage_bin"] = pd.cut(df_copy["voltage_V"], bins=6)
    pivot = df_copy.pivot_table(
        values="PFAS_degradation_pct",
        index="voltage_bin",
        columns="OLR_bin",
        aggfunc="mean",
    )
    fig = px.imshow(
        pivot.values,
        x=[str(c) for c in pivot.columns],
        y=[str(i) for i in pivot.index],
        color_continuous_scale="YlOrRd",
        title="PFAS Degradation Heatmap: Voltage × OLR",
        labels={"x": "OLR Bin (kg/m³/d)", "y": "Voltage Bin (V)", "color": "Degradation (%)"},
        template=TEMPLATE,
    )
    fig.update_layout(height=420)
    return fig


def dual_axis_performance(df: pd.DataFrame) -> go.Figure:
    """Dual-axis: degradation and fluoride release vs experiment."""
    fig = make_subplots(specs=[[{"secondary_y": True}]])
    fig.add_trace(
        go.Scatter(
            x=df["experiment_id"],
            y=df["PFAS_degradation_pct"],
            mode="lines+markers",
            name="PFAS Degradation (%)",
            marker=dict(size=5, color=COLOR_SCHEME[0]),
        ),
        secondary_y=False,
    )
    fig.add_trace(
        go.Scatter(
            x=df["experiment_id"],
            y=df["fluoride_release_mg_L"],
            mode="lines+markers",
            name="Fluoride Release (mg/L)",
            marker=dict(size=5, color=COLOR_SCHEME[1]),
        ),
        secondary_y=True,
    )
    fig.update_layout(
        title="Degradation & Fluoride Release Across Experiments",
        template=TEMPLATE,
        height=400,
        legend=dict(orientation="h", yanchor="bottom", y=1.02),
    )
    fig.update_yaxes(title_text="PFAS Degradation (%)", secondary_y=False)
    fig.update_yaxes(title_text="Fluoride Release (mg/L)", secondary_y=True)
    return fig


def stability_radar(df: pd.DataFrame) -> go.Figure:
    """Radar chart of average stability indicators."""
    cols = ["pH_drop", "current_instability_index"]
    vfa_norm = df["VFA_accumulation_mg_L"] / df["VFA_accumulation_mg_L"].max()
    orp_norm = df["ORP_drift_mV"].abs() / df["ORP_drift_mV"].abs().max()

    values = [
        df["pH_drop"].mean() / 1.5,
        vfa_norm.mean(),
        orp_norm.mean(),
        df["current_instability_index"].mean() / 0.5,
    ]
    categories = ["pH Drop", "VFA Accumulation", "ORP Drift", "Current Instability"]
    values.append(values[0])
    categories.append(categories[0])

    fig = go.Figure(data=go.Scatterpolar(
        r=values,
        theta=categories,
        fill="toself",
        fillcolor="rgba(255, 99, 71, 0.2)",
        line_color="tomato",
    ))
    fig.update_layout(
        polar=dict(radialaxis=dict(visible=True, range=[0, 1])),
        title="Reactor Stability Indicators (Normalized)",
        template=TEMPLATE,
        height=400,
    )
    return fig


def sensitivity_bar(sens_df: pd.DataFrame) -> go.Figure:
    """Bar chart of sensitivity analysis."""
    fig = px.bar(
        sens_df,
        x="feature",
        y="correlation_with_AI_score",
        color="correlation_with_AI_score",
        color_continuous_scale="RdYlGn",
        title="Sensitivity Analysis: Feature Correlation with AI Score",
        labels={"feature": "", "correlation_with_AI_score": "Correlation"},
        template=TEMPLATE,
    )
    fig.update_layout(height=400, xaxis_tickangle=-45)
    return fig


def energy_vs_degradation(df: pd.DataFrame) -> go.Figure:
    """Scatter: energy input vs degradation with instability flag."""
    fig = px.scatter(
        df,
        x="energy_input_kWh_d",
        y="PFAS_degradation_pct",
        color="instability_flag",
        symbol="instability_flag",
        title="Energy Input vs PFAS Degradation (Instability Highlighted)",
        labels={
            "energy_input_kWh_d": "Energy Input (kWh/d)",
            "PFAS_degradation_pct": "PFAS Degradation (%)",
            "instability_flag": "Instability",
        },
        color_discrete_map={0: COLOR_SCHEME[0], 1: "red"},
        template=TEMPLATE,
    )
    fig.update_layout(height=400)
    return fig


def optimization_pareto(df: pd.DataFrame) -> go.Figure:
    """Pareto front: degradation vs energy showing trade-off."""
    fig = px.scatter(
        df,
        x="energy_input_kWh_d",
        y="PFAS_degradation_pct",
        color="AI_score",
        size="fluoride_release_mg_L",
        hover_data=["voltage_V", "HRT_days", "OLR_kg_m3_d"],
        title="Optimization Landscape: Degradation vs Energy Trade-off",
        labels={
            "energy_input_kWh_d": "Energy Input (kWh/d)",
            "PFAS_degradation_pct": "PFAS Degradation (%)",
            "AI_score": "AI Score",
        },
        color_continuous_scale="Plasma",
        template=TEMPLATE,
    )
    fig.update_layout(height=450)
    return fig