import sys
from pathlib import Path

import pandas as pd
import streamlit as st

try:
    from common.method_registry import KNOWN_METHODS
    from common.leaderboard import _UPDATE_TYPE_ORDER
except ModuleNotFoundError:
    sys.path.append(str(Path(__file__).resolve().parent))
    from common.method_registry import KNOWN_METHODS
    from common.leaderboard import _UPDATE_TYPE_ORDER

st.set_page_config(page_title="Calibration Benchmark", 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("Calibration Benchmark")
st.markdown(
    "A benchmark comparing parameter-calibration methods on chaotic dynamical systems. "
    "Methods are ranked by **forward-model run efficiency** — how many forward-model "
    "evaluations are needed, on average across random seeds, to reach a target accuracy."
)

st.divider()

col1, col2 = st.columns(2, gap="large")

with col1:
    st.subheader("📊 Optimization Leaderboard")
    st.markdown(
        "Ranks methods by mean forward-model runs to reach an **RMSE target** on "
        "Lorenz-63 and Lorenz-96 benchmarks. "
        "Lower is better; failed runs are tracked separately as a failure rate."
    )
    st.page_link("pages/OptimizationLeaderboard.py", label="Go to Optimization Leaderboard →")

with col2:
    st.subheader("🎯 UQ Leaderboard")
    st.markdown(
        "Ranks methods by mean forward-model runs to reach an **uncertainty quantification "
        "target**. Same metric and benchmarks as the Optimization Leaderboard, evaluated "
        "at a UQ-specific convergence criterion."
    )
    st.page_link("pages/UQLeaderboard.py", label="Go to UQ Leaderboard →")

st.divider()

st.subheader("Benchmarks")
st.markdown(
    "Results are reported on four benchmark configurations of the [Lorenz system]"
    "(https://en.wikipedia.org/wiki/Lorenz_system), a standard testbed for "
    "data-assimilation and calibration algorithms:"
)

_media = Path(__file__).parent / "media"

st.markdown("- **L63** — Lorenz-63 — 3-variable chaotic attractor; learn 2 parameters, strongly nonlinear.")
st.image(
    str(_media / "posterior_ribbons_20_13_k5.png"),
    caption="Example prior-posterior & truth. L: Difference to true parameter. R: data-sample/output predictived distribution (state-mean [1:3],  state-covariance (diag [4:6] and off-diag [7:9]))",
    width=800,
)


st.markdown("- **L96** — Lorenz-96 (40-variable); learn 1-parameter constant forcing.")
st.image(
    str(_media / "posterior_ribbons_const-force_12_1_k3.png"),
    caption="Example prior-posterior & truth. L: Difference to true parameter. C: parameter-induced forcing of the L96 system. R: data-sample/output predictived distribution ([1:40] state-mean [41:80] state-std)",
    width=1200,
)

st.markdown("- **L96_SPATIAL_FORCING** — Lorenz-96 (40-variable) with spatially-varying forcing; learn 40 parameters; moderately correlated prior.")
st.image(
    str(_media / "posterior_ribbons_vec-force_65_1_k3.png"),
    caption="Example prior-posterior & truth. L: Difference to true parameters. C: parameter-induced forcing of the L96 system. R: data-sample/output predictived distribution ([1:40] state-mean [41:80] state-std)",
    width=1200,
)


st.markdown("- **L96_NN_FORCING** — Lorenz-96 (100-variable) with a neural-network forcing; Learn 61 parameters (weights and biases of the network). Reasonable prior given.")
st.image(
    str(_media / "posterior_ribbons_flux-force_80_1_k3.png"),
    caption="Example prior-posterior & truth. L: Difference to true parameters (weights). C: parameter-induced forcing of the L96 system. R: data-sample/output predictived distribution ([1:100] state-mean [101:200] state-std)",
    width=1200,
)

st.subheader("Method taxonomy")
st.markdown(
    "Every method carries four independent tags — how it searches, what update "
    "mechanism drives each step, what it's built to report, and whether/when it uses "
    "a surrogate model. See the **📘 Methods** page for citations and per-method "
    "performance charts."
)
_taxonomy_rows = [
    {
        "Algorithm": info["abbreviation"],
        "Method Goal": info["method_goal"],
        "Update Type": info["update_type"],
        "Parallelism": info["parallelism"],
        "Emulator Use": info["emulator_use"],
        "_update_type_order": _UPDATE_TYPE_ORDER.get(info["update_type"], 99),
    }
    for info in KNOWN_METHODS.values()
]
_taxonomy_df = (
    pd.DataFrame(_taxonomy_rows)
    .sort_values(["_update_type_order", "Algorithm"])
    .drop(columns="_update_type_order")
)

_UPDATE_TYPE_ROW_COLOR = {
    "kalman": "#d9f2e6",    # pastel green
    "gradient": "#dbe9fa",  # pastel blue
    "general": "#fde2c8",   # pastel orange
}


def _style_taxonomy_rows(df: pd.DataFrame) -> pd.DataFrame:
    result = pd.DataFrame("", index=df.index, columns=df.columns)
    for row in df.index:
        bg = _UPDATE_TYPE_ROW_COLOR.get(df.loc[row, "Update Type"], "")
        if bg:
            result.loc[row, :] = f"background-color: {bg}; color: #212529"
    return result


st.dataframe(
    _taxonomy_df.style.apply(_style_taxonomy_rows, axis=None),
    hide_index=True,
    use_container_width=True,
    height=int((len(_taxonomy_df) + 1) * 35.2 + 3),
    column_config={
        "Algorithm": st.column_config.TextColumn("Algorithm"),
        "Method Goal": st.column_config.TextColumn("Method Goal"),
        "Update Type": st.column_config.TextColumn("Update Type"),
        "Parallelism": st.column_config.TextColumn("Parallelism"),
        "Emulator Use": st.column_config.TextColumn("Emulator Use"),
    },
)
st.caption(
    "Note: Kalman methods are Bayesian in spirit too (they're approximate Gaussian "
    "posterior updates) — update type is about mechanism (gradient vs. Kalman vs. "
    "general), not whether a method is 'Bayesian'."
)

st.subheader("Key metric")
st.markdown(
    "The reported metric is the **mean number of forward-model evaluations** required "
    "to reach the target, averaged over random seeds. "
    "A value of **−1** indicates a failed run (did not reach the target); "
    "the **failure rate** shows the fraction of seeds that failed."
)