#!/usr/bin/env python3
from __future__ import annotations

import argparse
import html
from pathlib import Path
import sys
from typing import Dict, Iterable, List, Tuple

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from src.config_utils import prepare_output_dirs, read_config

PRIMARY_CASES = [
    ("pr_o", "RMAE_percent", True, "Precipitation RMAE (%)"),
    ("pr_o", "CORR", False, "Precipitation CORR"),
    ("t2_o", "RMSE", True, "T2m RMSE (°C)"),
    ("t2_o", "CORR", False, "T2m CORR"),
]


def _safe_wmean(g: pd.DataFrame) -> float:
    w = pd.to_numeric(g["n_grid"], errors="coerce").astype(float)
    x = pd.to_numeric(g["mean_value"], errors="coerce").astype(float)
    ok = np.isfinite(x) & np.isfinite(w) & (w > 0)
    if ok.sum() == 0:
        return np.nan
    return float(np.sum(x[ok] * w[ok]) / np.sum(w[ok]))


def build_weighted_domain(summary_all: pd.DataFrame) -> pd.DataFrame:
    """Create domain-weighted metrics by aggregating R1-R7 with n_grid weights.

    The ALL row is not used, so this is robust even if ALL was calculated by a
    simple unweighted mean over valid grid cells in an earlier version.
    """
    region_rows = summary_all[summary_all["region"].astype(str).str.match(r"^R[1-7]$", na=False)].copy()
    group_cols = [
        "source_model", "lead", "target", "spatial_mode",
        "predictor_set", "method", "metric",
    ]
    # Apply on only the value/weight columns to avoid pandas groupby.apply deprecation warnings.
    out = (
        region_rows.groupby(group_cols, dropna=False)[["mean_value", "n_grid"]]
        .apply(_safe_wmean)
        .reset_index(name="weighted_mean")
    )
    return out


def best_vs_raw(weighted: pd.DataFrame) -> pd.DataFrame:
    rows = []
    for target, metric, small_better, _title in PRIMARY_CASES:
        x = weighted[(weighted["target"] == target) & (weighted["metric"] == metric)].copy()
        if x.empty:
            continue
        for (source_model, lead), g in x.groupby(["source_model", "lead"]):
            raw = g[g["method"].astype(str).str.lower() == "raw"].copy()
            best = g.sort_values("weighted_mean", ascending=small_better).iloc[0]
            raw_value = float(raw["weighted_mean"].iloc[0]) if len(raw) else np.nan
            best_value = float(best["weighted_mean"])
            if np.isfinite(raw_value):
                if small_better:
                    improvement = 100.0 * (raw_value - best_value) / raw_value
                    absolute_change = raw_value - best_value
                else:
                    improvement = 100.0 * (best_value - raw_value) / abs(raw_value)
                    absolute_change = best_value - raw_value
            else:
                improvement = np.nan
                absolute_change = np.nan
            rows.append({
                "source_model": source_model,
                "lead": int(lead),
                "target": target,
                "metric": metric,
                "raw_value": raw_value,
                "best_value": best_value,
                "absolute_change": absolute_change,
                "improvement_percent": improvement,
                "best_spatial_mode": best["spatial_mode"],
                "best_predictor_set": best["predictor_set"],
                "best_method": best["method"],
            })
    return pd.DataFrame(rows).sort_values(["target", "metric", "source_model", "lead"])


def best_per_source_lead(weighted: pd.DataFrame) -> pd.DataFrame:
    rows = []
    for target, metric, small_better, _title in PRIMARY_CASES:
        x = weighted[(weighted["target"] == target) & (weighted["metric"] == metric)].copy()
        if x.empty:
            continue
        x = x.sort_values("weighted_mean", ascending=small_better)
        best = x.groupby(["source_model", "lead"], as_index=False).first()
        best["rank_metric"] = metric
        best["small_better"] = small_better
        rows.append(best)
    return pd.concat(rows, ignore_index=True) if rows else pd.DataFrame()


def best_by_method(weighted: pd.DataFrame) -> pd.DataFrame:
    """For each method and source/lead, keep its best spatial/predictor setup."""
    rows = []
    for target, metric, small_better, _title in PRIMARY_CASES:
        x = weighted[(weighted["target"] == target) & (weighted["metric"] == metric)].copy()
        if x.empty:
            continue
        x = x.sort_values("weighted_mean", ascending=small_better)
        best = x.groupby(["source_model", "lead", "target", "method", "metric"], as_index=False).first()
        best["small_better"] = small_better
        rows.append(best)
    return pd.concat(rows, ignore_index=True) if rows else pd.DataFrame()


def method_ranking(best_method_df: pd.DataFrame) -> pd.DataFrame:
    """Rank methods after each method has been allowed to choose its best setup."""
    rows = []
    for target, metric, small_better, _title in PRIMARY_CASES:
        x = best_method_df[(best_method_df["target"] == target) & (best_method_df["metric"] == metric)].copy()
        if x.empty:
            continue
        agg = x.groupby("method", as_index=False).agg(
            mean_score=("weighted_mean", "mean"),
            median_score=("weighted_mean", "median"),
            n_cases=("weighted_mean", "count"),
        )
        agg = agg.sort_values("mean_score", ascending=small_better).reset_index(drop=True)
        agg.insert(0, "rank", np.arange(1, len(agg) + 1))
        agg.insert(1, "target", target)
        agg.insert(2, "metric", metric)
        agg["small_better"] = small_better
        rows.append(agg)
    return pd.concat(rows, ignore_index=True) if rows else pd.DataFrame()


def ridge_comparison(best_method_df: pd.DataFrame) -> pd.DataFrame:
    """Compare each method with Ridge after method-wise best setup selection."""
    rows = []
    for target, metric, small_better, _title in PRIMARY_CASES:
        x = best_method_df[(best_method_df["target"] == target) & (best_method_df["metric"] == metric)].copy()
        if x.empty:
            continue
        for (source_model, lead), g in x.groupby(["source_model", "lead"]):
            rg = g[g["method"] == "ridge"]
            if rg.empty:
                continue
            ridge_val = float(rg["weighted_mean"].iloc[0])
            for _, row in g.iterrows():
                val = float(row["weighted_mean"])
                if small_better:
                    # Positive means the candidate improves upon Ridge.
                    rel_vs_ridge = 100.0 * (ridge_val - val) / ridge_val
                    abs_vs_ridge = ridge_val - val
                else:
                    rel_vs_ridge = 100.0 * (val - ridge_val) / abs(ridge_val)
                    abs_vs_ridge = val - ridge_val
                rows.append({
                    "source_model": source_model,
                    "lead": int(lead),
                    "target": target,
                    "metric": metric,
                    "method": row["method"],
                    "score": val,
                    "ridge_score": ridge_val,
                    "absolute_gain_vs_ridge": abs_vs_ridge,
                    "relative_gain_percent_vs_ridge": rel_vs_ridge,
                    "spatial_mode": row["spatial_mode"],
                    "predictor_set": row["predictor_set"],
                })
    return pd.DataFrame(rows)


def _fmt_df(df: pd.DataFrame, max_rows: int = 200) -> str:
    if df.empty:
        return "<p><em>No data.</em></p>"
    show = df.copy().head(max_rows)
    for c in show.columns:
        if pd.api.types.is_float_dtype(show[c]):
            show[c] = show[c].map(lambda v: "" if pd.isna(v) else f"{v:.4f}")
    return show.to_html(index=False, escape=False, classes="data-table")


def _plot_best_vs_raw(best_raw: pd.DataFrame, report_fig_dir: Path) -> List[Path]:
    out_files: List[Path] = []
    for target, metric, small_better, title in PRIMARY_CASES:
        d = best_raw[(best_raw["target"] == target) & (best_raw["metric"] == metric)].copy()
        if d.empty:
            continue
        d["label"] = d["source_model"].astype(str) + " L" + d["lead"].astype(str)
        x = np.arange(len(d))
        width = 0.38
        fig, ax = plt.subplots(figsize=(max(9, 0.55 * len(d)), 5), constrained_layout=True)
        ax.bar(x - width/2, d["raw_value"].astype(float), width, label="RAW")
        ax.bar(x + width/2, d["best_value"].astype(float), width, label="Best")
        ax.set_xticks(x)
        ax.set_xticklabels(d["label"], rotation=45, ha="right")
        ax.set_title(title + ": RAW vs best")
        ax.set_ylabel(metric)
        ax.legend()
        fname = report_fig_dir / f"best_vs_raw_{target}_{metric}.png"
        fig.savefig(fname, dpi=150)
        plt.close(fig)
        out_files.append(fname)
    return out_files


def _plot_method_ranking(rank: pd.DataFrame, report_fig_dir: Path) -> List[Path]:
    out_files: List[Path] = []
    for target, metric, small_better, title in PRIMARY_CASES:
        d = rank[(rank["target"] == target) & (rank["metric"] == metric)].copy()
        if d.empty:
            continue
        d = d.sort_values("mean_score", ascending=small_better)
        fig, ax = plt.subplots(figsize=(9, max(4, 0.35 * len(d))), constrained_layout=True)
        ax.barh(d["method"], d["mean_score"].astype(float))
        ax.invert_yaxis()
        ax.set_xlabel(metric)
        ax.set_title(title + ": method ranking, lower is better" if small_better else title + ": method ranking, higher is better")
        fname = report_fig_dir / f"method_ranking_{target}_{metric}.png"
        fig.savefig(fname, dpi=150)
        plt.close(fig)
        out_files.append(fname)
    return out_files


def write_html_report(
    report_file: Path,
    cfg: Dict,
    weighted: pd.DataFrame,
    best_raw: pd.DataFrame,
    best_lead: pd.DataFrame,
    best_method_df: pd.DataFrame,
    rank: pd.DataFrame,
    ridge_cmp: pd.DataFrame,
    fig_files: List[Path],
) -> None:
    report_file.parent.mkdir(parents=True, exist_ok=True)
    rel_figs = [p.relative_to(report_file.parent).as_posix() for p in fig_files]

    # Compact top tables.
    top_tables = []
    for target, metric, small_better, title in PRIMARY_CASES:
        d = weighted[(weighted["target"] == target) & (weighted["metric"] == metric)].copy()
        d = d.sort_values("weighted_mean", ascending=small_better).head(30)
        top_tables.append((title, d))

    # Candidate methods that beat Ridge in any source/lead case.
    beats_ridge = ridge_cmp[(ridge_cmp["method"] != "ridge") & (ridge_cmp["relative_gain_percent_vs_ridge"] > 0)].copy()
    beats_ridge = beats_ridge.sort_values("relative_gain_percent_vs_ridge", ascending=False)

    cfg_summary = {
        "data_dir": cfg.get("data_dir"),
        "mask_file": cfg.get("mask_file"),
        "source_models": cfg.get("source_models"),
        "leads": cfg.get("leads"),
        "targets": cfg.get("targets"),
        "spatial_modes": cfg.get("spatial_modes"),
        "predictor_sets": cfg.get("predictor_sets"),
        "methods": cfg.get("methods"),
        "train_years": cfg.get("train_years"),
        "test_years": cfg.get("test_years"),
    }

    html_parts = []
    html_parts.append("""<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<title>ML/DL Bias Correction Report</title>
<style>
body { font-family: Arial, Helvetica, sans-serif; margin: 24px; line-height: 1.45; color: #222; }
h1, h2, h3 { color: #1f2937; }
.notice { background: #f3f4f6; border-left: 4px solid #6b7280; padding: 10px 14px; margin: 14px 0; }
.ok { background: #ecfdf5; border-left: 4px solid #10b981; padding: 10px 14px; margin: 14px 0; }
.warn { background: #fffbeb; border-left: 4px solid #f59e0b; padding: 10px 14px; margin: 14px 0; }
.data-table { border-collapse: collapse; font-size: 12px; margin: 12px 0 24px 0; width: 100%; }
.data-table th, .data-table td { border: 1px solid #ddd; padding: 5px 7px; text-align: right; }
.data-table th { background: #f4f4f4; position: sticky; top: 0; }
.data-table td:nth-child(1), .data-table td:nth-child(2), .data-table td:nth-child(3), .data-table td:nth-child(4), .data-table td:nth-child(5), .data-table td:nth-child(6), .data-table td:nth-child(7) { text-align: left; }
img { max-width: 100%; border: 1px solid #ddd; margin: 8px 0 18px 0; }
code, pre { background: #f4f4f4; padding: 2px 4px; }
pre { padding: 12px; overflow-x: auto; }
.toc a { display: block; margin: 4px 0; }
</style>
</head>
<body>
""")
    html_parts.append("<h1>ML/DL Bias Correction Report</h1>")
    html_parts.append("<div class='notice'>This report is automatically generated from <code>outputs/metrics/summary_all.csv</code>. Metrics are calculated on the independent test period only.</div>")
    html_parts.append("<h2>Table of contents</h2><div class='toc'>")
    for anchor, name in [
        ("config", "Experiment configuration"),
        ("best-vs-raw", "Best correction vs RAW"),
        ("best-lead", "Best configuration per source model and lead"),
        ("method-ranking", "Method ranking"),
        ("ridge", "Comparison against Ridge benchmark"),
        ("top", "Top configurations"),
        ("figures", "Figures"),
    ]:
        html_parts.append(f"<a href='#{anchor}'>{html.escape(name)}</a>")
    html_parts.append("</div>")

    html_parts.append("<h2 id='config'>Experiment configuration</h2>")
    html_parts.append("<pre>" + html.escape(pd.Series(cfg_summary).to_string()) + "</pre>")

    html_parts.append("<h2 id='best-vs-raw'>Best correction vs RAW</h2>")
    html_parts.append("<div class='ok'>Positive improvement means the selected correction improves upon RAW. For CORR, the absolute change is usually more interpretable than the percentage change.</div>")
    for target, metric, _small, title in PRIMARY_CASES:
        d = best_raw[(best_raw["target"] == target) & (best_raw["metric"] == metric)].copy()
        html_parts.append(f"<h3>{html.escape(title)}</h3>")
        html_parts.append(_fmt_df(d, 200))

    html_parts.append("<h2 id='best-lead'>Best configuration per source model and lead</h2>")
    for target, metric, _small, title in PRIMARY_CASES:
        d = best_lead[(best_lead["target"] == target) & (best_lead["rank_metric"] == metric)].copy()
        html_parts.append(f"<h3>{html.escape(title)}</h3>")
        html_parts.append(_fmt_df(d, 200))

    html_parts.append("<h2 id='method-ranking'>Method ranking</h2>")
    html_parts.append("<div class='notice'>Each method is first allowed to choose its best spatial mode and predictor set for each source model and lead. The table then averages those best scores across all source-model/lead cases.</div>")
    for target, metric, _small, title in PRIMARY_CASES:
        d = rank[(rank["target"] == target) & (rank["metric"] == metric)].copy()
        html_parts.append(f"<h3>{html.escape(title)}</h3>")
        html_parts.append(_fmt_df(d, 200))

    html_parts.append("<h2 id='ridge'>Comparison against Ridge benchmark</h2>")
    html_parts.append("<div class='warn'>Positive gain means the candidate method beats Ridge for that source-model/lead case. Negative gain means Ridge remains better.</div>")
    if beats_ridge.empty:
        html_parts.append("<p><strong>No non-Ridge method beats Ridge in the selected primary metrics.</strong></p>")
    else:
        html_parts.append("<h3>Cases where a non-Ridge method beats Ridge</h3>")
        html_parts.append(_fmt_df(beats_ridge, 200))
    html_parts.append("<h3>Full Ridge comparison</h3>")
    html_parts.append(_fmt_df(ridge_cmp.sort_values(["target", "metric", "source_model", "lead", "method"]), 500))

    html_parts.append("<h2 id='top'>Top configurations</h2>")
    for title, d in top_tables:
        html_parts.append(f"<h3>{html.escape(title)}</h3>")
        html_parts.append(_fmt_df(d, 30))

    html_parts.append("<h2 id='figures'>Figures</h2>")
    for rel in rel_figs:
        html_parts.append(f"<h3>{html.escape(Path(rel).stem)}</h3><img src='{html.escape(rel)}' alt='{html.escape(rel)}'>")

    html_parts.append("</body></html>")
    report_file.write_text("\n".join(html_parts), encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description="Analyze experiment summaries and build an HTML report.")
    parser.add_argument("--config", default="config/experiment.yaml")
    args = parser.parse_args()

    cfg_path = Path(args.config).resolve()
    cfg = read_config(cfg_path)
    dirs = prepare_output_dirs(cfg, base_dir=cfg_path.parent.parent)
    report_dir = dirs["output"] / "report"
    report_fig_dir = report_dir / "figures"
    report_fig_dir.mkdir(parents=True, exist_ok=True)

    summary_file = dirs["metrics"] / "summary_all.csv"
    if not summary_file.exists():
        raise FileNotFoundError(f"Missing summary file: {summary_file}. Run scripts/02_collect_summary.py first.")
    summary = pd.read_csv(summary_file)

    weighted = build_weighted_domain(summary)
    weighted_file = dirs["metrics"] / "summary_weighted_domain.csv"
    weighted.to_csv(weighted_file, index=False)

    br = best_vs_raw(weighted)
    br_file = dirs["metrics"] / "best_vs_raw.csv"
    br.to_csv(br_file, index=False)

    bpl = best_per_source_lead(weighted)
    bpl_file = dirs["metrics"] / "best_per_source_lead.csv"
    bpl.to_csv(bpl_file, index=False)

    bbm = best_by_method(weighted)
    bbm_file = dirs["metrics"] / "best_by_method.csv"
    bbm.to_csv(bbm_file, index=False)

    rank = method_ranking(bbm)
    rank_file = dirs["metrics"] / "method_ranking.csv"
    rank.to_csv(rank_file, index=False)

    rc = ridge_comparison(bbm)
    rc_file = dirs["metrics"] / "ridge_comparison.csv"
    rc.to_csv(rc_file, index=False)

    figs = []
    figs.extend(_plot_best_vs_raw(br, report_fig_dir))
    figs.extend(_plot_method_ranking(rank, report_fig_dir))

    report_file = report_dir / "index.html"
    write_html_report(report_file, cfg, weighted, br, bpl, bbm, rank, rc, figs)

    print(f"Wrote {weighted_file}")
    print(f"Wrote {br_file}")
    print(f"Wrote {bpl_file}")
    print(f"Wrote {bbm_file}")
    print(f"Wrote {rank_file}")
    print(f"Wrote {rc_file}")
    print(f"Wrote HTML report: {report_file}")


if __name__ == "__main__":
    main()