"""
Closed-loop PFAS-SBEAD Optimization Pipeline — Streamlit Application.
AI-driven optimization for PFAS degradation using Sidestream Bioelectrochemical
Anaerobic Digestion (SBEAD) reactor system.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import streamlit as st

from utils.calculations import (
    FEATURE_COLUMNS,
    AI_SCORE_WEIGHTS,
    bayesian_next_recommendation,
    compute_shap_importance,
    sensitivity_analysis,
    train_degradation_model,
    train_fluoride_model,
    train_instability_detector,
    train_short_chain_classifier,
)
from utils.data_generator import generate_mass_balance_data, generate_sbead_dataset
from utils.visualizations import (
    ai_score_distribution,
    degradation_heatmap,
    degradation_scatter,
    dual_axis_performance,
    energy_vs_degradation,
    feature_importance_bar,
    mass_balance_sunburst,
    optimization_pareto,
    sensitivity_bar,
    stability_radar,
)

st.set_page_config(
    page_title="PFAS-SBEAD AI Pipeline",
    page_icon="⚗️",
    layout="wide",
    initial_sidebar_state="expanded",
)

DATA_PATH = Path(__file__).parent / "data"


@st.cache_data
def load_data() -> tuple[pd.DataFrame, pd.DataFrame]:
    exp_path = DATA_PATH / "sbead_experiments.csv"
    mb_path = DATA_PATH / "mass_balance.csv"
    if exp_path.exists() and mb_path.exists():
        df = pd.read_csv(exp_path)
        mb = pd.read_csv(mb_path)
    else:
        df = generate_sbead_dataset(120)
        mb = generate_mass_balance_data(df)
    return df, mb


@st.cache_resource
def train_models(df: pd.DataFrame):
    deg_model, deg_r2 = train_degradation_model(df)
    flu_model, flu_r2 = train_fluoride_model(df)
    sc_model, sc_acc = train_short_chain_classifier(df)
    iso_model = train_instability_detector(df)
    return {
        "degradation": (deg_model, deg_r2),
        "fluoride": (flu_model, flu_r2),
        "short_chain": (sc_model, sc_acc),
        "instability": iso_model,
    }


def render_sidebar(df: pd.DataFrame) -> dict:
    """Sidebar filters and controls."""
    st.sidebar.title("⚗️ PFAS-SBEAD Pipeline")
    st.sidebar.markdown("---")
    st.sidebar.subheader("Experiment Filters")

    olr_range = st.sidebar.slider(
        "OLR (kg/m³/d)", 1.0, 6.0, (1.0, 6.0), 0.1
    )
    voltage_range = st.sidebar.slider(
        "Voltage (V)", 0.2, 1.2, (0.2, 1.2), 0.05
    )
    hrt_range = st.sidebar.slider(
        "HRT (days)", 10.0, 30.0, (10.0, 30.0), 1.0
    )
    ph_range = st.sidebar.slider(
        "pH", 6.5, 8.0, (6.5, 8.0), 0.1
    )

    st.sidebar.markdown("---")
    st.sidebar.subheader("AI Objective Weights")
    w_deg = st.sidebar.slider("PFAS degradation weight", 0.0, 1.0, 0.40, 0.05)
    w_flu = st.sidebar.slider("Fluoride release weight", 0.0, 1.0, 0.30, 0.05)
    w_sc = st.sidebar.slider("Short-chain penalty", 0.0, 0.5, 0.15, 0.05)
    w_en = st.sidebar.slider("Energy penalty", 0.0, 0.5, 0.10, 0.05)
    w_inst = st.sidebar.slider("Instability penalty", 0.0, 0.5, 0.05, 0.05)

    mask = (
        (df["OLR_kg_m3_d"] >= olr_range[0])
        & (df["OLR_kg_m3_d"] <= olr_range[1])
        & (df["voltage_V"] >= voltage_range[0])
        & (df["voltage_V"] <= voltage_range[1])
        & (df["HRT_days"] >= hrt_range[0])
        & (df["HRT_days"] <= hrt_range[1])
        & (df["pH"] >= ph_range[0])
        & (df["pH"] <= ph_range[1])
    )

    return {
        "mask": mask,
        "weights": {"degradation": w_deg, "fluoride": w_flu, "short_chain": w_sc, "energy": w_en, "instability": w_inst},
    }


def render_kpi_header(df: pd.DataFrame) -> None:
    """Top KPI metric cards."""
    col1, col2, col3, col4, col5 = st.columns(5)
    col1.metric(
        "Avg Degradation",
        f"{df['PFAS_degradation_pct'].mean():.1f}%",
        delta=f"max {df['PFAS_degradation_pct'].max():.1f}%",
    )
    col2.metric(
        "Avg Fluoride Release",
        f"{df['fluoride_release_mg_L'].mean():.1f} mg/L",
    )
    col3.metric(
        "Avg AI Score",
        f"{df['AI_score'].mean():.3f}",
        delta=f"best {df['AI_score'].max():.3f}",
    )
    col4.metric(
        "Unstable Runs",
        f"{df['instability_flag'].sum()}/{len(df)}",
    )
    col5.metric(
        "Avg Energy",
        f"{df['energy_input_kWh_d'].mean():.3f} kWh/d",
    )


def render_overview_tab(df: pd.DataFrame, mb_df: pd.DataFrame) -> None:
    """Overview tab with main charts."""
    st.subheader("Experiment Overview")

    c1, c2 = st.columns(2)
    with c1:
        st.plotly_chart(degradation_scatter(df), use_container_width=True)
    with c2:
        st.plotly_chart(dual_axis_performance(df), use_container_width=True)

    c3, c4 = st.columns(2)
    with c3:
        st.plotly_chart(ai_score_distribution(df), use_container_width=True, key="overview_ai_dist")
    with c4:
        st.plotly_chart(mass_balance_sunburst(mb_df), use_container_width=True, key="overview_mass_balance")

    st.subheader("PFAS Degradation Heatmap")
    st.plotly_chart(degradation_heatmap(df), use_container_width=True)


def render_ai_models_tab(df: pd.DataFrame, models: dict) -> None:
    """AI Models tab showing training results and predictions."""
    st.subheader("AI Model Performance")

    deg_model, deg_r2 = models["degradation"]
    flu_model, flu_r2 = models["fluoride"]
    sc_model, sc_acc = models["short_chain"]

    mc1, mc2, mc3 = st.columns(3)
    mc1.metric("XGBoost Degradation (R²)", f"{deg_r2:.3f}")
    mc2.metric("RF Fluoride Release (R²)", f"{flu_r2:.3f}")
    mc3.metric("Short-Chain Classifier (Acc)", f"{sc_acc:.3f}")

    st.markdown("---")
    st.subheader("Feature Importance (PFAS Degradation Model)")
    imp_df = compute_shap_importance(deg_model, df)
    st.plotly_chart(feature_importance_bar(imp_df), use_container_width=True)

    st.subheader("Sensitivity Analysis")
    sens_df = sensitivity_analysis(df)
    st.plotly_chart(sensitivity_bar(sens_df), use_container_width=True)

    st.subheader("Instability Detection")
    iso_model = models["instability"]
    stability_cols = ["pH_drop", "VFA_accumulation_mg_L", "ORP_drift_mV", "current_instability_index"]
    preds = iso_model.predict(df[stability_cols])
    n_anomalies = int((preds == -1).sum())
    st.info(f"Isolation Forest detected **{n_anomalies}** potentially unstable experiments out of {len(df)}.")
    st.plotly_chart(stability_radar(df), use_container_width=True, key="models_stability_radar")


def render_optimization_tab(df: pd.DataFrame) -> None:
    """Optimization and Bayesian recommendation tab."""
    st.subheader("Closed-Loop Optimization")
    st.markdown(
        "The pipeline uses Bayesian-inspired optimization to recommend the next best "
        "reactor condition, balancing exploitation (high-performing conditions) with "
        "exploration (uncertain regions)."
    )

    c1, c2 = st.columns(2)
    with c1:
        st.plotly_chart(optimization_pareto(df), use_container_width=True)
    with c2:
        st.plotly_chart(energy_vs_degradation(df), use_container_width=True)

    st.markdown("---")
    st.subheader("Next Experiment Recommendation")
    rec = bayesian_next_recommendation(df)

    rc1, rc2, rc3, rc4 = st.columns(4)
    rc1.metric("Predicted Degradation", f"{rec['predicted_degradation_pct']:.1f}%")
    rc2.metric("Predicted Fluoride", f"{rec['predicted_fluoride_release']:.1f} mg/L")
    rc3.metric("Expected AI Score", f"{rec['expected_ai_score']:.3f}")
    rc4.metric("Confidence", f"{rec['confidence']:.0%}")

    st.markdown("##### Recommended Operating Conditions")
    rec_display = {k: v for k, v in rec.items() if k in FEATURE_COLUMNS}
    rec_df = pd.DataFrame([rec_display])
    st.dataframe(rec_df.round(3), use_container_width=True, hide_index=True)

    with st.expander("Optimization Strategy"):
        st.markdown("""
**AI Score Function:**

`AI Score = 0.40 × PFAS degradation + 0.30 × fluoride release - 0.15 × short-chain risk - 0.10 × energy input - 0.05 × instability`

**Objective:** Maximize PFAS degradation (>40%) and fluoride release while minimizing
short-chain PFAS accumulation, energy input, and reactor instability.

**Method:** The optimizer identifies top-performing experiments, computes mean conditions,
and recommends the next trial from the region of highest expected improvement.
        """)


def render_mass_balance_tab(df: pd.DataFrame, mb_df: pd.DataFrame) -> None:
    """Mass balance analysis tab."""
    st.subheader("PFAS Mass Balance")
    st.markdown(
        "PFAS degradation must be distinguished from adsorption. The mass balance is:\n\n"
        "`Initial PFAS = Remaining in water + Adsorbed on sludge + Adsorbed on electrode "
        "+ Short-chain products + Mineralized PFAS`"
    )

    st.plotly_chart(mass_balance_sunburst(mb_df), use_container_width=True, key="tab_mass_balance_sunburst")

    st.subheader("Mass Balance Details")
    st.dataframe(
        mb_df.round(2),
        use_container_width=True,
        hide_index=True,
        height=400,
    )

    st.subheader("Mass Balance Closure")
    fig_closure = px.histogram(
        mb_df,
        x="mass_balance_closure_pct",
        nbins=20,
        title="Mass Balance Closure Distribution",
        labels={"mass_balance_closure_pct": "Closure (%)"},
        template="plotly_white",
    )
    st.plotly_chart(fig_closure, use_container_width=True)


def render_reactor_stability_tab(df: pd.DataFrame) -> None:
    """Reactor stability monitoring tab."""
    st.subheader("Reactor Stability Monitoring")

    c1, c2 = st.columns(2)
    with c1:
        fig_ph = px.scatter(
            df,
            x="experiment_id",
            y="pH_drop",
            color="instability_flag",
            title="pH Drop Across Experiments",
            labels={"pH_drop": "pH Drop", "experiment_id": "Experiment"},
            color_discrete_map={0: "green", 1: "red"},
            template="plotly_white",
        )
        fig_ph.add_hline(y=0.8, line_dash="dash", line_color="red",
                         annotation_text="Instability threshold")
        st.plotly_chart(fig_ph, use_container_width=True)

    with c2:
        fig_vfa = px.scatter(
            df,
            x="experiment_id",
            y="VFA_accumulation_mg_L",
            color="instability_flag",
            title="VFA Accumulation",
            labels={"VFA_accumulation_mg_L": "VFA (mg/L)", "experiment_id": "Experiment"},
            color_discrete_map={0: "green", 1: "red"},
            template="plotly_white",
        )
        fig_vfa.add_hline(y=1200, line_dash="dash", line_color="red",
                          annotation_text="High VFA threshold")
        st.plotly_chart(fig_vfa, use_container_width=True)

    st.plotly_chart(stability_radar(df), use_container_width=True, key="stability_tab_radar")

    st.subheader("Instability Detection Summary")
    stable = df[df["instability_flag"] == 0]
    unstable = df[df["instability_flag"] == 1]
    sc1, sc2 = st.columns(2)
    with sc1:
        st.success(f"**Stable experiments:** {len(stable)}")
        if not stable.empty:
            st.dataframe(
                stable[["experiment_id", "pH_drop", "VFA_accumulation_mg_L",
                         "current_instability_index", "AI_score"]].describe().round(3),
                use_container_width=True,
            )
    with sc2:
        st.error(f"**Unstable experiments:** {len(unstable)}")
        if not unstable.empty:
            st.dataframe(
                unstable[["experiment_id", "pH_drop", "VFA_accumulation_mg_L",
                           "current_instability_index", "AI_score"]].describe().round(3),
                use_container_width=True,
            )


def render_prediction_tab(df: pd.DataFrame, models: dict) -> None:
    """Interactive prediction tab for new experiments."""
    st.subheader("Predict New Experiment")
    st.markdown("Adjust reactor parameters below to predict PFAS degradation performance.")

    with st.form("prediction_form"):
        pc1, pc2, pc3, pc4 = st.columns(4)
        with pc1:
            olr = st.number_input("OLR (kg/m³/d)", 1.0, 6.0, 3.5, 0.1)
            hrt = st.number_input("HRT (days)", 10.0, 30.0, 20.0, 1.0)
            ph = st.number_input("pH", 6.5, 8.0, 7.2, 0.1)
            temp = st.number_input("Temperature (°C)", 30.0, 42.0, 37.0, 0.5)
        with pc2:
            cod = st.number_input("COD (mg/L)", 2000.0, 8000.0, 5000.0, 100.0)
            vfa = st.number_input("VFA (mg/L)", 100.0, 1500.0, 500.0, 50.0)
            alk = st.number_input("Alkalinity (mg CaCO₃/L)", 1500.0, 5000.0, 3000.0, 100.0)
        with pc3:
            volt = st.number_input("Voltage (V)", 0.2, 1.2, 0.7, 0.05)
            curr = st.number_input("Current (A)", 0.1, 3.6, 1.5, 0.1)
            cd = st.number_input("Current Density (A/m²)", 0.1, 5.0, 1.5, 0.1)
            cond = st.number_input("Conductivity (mS/cm)", 1.0, 8.0, 4.0, 0.5)
        with pc4:
            ea = st.number_input("Electrode Area (m²)", 0.5, 2.0, 1.0, 0.1)
            es = st.number_input("Electrode Spacing (cm)", 1.0, 5.0, 3.0, 0.5)
            init_pfas = st.number_input("Initial PFAS (µg/L)", 50.0, 500.0, 200.0, 10.0)

        submitted = st.form_submit_button("Predict", type="primary", use_container_width=True)

    if submitted:
        features = np.array([[olr, hrt, ph, temp, cod, vfa, alk, volt, curr, cd, cond, ea, es, init_pfas]])
        deg_model, _ = models["degradation"]
        flu_model, _ = models["fluoride"]
        sc_model, _ = models["short_chain"]

        pred_deg = float(deg_model.predict(features)[0])
        pred_flu = float(flu_model.predict(features)[0])
        pred_sc_risk = int(sc_model.predict(features)[0])

        st.markdown("---")
        st.subheader("Prediction Results")
        pr1, pr2, pr3 = st.columns(3)
        pr1.metric("Predicted PFAS Degradation", f"{pred_deg:.1f}%")
        pr2.metric("Predicted Fluoride Release", f"{pred_flu:.1f} mg/L")
        pr3.metric(
            "Short-Chain Risk",
            "HIGH" if pred_sc_risk == 1 else "LOW",
            delta="⚠️" if pred_sc_risk == 1 else "✓",
            delta_color="inverse" if pred_sc_risk == 1 else "normal",
        )

        energy = volt * curr * 24 / 1000
        deg_norm = pred_deg / 65.0
        flu_norm = pred_flu / (df["fluoride_release_mg_L"].max() or 1)
        en_norm = energy / (df["energy_input_kWh_d"].max() or 1)
        ai_score = np.clip(
            0.40 * deg_norm + 0.30 * flu_norm - 0.15 * pred_sc_risk - 0.10 * en_norm - 0.05 * 0.1,
            0, 1
        )
        st.metric("Predicted AI Score", f"{ai_score:.3f}")


def render_data_tab(df: pd.DataFrame) -> None:
    """Raw data exploration tab."""
    st.subheader("Experiment Data Explorer")
    st.dataframe(df, use_container_width=True, hide_index=True, height=500)

    csv = df.to_csv(index=False).encode("utf-8")
    st.download_button(
        "Download Data (CSV)",
        csv,
        "pfas_sbead_experiments.csv",
        "text/csv",
        use_container_width=True,
    )


def main() -> None:
    df, mb_df = load_data()
    sidebar_config = render_sidebar(df)
    filtered_df = df[sidebar_config["mask"]].reset_index(drop=True)
    filtered_mb = mb_df[mb_df["experiment_id"].isin(filtered_df["experiment_id"])].reset_index(drop=True)

    st.title("Closed-Loop PFAS-SBEAD Optimization Pipeline")
    st.caption(
        "AI-driven optimization for PFAS degradation using Sidestream Bioelectrochemical "
        "Anaerobic Digestion (SBEAD). Maximizes degradation and fluoride release while "
        "minimizing short-chain accumulation, energy input, and reactor instability."
    )

    render_kpi_header(filtered_df)
    st.markdown("---")

    models = train_models(df)

    tab_overview, tab_models, tab_optim, tab_mass, tab_stability, tab_predict, tab_data = st.tabs([
        "📈 Overview",
        "🤖 AI Models",
        "🎯 Optimization",
        "⚖️ Mass Balance",
        "🛡️ Stability",
        "🔮 Predict",
        "📊 Data",
    ])

    with tab_overview:
        render_overview_tab(filtered_df, filtered_mb)
    with tab_models:
        render_ai_models_tab(filtered_df, models)
    with tab_optim:
        render_optimization_tab(filtered_df)
    with tab_mass:
        render_mass_balance_tab(filtered_df, filtered_mb)
    with tab_stability:
        render_reactor_stability_tab(filtered_df)
    with tab_predict:
        render_prediction_tab(filtered_df, models)
    with tab_data:
        render_data_tab(filtered_df)


if __name__ == "__main__":
    main()