#!/usr/bin/env python
"""Phase 3/4 evaluation: every metric and every paired contrast is computed from saved prediction files.
Outputs (artifacts/evaluation + results/): metrics_all.csv, contrasts.csv, horizon_ladder.csv, secondary_targets.csv,
firm_holdout.csv, walkforward_by_year.csv, reliability.csv, continuous.csv, power.json."""
from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import polars as pl

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from wp39.inference import (binary_metrics, continuous_metrics, dm_gw_test, holm, paired_block_bootstrap,  # noqa: E402
                            pointwise_logloss, quarter_blocks, reliability_table)
from wp39.pipeline import Paths, get_logger, load_config  # noqa: E402
from wp39.provenance import now_iso, write_json  # noqa: E402

CONTRASTS = [("B5", "B2"), ("B4", "B2"), ("B3", "B2"), ("B5", "B4"), ("B5", "B3"), ("B5all", "B2"), ("B5all", "B5"), ("B2", "B1"), ("B1", "B0")]
FAMILIES = {"freq_vs_daily": [("B5", "B2"), ("B4", "B2"), ("B3", "B2")], "fine_vs_coarse": [("B5", "B4"), ("B5", "B3")]}


def load_preds(P: Paths) -> pl.DataFrame:
    files = sorted((P.evaluation / "preds").glob("*.parquet"))
    frames = []
    for f in files:
        d = pl.read_parquet(f)
        target = [c for c in d.columns if c.startswith("y_") or c == "S_opinc"][0]
        d = d.rename({target: "y"}).with_columns(pl.lit(target).alias("target"))
        if "fold" not in d.columns:
            d = d.with_columns(pl.lit(None, dtype=pl.Int64).alias("fold"))
        d = d.with_columns([pl.col("fold").cast(pl.Int64), pl.col("h").cast(pl.Int64), pl.col("y").cast(pl.Float64), pl.col("p").cast(pl.Float64)])
        frames.append(d.select(["entity_id", "symbol", "fiscal_period_end", "h", "cutoff", "A_q", "sic_division", "release_time",
                                "y", "p", "model", "part", "tag", "target", "fold"]))
    return pl.concat(frames)


def contrast(a: pl.DataFrame, b: pl.DataFrame, n_boot: int, seed: int) -> dict:
    key = ["entity_id", "fiscal_period_end", "cutoff"]
    j = a.select(key + ["y", "p", "A_q"]).join(b.select(key + [pl.col("p").alias("pb")]), on=key, how="inner")
    y = j["y"].to_numpy().astype(int); pa, pb = j["p"].to_numpy(), j["pb"].to_numpy()
    d_ll = pointwise_logloss(y, pa) - pointwise_logloss(y, pb)
    d_br = (pa - y) ** 2 - (pb - y) ** 2
    blocks = quarter_blocks(j["A_q"].to_numpy())
    bll, bbr = paired_block_bootstrap(d_ll, blocks, n_boot, seed), paired_block_bootstrap(d_br, blocks, n_boot, seed)
    gw = dm_gw_test(d_ll, blocks)
    return {"n": int(len(y)), "d_logloss": bll["mean"], "ci_lo": bll["ci_lo"], "ci_hi": bll["ci_hi"], "p_boot": bll["p_boot"],
            "gw_t": gw["t"], "gw_p": gw["p"], "n_quarters": gw["T"], "d_brier": bbr["mean"], "brier_ci_lo": bbr["ci_lo"], "brier_ci_hi": bbr["ci_hi"]}


def main() -> int:
    P = Paths.default()
    log = get_logger("07_eval", P.logs / "07_eval.log")
    cfg = load_config("research")
    hp, nb, seed = cfg["cutoffs"]["primary_horizon_td"], cfg["inference"]["n_bootstrap"], cfg["seed"]
    preds = load_preds(P)
    log.info("prediction rows: %d (%d files)", preds.height, len(list((P.evaluation / "preds").glob("*.parquet"))))

    # 1) metrics for every binary prediction set
    rows = []
    grp = preds.filter(pl.col("target") != "S_opinc").with_columns(pl.col("cutoff").dt.year().alias("year"))
    for keys, g in grp.group_by(["h", "target", "model", "part", "tag"], maintain_order=True):
        h, target, model, part, tag = keys
        m = binary_metrics(g["y"].to_numpy().astype(int), g["p"].to_numpy())
        rows.append({"h": h, "target": target, "model": model, "part": part, "tag": tag, **m})
    metrics = pd.DataFrame(rows).sort_values(["target", "h", "tag", "part", "model"])
    metrics.to_csv(P.evaluation / "metrics_all.csv", index=False)

    # 2) primary contrasts (h = hp, primary target, chronological test) for logistic and LightGBM
    out = []
    base = preds.filter((pl.col("h") == hp) & (pl.col("target") == "y_adverse_opinc") & (pl.col("part") == "test") & (pl.col("tag") == "chrono"))
    for prefix in ("", "L"):
        pv = {}
        for a, b in CONTRASTS:
            A, B = base.filter(pl.col("model") == prefix + a), base.filter(pl.col("model") == (prefix + b if b != "B0" else "B0"))
            if A.height == 0 or B.height == 0:
                continue
            c = contrast(A, B, nb, seed); c.update({"family": "logistic" if prefix == "" else "lightgbm", "a": prefix + a, "b": prefix + b})
            out.append(c); pv[(a, b)] = c["gw_p"]
        for fam, pairs in FAMILIES.items():
            adj = holm({f"{a}-{b}": pv[(a, b)] for a, b in pairs if (a, b) in pv})
            for r in out:
                if r["family"] == ("logistic" if prefix == "" else "lightgbm") and f"{r['a'][len(prefix):]}-{r['b'][len(prefix):]}" in adj:
                    r["holm_family"], r["gw_p_holm"] = fam, adj[f"{r['a'][len(prefix):]}-{r['b'][len(prefix):]}"]
    contrasts = pd.DataFrame(out); contrasts.to_csv(P.evaluation / "contrasts.csv", index=False)
    log.info("primary contrasts:\n%s", contrasts[["family", "a", "b", "n", "d_logloss", "ci_lo", "ci_hi", "gw_p"]].to_string())

    # 3) horizon ladder (B5 vs B2) logistic chrono + LightGBM fixed
    lad = []
    for h in sorted(preds["h"].unique().to_list()):
        for fam, tag, a, b in (("logistic", "chrono", "B5", "B2"), ("lightgbm", "fixed", "LB5", "LB2")):
            s = preds.filter((pl.col("h") == h) & (pl.col("target") == "y_adverse_opinc") & (pl.col("part") == "test") & (pl.col("tag") == tag))
            A, B = s.filter(pl.col("model") == a), s.filter(pl.col("model") == b)
            if A.height and B.height:
                c = contrast(A, B, nb, seed); c.update({"h": h, "family": fam})
                c["auc_a"] = binary_metrics(A["y"].to_numpy().astype(int), A["p"].to_numpy()).get("auc")
                c["auc_b"] = binary_metrics(B["y"].to_numpy().astype(int), B["p"].to_numpy()).get("auc")
                lad.append(c)
    ladder = pd.DataFrame(lad); ladder.to_csv(P.evaluation / "horizon_ladder.csv", index=False)
    h4 = {}
    for fam, g in ladder.groupby("family"):
        g = g.sort_values("h", ascending=False)
        from scipy.stats import spearmanr
        rho, p = spearmanr(-g["h"], -g["d_logloss"])  # gain (= -d) should rise as h falls
        h4[fam] = {"spearman_gain_vs_minus_h": float(rho), "p": float(p), "horizons": g["h"].tolist(), "d_logloss": g["d_logloss"].round(5).tolist()}
    write_json(P.evaluation / "h4_timing.json", h4)

    # 4) secondary targets (B5 vs B2) logistic + lightgbm at hp
    sec = []
    for target in ("y_adverse_netinc", "y_roa_decline", "y_ocf_negative", "y_liquidity_deterioration", "y_abs_surprise", "y_adverse_opinc"):
        for fam, a, b in (("logistic", "B5", "B2"), ("lightgbm", "LB5", "LB2")):
            s = preds.filter((pl.col("h") == hp) & (pl.col("target") == target) & (pl.col("part") == "test") & (pl.col("tag") == "chrono"))
            A, B = s.filter(pl.col("model") == a), s.filter(pl.col("model") == b)
            if A.height and B.height:
                c = contrast(A, B, nb, seed); c.update({"target": target, "family": fam, "base_rate": float(A["y"].mean()),
                                                        "auc_a": binary_metrics(A["y"].to_numpy().astype(int), A["p"].to_numpy()).get("auc"),
                                                        "auc_b": binary_metrics(B["y"].to_numpy().astype(int), B["p"].to_numpy()).get("auc")})
                sec.append(c)
    pd.DataFrame(sec).to_csv(P.evaluation / "secondary_targets.csv", index=False)

    # 5) firm hold-out
    fh = []
    for tag in ("firmhold_heldout", "firmhold_seen"):
        for fam, a, b in (("logistic", "B5", "B2"), ("lightgbm", "LB5", "LB2")):
            s = preds.filter((pl.col("h") == hp) & (pl.col("target") == "y_adverse_opinc") & (pl.col("part") == "test") & (pl.col("tag") == tag))
            A, B = s.filter(pl.col("model") == a), s.filter(pl.col("model") == b)
            if A.height and B.height:
                c = contrast(A, B, nb, seed); c.update({"tag": tag, "family": fam}); fh.append(c)
    pd.DataFrame(fh).to_csv(P.evaluation / "firm_holdout.csv", index=False)

    # 6) walk-forward by year
    wf = []
    oof = preds.filter((pl.col("tag") == "walkforward") & (pl.col("h") == hp)).with_columns(pl.col("cutoff").dt.year().alias("year"))
    for keys, g in oof.group_by(["model", "year"], maintain_order=True):
        m = binary_metrics(g["y"].to_numpy().astype(int), g["p"].to_numpy()); wf.append({"model": keys[0], "year": keys[1], **m})
    wfd = pd.DataFrame(wf).sort_values(["model", "year"]); wfd.to_csv(P.evaluation / "walkforward_by_year.csv", index=False)
    wfc = []
    for fam, a, b in (("logistic", "B5", "B2"), ("lightgbm", "LB5", "LB2")):
        A, B = oof.filter(pl.col("model") == a), oof.filter(pl.col("model") == b)
        if A.height and B.height:
            for y in sorted(A["year"].unique().to_list()):
                c = contrast(A.filter(pl.col("year") == y), B.filter(pl.col("year") == y), 300, seed); c.update({"family": fam, "year": y}); wfc.append(c)
            c = contrast(A, B, nb, seed); c.update({"family": fam, "year": "all"}); wfc.append(c)
    pd.DataFrame(wfc).to_csv(P.evaluation / "walkforward_contrasts.csv", index=False)

    # 7) reliability (test, chrono) for B2, B5, LB2, LB5
    rel = []
    for model in ("B2", "B5", "LB2", "LB5"):
        s = base.filter(pl.col("model") == model)
        if s.height:
            t = reliability_table(s["y"].to_numpy().astype(int), s["p"].to_numpy()); t["model"] = model; rel.append(t)
    if rel:
        pd.concat(rel).to_csv(P.evaluation / "reliability.csv", index=False)

    # 8) continuous target
    cont = []
    fitlog = json.load(open(P.evaluation / "baselines_fit_log.json"))["fits"]
    for model in ("B2", "B5"):
        s = preds.filter((pl.col("target") == "S_opinc") & (pl.col("model") == model) & (pl.col("tag") == "chrono_cont"))
        if s.height:
            rec = next(r for r in fitlog if r.get("tag") == "chrono_cont" and r["model"] == model)
            lo, hi = rec["winsor"]
            y = np.clip(s["y"].to_numpy(), lo, hi)
            cont.append({"model": model, **continuous_metrics(y, s["p"].to_numpy(), rec["train_mean"])})
    if cont:
        A = preds.filter((pl.col("target") == "S_opinc") & (pl.col("model") == "B5") & (pl.col("tag") == "chrono_cont"))
        B = preds.filter((pl.col("target") == "S_opinc") & (pl.col("model") == "B2") & (pl.col("tag") == "chrono_cont"))
        key = ["entity_id", "fiscal_period_end", "cutoff"]
        j = A.join(B.select(key + [pl.col("p").alias("pb")]), on=key)
        lo, hi = next(r for r in fitlog if r.get("tag") == "chrono_cont" and r["model"] == "B5")["winsor"]
        y = np.clip(j["y"].to_numpy(), lo, hi)
        d = (j["p"].to_numpy() - y) ** 2 - (j["pb"].to_numpy() - y) ** 2
        bb = paired_block_bootstrap(d, quarter_blocks(j["A_q"].to_numpy()), nb, seed)
        cont.append({"model": "B5-B2 d_mse", **bb})
    pd.DataFrame(cont).to_csv(P.evaluation / "continuous.csv", index=False)

    # 9) power (validation period, logistic B5 vs B2): MDE at 80 % power, alpha 5 %, quarter clustering
    v = preds.filter((pl.col("h") == hp) & (pl.col("target") == "y_adverse_opinc") & (pl.col("part") == "validation") & (pl.col("tag") == "chrono"))
    A, B = v.filter(pl.col("model") == "B5"), v.filter(pl.col("model") == "B2")
    power = {}
    if A.height and B.height:
        key = ["entity_id", "fiscal_period_end", "cutoff"]
        j = A.join(B.select(key + [pl.col("p").alias("pb")]), on=key)
        y = j["y"].to_numpy().astype(int)
        d = pointwise_logloss(y, j["p"].to_numpy()) - pointwise_logloss(y, j["pb"].to_numpy())
        bl = quarter_blocks(j["A_q"].to_numpy()); ub = np.unique(bl)
        dq = np.array([d[bl == b].mean() for b in ub]); se = dq.std(ddof=1) / np.sqrt(len(ub))
        n_test_q = preds.filter((pl.col("part") == "test") & (pl.col("tag") == "chrono") & (pl.col("h") == hp))["A_q"].to_numpy()
        T_test = len(np.unique(quarter_blocks(n_test_q))) if len(n_test_q) else 23
        se_test = dq.std(ddof=1) / np.sqrt(T_test)
        power = {"validation_quarters": int(len(ub)), "sd_quarterly_diff": float(dq.std(ddof=1)), "test_quarters_assumed": int(T_test),
                 "mde_logloss_80pct_power": float((1.96 + 0.84) * se_test), "validation_mean_diff": float(d.mean())}
    write_json(P.evaluation / "power.json", power)

    # copies to results/
    for f in ("metrics_all", "contrasts", "horizon_ladder", "secondary_targets", "firm_holdout", "walkforward_by_year",
              "walkforward_contrasts", "reliability", "continuous"):
        src = P.evaluation / f"{f}.csv"
        if src.exists():
            (P.results / f"{f}.csv").write_text(src.read_text())
    write_json(P.results / "evaluation_summary.json", {"run_at": now_iso(), "power": power, "h4": h4,
                                                      "primary": contrasts.to_dict(orient="records") if len(contrasts) else []})
    print(contrasts[["family", "a", "b", "n", "d_logloss", "ci_lo", "ci_hi", "gw_p"]].to_string() if len(contrasts) else "no contrasts")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())