#!/usr/bin/env python3
"""Train LightGBM ML scorer for PEAD long, save booster + feature manifest.

Outputs:
- data/ml_models/lgbm_v1_pead_long.txt           (LGBM booster, text format)
- data/ml_models/lgbm_v1_pead_long_manifest.json (feature schema + category vocab)
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd
import pyarrow.parquet as pq
from lightgbm import LGBMClassifier

NUMERIC = [
    "reaction_day_return", "gap_size", "close_location", "volume_ratio_20d",
    "document_quality_score", "parse_confidence_overall", "oneoff_penalty",
    "market_cap_proxy", "avg_dollar_volume_20d",
    "pre_event_volatility_20d", "pre_event_rsi_14", "pre_event_bb_position",
    "pre_event_obv_slope_20d",
    "macro_vix", "macro_hy_spread", "macro_t10y2y", "prior_event_fwd5d",
    "lm_positive_pct", "lm_negative_pct", "lm_net_sentiment",
    "earnings_surprise_pct",
    "sue_lag_1_pct", "sue_lag_2_pct", "sue_lag_3_pct",
    "sue_hist_mean_4q", "sue_hist_mean_8q", "sue_hist_mean_12q",
    "sue_hist_pos_rate_4q", "sue_hist_pos_rate_12q",
    "sue_hist_latest_pct", "sue_hist_streak_pos",
    "peer_sector_event_count_365d", "peer_sector_surprise_median_365d",
    "peer_sector_surprise_mean_365d", "peer_sector_surprise_pos_rate_365d",
    "peer_relative_surprise_pct_365d",
    "peer_sector_sue_hist_mean_4q_median_365d",
    "peer_sector_sue_hist_mean_4q_mean_365d",
    "peer_relative_sue_hist_mean_4q_365d",
    "peer_sector_sue_hist_pos_rate_4q_mean_365d",
    "pre_event_short_ratio", "pre_event_sector_momentum_20d",
    "prior_catalyst_count_60d", "prior_catalyst_type_diversity_60d",
    "price_vs_sma20", "pre_event_momentum_20d",
]
CAT = ["event_type", "event_direction", "guidance_status", "sector"]
SNAPSHOT = Path("data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_canonical_ftb_fix_v2")
OUT_DIR = Path("data/ml_models")


def load(split: str) -> pd.DataFrame:
    schema = pq.read_schema(SNAPSHOT / f"{split}.parquet")
    cols = [c for c in NUMERIC + CAT + ["mfe_10d", "mae_10d", "fwd_return_10d"] if c in schema.names]
    df = pq.read_table(SNAPSHOT / f"{split}.parquet", columns=cols).to_pandas()
    for c in NUMERIC:
        if c not in df.columns:
            df[c] = np.nan
        df[c] = pd.to_numeric(df[c], errors="coerce")
    for c in CAT:
        if c not in df.columns:
            df[c] = "missing"
        df[c] = df[c].fillna("missing").astype(str)
    df["target"] = ((df["mfe_10d"] - df["mae_10d"]) > 0.08).astype(int)
    return df


def main() -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    train = load("train")

    cat_vocab = {c: sorted(train[c].dropna().unique().tolist()) for c in CAT}
    # encode categories as integer codes for stable serialisation
    for c in CAT:
        codes = {v: i for i, v in enumerate(cat_vocab[c])}
        train[c] = train[c].map(codes).fillna(-1).astype("int32")
    feats = NUMERIC + CAT
    model = LGBMClassifier(
        n_estimators=400, learning_rate=0.03, num_leaves=31,
        min_child_samples=50, subsample=0.8, colsample_bytree=0.8,
        objective="binary", random_state=42, verbosity=-1,
    )
    model.fit(train[feats], train["target"], categorical_feature=CAT)

    booster = model.booster_
    model_path = OUT_DIR / "lgbm_v1_pead_long.txt"
    booster.save_model(str(model_path))

    manifest = {
        "numeric_features": NUMERIC,
        "categorical_features": CAT,
        "categorical_vocab": cat_vocab,
        "snapshot": str(SNAPSHOT),
        "n_estimators": 400,
        "lgbm_version": __import__("lightgbm").__version__,
    }
    manifest_path = OUT_DIR / "lgbm_v1_pead_long_manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2))

    # quick eval: re-predict on train + valid + test
    for split in ["train", "valid", "test"]:
        df = load(split)
        for c in CAT:
            codes = {v: i for i, v in enumerate(cat_vocab[c])}
            df[c] = df[c].map(codes).fillna(-1).astype("int32")
        prob = model.predict_proba(df[feats])[:, 1]
        print(f"{split}: mean={prob.mean():.4f} std={prob.std():.4f} p10={np.percentile(prob,10):.3f} p90={np.percentile(prob,90):.3f}")

    print(f"\nSaved booster: {model_path}")
    print(f"Saved manifest: {manifest_path}")


if __name__ == "__main__":
    main()