from pathlib import Path
import sys

import altair as alt
import streamlit as st

try:
    from data_store import load_metric_store
except ModuleNotFoundError:
    sys.path.append(str(Path(__file__).resolve().parents[1]))
    from data_store import load_metric_store

st.set_page_config(page_title="Method Details", page_icon="📘", layout="wide")

# Sidebar navigation
st.sidebar.title("Navigation")
st.sidebar.page_link("streamlit_app.py", label="Home", icon="🏠")
st.sidebar.page_link("pages/OptimizationLeaderboard.py", label="Optimization Leaderboard", icon="📊")
st.sidebar.page_link("pages/UQLeaderboard.py", label="UQ Leaderboard", icon="🎯")
st.sidebar.page_link("pages/MethodDetails.py", label="Methods", icon="📘")
st.sidebar.page_link("pages/RawData.py", label="Get Data", icon="🧾")

st.title("Method Details")
st.caption("Method-level view from NetCDF forward-model-run metrics averaged over random seeds")

# Read selected method from query params
abbr = st.query_params.get("method")
if not isinstance(abbr, str):
    abbr = None

metric_store = load_metric_store()
if metric_store.empty:
    st.warning("No metric data found. Expected NetCDF files in `data/` with a `metric` variable.")
    st.stop()

methods_df = metric_store[["Method", "abbreviation"]].drop_duplicates().sort_values("abbreviation")

abbrs = methods_df["abbreviation"].tolist()

# Metadata dictionary (can be moved to a JSON/YAML later)
method_meta = {
    "TEKI": {
        "citation": "Chada et al., SIAM/ASA J. UQ, 2020",
        "url": "https://doi.org/10.1137/17M114402X",
        "summary": "EKI variant with Tikhonov regularization that adds a penalty term to prevent ensemble collapse and improve stability on nonlinear problems.",
    },
    "ETKI": {
        "citation": "Schillings & Stuart, Numer. Math., 2017",
        "url": "https://clima.github.io/EnsembleKalmanProcesses.jl/dev/",
        "summary": "Ensemble Transform Kalman Inversion — applies an ensemble-space transform update that preserves the ensemble mean while reducing variance inflation.",
    },
    "IEKF": {
        "citation": "Iglesias, Inverse Problems, 2016",
        "url": "https://doi.org/10.1088/0266-5611/32/2/025002",
        "summary": "Regularizing iterative ensemble Kalman method; repeatedly refines the ensemble around a regularized Gauss-Newton step for nonlinear inverse problems.",
    },
    "UKI": {
        "citation": "Huang, Huang & Stuart, Physica D, 2022",
        "url": "https://clima.github.io/EnsembleKalmanProcesses.jl/dev/",
        "summary": "Unscented Kalman Inversion — propagates a deterministic set of sigma points through the forward model to estimate mean and covariance without linearization.",
    },
    "ABC": {
        "citation": "Approximate Bayesian Calibration",
        "url": "https://example.com/abc",
        "summary": "Sample without exact likelihoods until error falls below a target convergence.",
    },
    "HM": {
        "citation": "Williamson et al. 2013; King et al. 2025",
        "url": "https://example.com/hm",
        "summary": "Iterative constraint of parameter space using wave reductions.",
    },
    "CES-EKI-DMC": {
        "citation": "Cleary et al., J. Comput. Phys., 2021",
        "url": "https://doi.org/10.1016/j.jcp.2020.109716",
        "summary": "Calibrate-Emulate-Sample: uses EKI with a DataMisfitController to select training points, builds a GP emulator of the forward model, then samples the posterior via MCMC.",
    },
    "CES-EKI-CONST": {
        "citation": "Cleary et al., J. Comput. Phys., 2021",
        "url": "https://doi.org/10.1016/j.jcp.2020.109716",
        "summary": "Calibrate-Emulate-Sample: uses EKI with a constant (fixed) timestep scheduler to select training points, builds a GP emulator of the forward model, then samples the posterior via MCMC.",
    },
    "CES-IEKF-CONST": {
        "citation": "Cleary et al., J. Comput. Phys., 2021; Iglesias, Inverse Problems, 2016",
        "url": "https://doi.org/10.1016/j.jcp.2020.109716",
        "summary": "Calibrate-Emulate-Sample: uses IEKF with a constant (fixed) timestep scheduler to select training points, builds a GP emulator of the forward model, then samples the posterior via MCMC.",
    },
    "ADAM": {
        "citation": "Kingma & Ba, ICLR, 2015",
        "url": "https://doi.org/10.48550/arXiv.1412.6980",
        "summary": "Adaptive Moment Estimation — gradient-based optimizer that adapts per-parameter learning rates using first and second moment estimates of the gradient.",
    },
    "LM": {
        "citation": "Levenberg, 1944; Marquardt, 1963; Fletcher, 1971",
        "url": "https://doi.org/10.1090/qam/10666",
        "summary": "Levenberg-Marquardt — damped least-squares algorithm that interpolates between gradient descent and Gauss-Newton steps for efficient nonlinear least-squares minimization.",
    },
}

# Selection UI (defaults to query param if valid)
default_idx = 0
if isinstance(abbr, str) and abbr in abbrs:
    default_idx = abbrs.index(abbr)
sel = str(st.selectbox("Choose a method", options=abbrs, index=default_idx))

# Persist selection to URL
st.query_params["method"] = sel

# Display details
row = methods_df.loc[methods_df["abbreviation"] == sel].iloc[0]
meta = method_meta.get(sel, {})

st.subheader(f"{sel} — {row['Method']}")

col1, col2, col3 = st.columns(3)
with col1:
    with st.container(border=True):
        st.markdown("**Citation**")
        st.write(meta.get("citation", "Citation pending"))
with col2:
    with st.container(border=True):
        st.markdown("**Link**")
        st.link_button("Open reference", meta.get("url", "https://example.com"))
with col3:
    with st.container(border=True):
        st.markdown("**Summary**")
        st.write(meta.get("summary", "Summary pending"))

st.markdown("**Benchmark Slice**")
slice_df = metric_store[metric_store["abbreviation"] == sel].copy()
slice_df = slice_df.sort_values(["benchmark", "rmse_target", "ensemble_size"])

target_options = sorted(slice_df["rmse_target"].astype(str).unique().tolist())
selected_target = st.radio("RMSE Target Level", options=target_options, horizontal=True)

best_table_view = slice_df[slice_df["rmse_target"].astype(str) == selected_target]

best_idx = best_table_view.groupby("benchmark")["metric"].idxmin()
best_ensemble_df = best_table_view.loc[best_idx, ["benchmark", "rmse_target", "ensemble_size", "metric", "failure_rate"]].rename(
    columns={"metric": "Mean Forward Model Runs", "failure_rate": "Failure Rate (%)", "ensemble_size": "Optimal Ensemble Size", "rmse_target": "RMSE Target"}
)
st.dataframe(best_ensemble_df, hide_index=True, use_container_width=True)

if sel == "HM":
    st.markdown("### Failure Analysis")
    st.write(f"Failed calibration rate for History Matching at RMSE Target {selected_target} (as % of random seeds).")
    
    chart = (
        alt.Chart(best_table_view)
        .mark_bar()
        .encode(
            x=alt.X("ensemble_size:O", title="Ensemble Size"),
            y=alt.Y("mean(failure_rate):Q", title="Failure Rate (%)", scale=alt.Scale(domain=[0, 100])),
            color="benchmark:N",
            tooltip=["benchmark", "ensemble_size", alt.Tooltip("mean(failure_rate):Q", format=".1f", title="Failure Rate (%)")]
        )
        .properties(height=350)
        .interactive()
    )
    st.altair_chart(chart, use_container_width=True)
best_by_benchmark = (
    best_table_view.loc[best_idx, ["benchmark", "metric", "ensemble_size"]]
    .rename(
        columns={
            "metric": "Best Mean Forward Model Runs",
            "ensemble_size": "Optimal Ensemble Size",
        }
    )
    .sort_values("benchmark")
)

st.markdown("**Best by Benchmark**")
st.dataframe(
    best_by_benchmark,
    hide_index=True,
    use_container_width=True,
    column_config={
        "benchmark": st.column_config.TextColumn("Benchmark"),
        "Best Mean Forward Model Runs": st.column_config.NumberColumn("Best Mean Forward Model Runs", format="%.4f"),
        "Optimal Ensemble Size": st.column_config.NumberColumn("Optimal Ensemble Size"),
    },
)

st.markdown("**Scaling by Benchmark**")
chart_view = slice_df[slice_df["rmse_target"].astype(str) == selected_target]
chart_df = (
    chart_view.groupby(["benchmark", "ensemble_size"], as_index=False)
    .agg(mean_forward_runs=("metric", "mean"))
)
chart = (
    alt.Chart(chart_df)
    .mark_line(point=True)
    .encode(
        x=alt.X("ensemble_size:Q", title="Ensemble Size"),
        y=alt.Y("mean_forward_runs:Q", title="Mean Forward Model Runs"),
        color=alt.Color("benchmark:N", title="Benchmark"),
        tooltip=["benchmark", "ensemble_size", alt.Tooltip("mean_forward_runs:Q", format=".4f")],
    )
)
st.altair_chart(chart, use_container_width=True)

st.markdown("**All Averaged Configurations for Method**")
st.dataframe(
    slice_df[["benchmark", "algorithm_alias", "rmse_target", "ensemble_size", "metric"]],
    hide_index=True,
    use_container_width=True,
    column_config={
        "benchmark": st.column_config.TextColumn("Benchmark"),
        "algorithm_alias": st.column_config.TextColumn("Source Alias"),
        "rmse_target": st.column_config.TextColumn("RMSE Target"),
        "ensemble_size": st.column_config.NumberColumn("Ensemble Size"),
        "metric": st.column_config.NumberColumn("Mean Forward Model Runs", format="%.4f"),
    },
)

st.page_link("pages/OptimizationLeaderboard.py", label="⬅ Back to Optimization Leaderboard", icon="↩️")