#!/usr/bin/env python
"""Phase 4: LightGBM (B6) with an identical random-search budget for every information set on the chronological split
(primary horizon, primary target + secondary targets), horizon ladder and walk-forward refits with a fixed configuration
(no tuning inside folds, so earlier folds never see later data). No neural/text models (DECISIONS D-006)."""
from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import polars as pl

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from wp39.models import INFO_SETS, Preprocessor, columns_for, fit_lightgbm, predict_proba  # noqa: E402
from wp39.pipeline import Paths, get_logger, load_config  # noqa: E402
from wp39.provenance import now_iso, write_json  # noqa: E402
from wp39.samples import feature_columns  # noqa: E402
from wp39.splits import chronological, firm_holdout_mask, walk_forward  # noqa: E402

KEY = ["entity_id", "symbol", "fiscal_period_end", "h", "cutoff", "A_q", "sic_division", "release_time"]
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 save(P, df, p, design, h, target, model, part, tag, extra=None):
    out = df.select(KEY + [target]).with_columns([pl.Series("p", p), pl.lit(model).alias("model"), pl.lit(part).alias("part"), pl.lit(tag).alias("tag")])
    if extra is not None:
        out = out.with_columns(extra)
    out.write_parquet(P.evaluation / "preds" / f"{design}_h{h}_{target}_{model}_{part}_{tag}.parquet")


def tuned(P, cfg, mcfg, fc, design, h, target, sp, tag, sets, log, record):
    ytr, yva = sp["train"][target].to_numpy().astype(int), sp["validation"][target].to_numpy().astype(int)
    for name in sets:
        cols = columns_for(name, fc); pp = Preprocessor(cols).fit(sp["train"])
        m, info = fit_lightgbm(pp.transform(sp["train"]), ytr, pp.transform(sp["validation"]), yva, mcfg["lightgbm"], cfg["seed"], n_jobs=12)
        save(P, sp["validation"], predict_proba(m, pp.transform(sp["validation"])), design, h, target, "L" + name, "validation", tag)
        save(P, sp["test"], predict_proba(m, pp.transform(sp["test"])), design, h, target, "L" + name, "test", tag)
        imp = dict(zip(cols, (m.booster_.feature_importance("gain") / max(m.booster_.feature_importance("gain").sum(), 1)).round(4).tolist()))
        record.append({"design": design, "h": h, "target": target, "model": "L" + name, "tag": tag, "n_features": len(cols),
                       "val_score": info["val_score"], "params": info["params"], "trials": len(info["trials"]), "importance_gain": imp})
        log.info("%s h=%d %s L%s %s: val_ll=%.4f params=%s", design, h, target, name, tag, info["val_score"], info["params"])


def fixed_fit(train, test, cols, target, seed):
    import lightgbm as lgb
    pp = Preprocessor(cols).fit(train)
    # early stopping on the last 4 calendar quarters of the training window (purged by construction: still before test)
    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)), int(m.best_iteration_ or 0)


def main() -> int:
    P = Paths.default()
    log = get_logger("06_advanced", P.logs / "06_advanced.log")
    cfg, mcfg = load_config("research"), load_config("models")
    hp = cfg["cutoffs"]["primary_horizon_td"]
    record = []
    df = pl.read_parquet(P.processed / "sample_event.parquet")
    fc = feature_columns(df)
    # 1) primary: tuned LightGBM, all information sets, h = 20, primary target
    d = df.filter((pl.col("h") == hp) & pl.col("y_adverse_opinc").is_not_null())
    sp = chronological(d, cfg["splits"])
    tuned(P, cfg, mcfg, fc, "event", hp, "y_adverse_opinc", sp, "chrono", list(INFO_SETS), log, record)
    # 2) secondary targets (B2 vs B5 only), tuned
    for target in ("y_adverse_netinc", "y_roa_decline", "y_ocf_negative", "y_liquidity_deterioration", "y_abs_surprise"):
        dt_ = df.filter((pl.col("h") == hp) & pl.col(target).is_not_null())
        tuned(P, cfg, mcfg, fc, "event", hp, target, chronological(dt_, cfg["splits"]), "chrono", ["B2", "B5"], log, record)
    # 3) horizon ladder (fixed config): B2 vs B5 at every h on the chronological split
    for h in sorted(df["h"].unique().to_list()):
        dh = df.filter((pl.col("h") == h) & pl.col("y_adverse_opinc").is_not_null())
        sph = chronological(dh, cfg["splits"])
        for name in ("B2", "B5"):
            p, it = fixed_fit(pl.concat([sph["train"], sph["validation"]]), sph["test"], columns_for(name, fc), "y_adverse_opinc", cfg["seed"])
            save(P, sph["test"], p, "event", h, "y_adverse_opinc", "L" + name, "test", "fixed")
            record.append({"design": "event", "h": h, "target": "y_adverse_opinc", "model": "L" + name, "tag": "fixed", "best_iteration": it})
    # 4) firm hold-out (fixed config)
    hold = firm_holdout_mask(d["entity_id"], cfg["splits"]["firm_holdout_fraction"])
    d2 = d.with_columns(pl.Series("_hold", hold))
    sp2 = chronological(d2.filter(~pl.col("_hold")), cfg["splits"])
    te_hold = d2.filter(pl.col("_hold") & (pl.col("cutoff") >= sp2["test"]["cutoff"].min()))
    for name in ("B2", "B5"):
        trv = pl.concat([sp2["train"], sp2["validation"]])
        p, _ = fixed_fit(trv, te_hold, columns_for(name, fc), "y_adverse_opinc", cfg["seed"]); save(P, te_hold, p, "event", hp, "y_adverse_opinc", "L" + name, "test", "firmhold_heldout")
        p, _ = fixed_fit(trv, sp2["test"], columns_for(name, fc), "y_adverse_opinc", cfg["seed"]); save(P, sp2["test"], p, "event", hp, "y_adverse_opinc", "L" + name, "test", "firmhold_seen")
    # 5) walk-forward out-of-fold (fixed config) for B1..B5 at h = 20 -> LFDS for the panel regressions / H4 / H5
    for name in ("B1", "B2", "B3", "B4", "B5"):
        parts = []
        for y, tr, te in walk_forward(d, cfg["splits"], 2026):
            p, it = fixed_fit(tr, te, columns_for(name, fc), "y_adverse_opinc", cfg["seed"])
            parts.append(te.select(KEY + ["y_adverse_opinc"]).with_columns([pl.Series("p", p), pl.lit(y).alias("fold")]))
        oof = pl.concat(parts).with_columns([pl.lit("L" + name).alias("model"), pl.lit("oof").alias("part"), pl.lit("walkforward").alias("tag")])
        oof.write_parquet(P.evaluation / "preds" / f"event_h{hp}_y_adverse_opinc_L{name}_oof_walkforward.parquet")
        log.info("walk-forward L%s: %d oof rows", name, oof.height)
    # 6) calendar design: walk-forward OOF for B2/B5 (scores for H5 portfolio sorts)
    dc = pl.read_parquet(P.processed / "sample_calendar.parquet").with_columns(pl.lit(0).alias("h"))
    fcc = feature_columns(dc)
    for name in ("B2", "B5"):
        parts = []
        for y, tr, te in walk_forward(dc, cfg["splits"], 2026):
            p, _ = fixed_fit(tr, te, columns_for(name, fcc), "y_adverse_opinc", cfg["seed"])
            parts.append(te.select(KEY + ["y_adverse_opinc"]).with_columns([pl.Series("p", p), pl.lit(y).alias("fold")]))
        oof = pl.concat(parts).with_columns([pl.lit("L" + name).alias("model"), pl.lit("oof").alias("part"), pl.lit("walkforward").alias("tag")])
        oof.write_parquet(P.evaluation / "preds" / f"calendar_h0_y_adverse_opinc_L{name}_oof_walkforward.parquet")
    write_json(P.evaluation / "advanced_fit_log.json", {"run_at": now_iso(), "fixed_config": FIXED, "fits": record})
    print(json.dumps({"fits": len(record)}, indent=1))
    return 0


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