#!/usr/bin/env python
"""Phase 5: negative controls, subgroup stress tests (H3), panel regressions (H3), pricing tests (H5)."""
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, paired_block_bootstrap, pointwise_logloss, quarter_blocks, two_way_cluster_ols  # noqa: E402
from wp39.models import Preprocessor, columns_for, predict_proba  # noqa: E402
from wp39.pipeline import Paths, get_logger, load_config  # noqa: E402
from wp39.portfolio import car_window, download_french, factor_alpha, monthly_decile_portfolios, read_french_daily  # noqa: E402
from wp39.provenance import now_iso, write_json  # noqa: E402
from wp39.robustness import add_terciles, shuffle_labels_within_quarter, subgroup_masks  # noqa: E402
from wp39.samples import feature_columns  # noqa: E402
from wp39.splits import chronological  # noqa: E402

KEY = ["entity_id", "fiscal_period_end", "cutoff"]
FIXED = {"num_leaves": 15, "learning_rate": 0.03, "min_child_samples": 100, "feature_fraction": 0.7, "bagging_fraction": 0.8,
         "bagging_freq": 1, "lambda_l2": 1.0, "n_estimators": 600, "verbose": -1, "n_jobs": 12}


def fixed_fit(train, test, cols, target, seed):
    import lightgbm as lgb
    pp = Preprocessor(cols).fit(train)
    cut = train["cutoff"].quantile(0.85)
    tr, va = train.filter(pl.col("cutoff") < cut), train.filter(pl.col("cutoff") >= cut)
    m = lgb.LGBMClassifier(random_state=seed, **FIXED)
    m.fit(pp.transform(tr), tr[target].to_numpy().astype(int), eval_set=[(pp.transform(va), va[target].to_numpy().astype(int))],
          callbacks=[lgb.early_stopping(50, verbose=False)])
    return predict_proba(m, pp.transform(test))


def paired(A: pl.DataFrame, B: pl.DataFrame, y: str, nb: int, seed: int) -> dict:
    j = A.select(KEY + [y, "p", "A_q"]).join(B.select(KEY + [pl.col("p").alias("pb")]), on=KEY)
    yy = j[y].to_numpy().astype(int)
    d = pointwise_logloss(yy, j["p"].to_numpy()) - pointwise_logloss(yy, j["pb"].to_numpy())
    r = paired_block_bootstrap(d, quarter_blocks(j["A_q"].to_numpy()), nb, seed)
    r["auc_a"] = binary_metrics(yy, j["p"].to_numpy()).get("auc"); r["auc_b"] = binary_metrics(yy, j["pb"].to_numpy()).get("auc")
    r["base_rate"] = float(yy.mean())
    return r


def main() -> int:
    P = Paths.default()
    log = get_logger("08_robust", P.logs / "08_robust.log")
    cfg = load_config("research"); hp, nb, seed = cfg["cutoffs"]["primary_horizon_td"], cfg["inference"]["n_bootstrap"], cfg["seed"]
    y = "y_adverse_opinc"
    df = pl.read_parquet(P.processed / "sample_event.parquet").filter((pl.col("h") == hp) & pl.col(y).is_not_null())
    fc = feature_columns(df)
    sp = chronological(df, cfg["splits"])
    trv = pl.concat([sp["train"], sp["validation"]])

    # ---------------------------------------------------------------- placebo 1: labels shuffled within calendar quarter
    plc = []
    for s in range(3):
        shuf_tr = shuffle_labels_within_quarter(trv, y, seed + s); shuf_te = shuffle_labels_within_quarter(sp["test"], y, seed + 100 + s)
        A = shuf_te.with_columns(pl.Series("p", fixed_fit(shuf_tr, shuf_te, columns_for("B5", fc), y, seed)))
        B = shuf_te.with_columns(pl.Series("p", fixed_fit(shuf_tr, shuf_te, columns_for("B2", fc), y, seed)))
        r = paired(A, B, y, nb, seed); r.update({"placebo": "shuffled_labels", "draw": s}); plc.append(r)
        log.info("placebo shuffled labels draw %d: d=%.5f [%.5f, %.5f] auc5=%.3f", s, r["mean"], r["ci_lo"], r["ci_hi"], r["auc_a"])
    pd.DataFrame(plc).to_csv(P.evaluation / "placebo_shuffled.csv", index=False)

    # ---------------------------------------------------------------- subgroups on the chronological LightGBM test predictions
    preds = {m: pl.read_parquet(P.evaluation / "preds" / f"event_h{hp}_{y}_{m}_test_chrono.parquet") for m in ("LB2", "LB5", "B2", "B5")}
    tr_t, te_t = add_terciles(sp["train"], [sp["test"]])
    te_t = te_t.with_columns([
        ((pl.col("cutoff") - pl.col("A_prev")).dt.total_days() >= 35).fill_null(True).alias("window_clean"),
        (pl.col("w1m_present_share") >= 0.98).alias("dense_bars"),
    ])
    extra = {"window_clean": pl.col("window_clean"), "window_straddles_prev_release": ~pl.col("window_clean"),
             "dense_bars": pl.col("dense_bars"), "sparse_bars": ~pl.col("dense_bars")}
    sub = []
    for fam, a, b in (("lightgbm", "LB5", "LB2"), ("logistic", "B5", "B2")):
        for name, mask in {**subgroup_masks(te_t), **extra}.items():
            ids = te_t.filter(mask).select(KEY)
            A, B = preds[a].join(ids, on=KEY), preds[b].join(ids, on=KEY)
            if A.height < 500:
                continue
            r = paired(A, B, y, nb, seed); r.update({"family": fam, "subgroup": name}); sub.append(r)
    subdf = pd.DataFrame(sub); subdf.to_csv(P.evaluation / "subgroups.csv", index=False)
    log.info("subgroups:\n%s", subdf[subdf.family == "lightgbm"][["subgroup", "n", "mean", "ci_lo", "ci_hi", "auc_a", "auc_b"]].to_string())

    # ---------------------------------------------------------------- H3 panel regressions on walk-forward OOF signals
    oof5 = pl.read_parquet(P.evaluation / "preds" / f"event_h{hp}_{y}_LB5_oof_walkforward.parquet").select(KEY + [pl.col("p").alias("lfds5")])
    oof2 = pl.read_parquet(P.evaluation / "preds" / f"event_h{hp}_{y}_LB2_oof_walkforward.parquet").select(KEY + [pl.col("p").alias("lfds2")])
    pan = df.join(oof5, on=KEY).join(oof2, on=KEY)
    pan = pan.with_columns([
        (pl.col("cutoff").dt.year().cast(pl.Utf8) + "Q" + ((pl.col("cutoff").dt.month() - 1) // 3 + 1).cast(pl.Utf8)).alias("cq"),
    ])
    pan = pan.with_columns((pl.col("sic_division") + "|" + pl.col("cq")).alias("sector_q"))
    # ex-ante cross-sectional terciles within calendar quarter
    for name, col in (("size", "f_log_mcap"), ("illiq", "wd_amihud60"), ("lev", "f_leverage")):
        pan = pan.with_columns(((pl.col(col).rank("average").over("cq") - 1) / pl.len().over("cq") * 3).floor().clip(0, 2).cast(pl.Int8).alias(f"t_{name}"))
    pd_ = pan.to_pandas()
    lo, hi = np.quantile(pd_["S_opinc"].dropna(), [0.01, 0.99]); pd_["S_w"] = pd_["S_opinc"].clip(lo, hi)
    pd_["d_lfds"] = pd_["lfds5"] - pd_["lfds2"]
    for name in ("size", "illiq", "lev"):
        for k in (0, 2):
            pd_[f"lfds5_x_{name}{k}"] = pd_["lfds5"] * (pd_[f"t_{name}"] == k)
            pd_[f"dlfds_x_{name}{k}"] = pd_["d_lfds"] * (pd_[f"t_{name}"] == k)
    regs = {}
    controls = ["f_roa", "f_leverage", "f_log_mcap", "wd_ret20", "wd_vol20", "last_surprise", "mean_surprise4"]
    regs["S_on_lfds5"] = two_way_cluster_ols(pd_, "S_w", ["lfds5"] + controls, "entity_id", "sector_q", "entity_id", "cq")
    regs["S_on_lfds2_and_dlfds"] = two_way_cluster_ols(pd_, "S_w", ["lfds2", "d_lfds"] + controls, "entity_id", "sector_q", "entity_id", "cq")
    regs["y_on_lfds5"] = two_way_cluster_ols(pd_.assign(yb=pd_[y].astype(float)), "yb", ["lfds5"] + controls, "entity_id", "sector_q", "entity_id", "cq")
    for name in ("size", "illiq", "lev"):
        regs[f"dlfds_x_{name}"] = two_way_cluster_ols(pd_, "S_w", ["lfds2", "d_lfds", f"dlfds_x_{name}0", f"dlfds_x_{name}2"] + controls,
                                                    "entity_id", "sector_q", "entity_id", "cq")
    allreg = pd.concat([r.assign(spec=k) for k, r in regs.items()]); allreg.to_csv(P.evaluation / "panel_regressions.csv", index=False)
    log.info("panel:\n%s", allreg[allreg.spec.isin(["S_on_lfds5", "S_on_lfds2_and_dlfds"])][["spec", "var", "coef", "se", "t", "p", "n"]].head(12).to_string())

    # ---------------------------------------------------------------- H5a: announcement-window CAR on the ex-ante signal
    ext = P.cache / "external"
    ff = read_french_daily(download_french(load_config("data_contract")["external"]["ff5_daily"], ext / "ff5_daily.zip"))
    mom = read_french_daily(download_french(load_config("data_contract")["external"]["mom_daily"], ext / "mom_daily.zip"))
    ff = ff.join(mom.iloc[:, :1].rename(columns={mom.columns[0]: "Mom"}), how="left")
    daily = pl.concat([pl.read_parquet(p).select(["symbol", "session_date", "d_ret_cc", "d_close", "d_dollar_vol"]) for p in (P.cache / "features").glob("*.parquet")])
    syms = set(pan["symbol"].unique().to_list())
    daily = daily.filter(pl.col("symbol").is_in(list(syms)))
    car = car_window(daily, ff, pan.select(["entity_id", "symbol", "cutoff", "A_q", "fiscal_period_end"]))
    pc = pd_.merge(car.to_pandas(), on=["entity_id", "symbol", "cutoff", "A_q", "fiscal_period_end"], how="inner")
    lo, hi = np.quantile(pc["car"].dropna(), [0.01, 0.99]); pc["car_w"] = pc["car"].clip(lo, hi)
    h5 = {"car_on_lfds5": two_way_cluster_ols(pc, "car_w", ["lfds5"] + controls, "entity_id", "sector_q", "entity_id", "cq"),
          "car_on_lfds2_dlfds": two_way_cluster_ols(pc, "car_w", ["lfds2", "d_lfds"] + controls, "entity_id", "sector_q", "entity_id", "cq")}
    pd.concat([r.assign(spec=k) for k, r in h5.items()]).to_csv(P.evaluation / "h5_car_regressions.csv", index=False)
    log.info("H5 CAR:\n%s", h5["car_on_lfds2_dlfds"][["var", "coef", "se", "t", "p", "n"]].head(3).to_string())

    # ---------------------------------------------------------------- H5b: monthly decile sorts (calendar design, OOF scores)
    port_out = {}
    for m in ("LB5", "LB2"):
        f = P.evaluation / "preds" / f"calendar_h0_{y}_{m}_oof_walkforward.parquet"
        if not f.exists():
            continue
        sig = pl.read_parquet(f).select(["symbol", "cutoff", pl.col("p").alias("score")]).to_pandas()
        res = monthly_decile_portfolios(sig, daily, ff, cost_bps=20.0)
        port = res["portfolios"]; port.to_csv(P.evaluation / f"portfolios_{m}.csv", index=False)
        port_out[m] = {"gross": factor_alpha(port, ff, 1, 10, net=False), "net20bps": factor_alpha(port, ff, 1, 10, net=True),
                       "decile_means": port.groupby("decile")["ret"].mean().round(5).to_dict(),
                       "avg_turnover": float(port["turnover"].mean())}
        log.info("portfolio %s: %s", m, port_out[m]["gross"])
    write_json(P.evaluation / "h5_portfolios.json", port_out)

    for f in ("placebo_shuffled", "subgroups", "panel_regressions", "h5_car_regressions"):
        (P.results / f"{f}.csv").write_text((P.evaluation / f"{f}.csv").read_text())
    write_json(P.results / "robustness_summary.json", {"run_at": now_iso(), "placebo": plc, "portfolios": port_out})
    return 0


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