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

import argparse
import json
from datetime import datetime, timezone
from pathlib import Path

import pandas as pd
import pyarrow.parquet as pq
from lightgbm import LGBMClassifier, LGBMRegressor
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler


SNAPSHOT_ROOT = Path("data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_canonical")
NUMERIC_FEATURES = [
    "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_entropy_60d",
    "pre_event_gravitational_pull",
    "pre_event_hurst_60d",
    "pre_event_market_temperature",
    "pre_event_ou_theta_60d",
    "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_lag_4_pct",
    "sue_lag_5_pct",
    "sue_lag_6_pct",
    "sue_lag_7_pct",
    "sue_lag_8_pct",
    "sue_lag_9_pct",
    "sue_lag_10_pct",
    "sue_lag_11_pct",
    "sue_lag_12_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",
]
CAT_FEATURES = ["event_type", "event_direction", "guidance_status", "sector"]
ALL_COLUMNS = NUMERIC_FEATURES + CAT_FEATURES + [
    "reported_eps",
    "estimated_eps",
    "fwd_return_10d",
    "fwd_return_20d",
    "mfe_10d",
    "mae_10d",
]


def _availability_report(columns: set[str]) -> dict[str, dict[str, object]]:
    requested = {
        "v13e_scoring_core": {
            "required": [
                "reaction_day_return",
                "gap_size",
                "close_location",
                "volume_ratio_20d",
                "document_quality_score",
                "parse_confidence_overall",
                "oneoff_penalty",
            ],
            "nice_to_have": [
                "market_cap_proxy",
                "avg_dollar_volume_20d",
            ],
        },
        "historical_12q_sue_vector": {
            "required": [],
            "nice_to_have": [
                "reported_eps",
                "estimated_eps",
                "earnings_surprise_pct",
                "sue_lag_1_pct",
                "sue_lag_2_pct",
                "sue_lag_3_pct",
                "sue_lag_4_pct",
                "sue_lag_5_pct",
                "sue_lag_6_pct",
                "sue_lag_7_pct",
                "sue_lag_8_pct",
                "sue_lag_9_pct",
                "sue_lag_10_pct",
                "sue_lag_11_pct",
                "sue_lag_12_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",
            ],
        },
        "transcript_sentiment_surprise": {
            "required": [],
            "nice_to_have": ["lm_positive_pct", "lm_negative_pct", "lm_net_sentiment"],
        },
        "three_day_early_reaction": {
            "required": [],
            "nice_to_have": [],
        },
        "insider_form4_persistence": {
            "required": [],
            "nice_to_have": [],
        },
        "peer_relative_surprise": {
            "required": [],
            "nice_to_have": [
                "sector",
                "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",
            ],
        },
        "options_skew_iv_term_structure": {
            "required": [],
            "nice_to_have": [],
        },
        "technical_macro_context": {
            "required": [
                "pre_event_entropy_60d",
                "pre_event_hurst_60d",
                "pre_event_gravitational_pull",
                "pre_event_market_temperature",
                "pre_event_ou_theta_60d",
                "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",
            ],
            "nice_to_have": [],
        },
    }
    report: dict[str, dict[str, object]] = {}
    for name, spec in requested.items():
        present = [c for c in spec["required"] + spec["nice_to_have"] if c in columns]
        missing_required = [c for c in spec["required"] if c not in columns]
        missing_optional = [c for c in spec["nice_to_have"] if c not in columns]
        if spec["required"] and not missing_required:
            status = "available"
        elif present:
            status = "partial"
        else:
            status = "missing"
        report[name] = {
            "status": status,
            "present": present,
            "missing_required": missing_required,
            "missing_optional": missing_optional,
        }
    return report


def _load_split(snapshot_root: Path, split: str) -> pd.DataFrame:
    parquet_path = snapshot_root / f"{split}.parquet"
    schema = pq.read_schema(parquet_path)
    present = [column for column in ALL_COLUMNS if column in schema.names]
    df = pq.read_table(parquet_path, columns=present).to_pandas()
    for column in ALL_COLUMNS:
        if column not in df.columns:
            df[column] = pd.NA
    df["target_reg_10d"] = df["fwd_return_10d"]
    df["target_cls_strong_drift_10d"] = ((df["mfe_10d"] - df["mae_10d"]) > 0.08).astype(int)
    for col in CAT_FEATURES:
        df[col] = df[col].fillna("missing").astype("category")
    for col in NUMERIC_FEATURES:
        df[col] = pd.to_numeric(df[col], errors="coerce")
    return df


def _bucket_summary(pred: pd.Series, target: pd.Series) -> dict[str, float | None]:
    out = pd.DataFrame({"pred": pred, "target": target}).dropna()
    if out.empty:
        return {
            "mean_target": None,
            "top_quintile_mean": None,
            "bottom_quintile_mean": None,
            "spread": None,
            "corr": None,
        }
    out["bucket"] = pd.qcut(out["pred"], 5, labels=False, duplicates="drop")
    top = out[out["bucket"] == out["bucket"].max()]["target"].mean()
    bottom = out[out["bucket"] == out["bucket"].min()]["target"].mean()
    corr = out["pred"].corr(out["target"])
    return {
        "mean_target": round(float(out["target"].mean()), 4),
        "top_quintile_mean": round(float(top), 4),
        "bottom_quintile_mean": round(float(bottom), 4),
        "spread": round(float(top - bottom), 4),
        "corr": None if pd.isna(corr) else round(float(corr), 4),
    }


def _top_decile_summary(pred: pd.Series, target: pd.Series, extra: pd.Series | None = None) -> dict[str, float | int | None]:
    out = pd.DataFrame({"pred": pred, "target": target}).dropna()
    if extra is not None:
        out["extra"] = extra.loc[out.index]
    if out.empty:
        return {"count": 0, "hit_rate": None, "extra_mean": None}
    threshold = out["pred"].quantile(0.9)
    top = out[out["pred"] >= threshold]
    payload: dict[str, float | int | None] = {
        "count": int(len(top)),
        "hit_rate": round(float(top["target"].mean()), 4),
    }
    if "extra" in top:
        payload["extra_mean"] = round(float(top["extra"].mean()), 4)
    return payload


def run_probe(snapshot_root: Path) -> dict[str, object]:
    train = _load_split(snapshot_root, "train")
    valid = _load_split(snapshot_root, "valid")
    test = _load_split(snapshot_root, "test")

    availability = _availability_report(set(train.columns))

    num_pipe = Pipeline(
        [
            ("impute", SimpleImputer(strategy="median")),
            ("scale", StandardScaler()),
        ]
    )
    cat_pipe = Pipeline(
        [
            ("impute", SimpleImputer(strategy="most_frequent")),
            ("onehot", OneHotEncoder(handle_unknown="ignore")),
        ]
    )
    prep = ColumnTransformer(
        [
            ("num", num_pipe, NUMERIC_FEATURES),
            ("cat", cat_pipe, CAT_FEATURES),
        ]
    )
    logit = Pipeline(
        [
            ("prep", prep),
            ("model", LogisticRegression(max_iter=2000, class_weight="balanced")),
        ]
    )
    logit.fit(train[NUMERIC_FEATURES + CAT_FEATURES], train["target_cls_strong_drift_10d"])

    lgbm_cls = 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,
    )
    lgbm_cls.fit(
        train[NUMERIC_FEATURES + CAT_FEATURES],
        train["target_cls_strong_drift_10d"],
        categorical_feature=CAT_FEATURES,
    )

    lgbm_reg = LGBMRegressor(
        n_estimators=300,
        learning_rate=0.05,
        num_leaves=31,
        min_child_samples=40,
        subsample=0.8,
        colsample_bytree=0.8,
        objective="regression",
        random_state=42,
        verbosity=-1,
    )
    lgbm_reg.fit(
        train[NUMERIC_FEATURES + CAT_FEATURES],
        train["target_reg_10d"],
        categorical_feature=CAT_FEATURES,
    )

    report: dict[str, object] = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "snapshot_root": str(snapshot_root),
        "target_notes": {
            "classification": "strong_drift_10d := (mfe_10d - mae_10d) > 8%",
            "regression": "fwd_return_10d used as proxy for direct return prediction; 14d label is not materialized yet",
        },
        "feature_availability": availability,
        "models": {},
    }

    for split_name, df in [("valid", valid), ("test", test)]:
        X = df[NUMERIC_FEATURES + CAT_FEATURES]
        cls_target = df["target_cls_strong_drift_10d"]
        reg_target = df["target_reg_10d"]

        logit_prob = pd.Series(logit.predict_proba(X)[:, 1], index=df.index)
        lgbm_prob = pd.Series(lgbm_cls.predict_proba(X)[:, 1], index=df.index)
        reg_pred = pd.Series(lgbm_reg.predict(X), index=df.index)

        report["models"].setdefault("logistic_strong_drift_10d", {})[split_name] = {
            "auc": round(float(roc_auc_score(cls_target, logit_prob)), 4),
            "bucket_summary": _bucket_summary(logit_prob, cls_target),
            "top_decile": _top_decile_summary(logit_prob, cls_target, reg_target),
        }
        report["models"].setdefault("lightgbm_strong_drift_10d", {})[split_name] = {
            "auc": round(float(roc_auc_score(cls_target, lgbm_prob)), 4),
            "bucket_summary": _bucket_summary(lgbm_prob, cls_target),
            "top_decile": _top_decile_summary(lgbm_prob, cls_target, reg_target),
        }
        report["models"].setdefault("lightgbm_return_10d_regression", {})[split_name] = {
            "bucket_summary": _bucket_summary(reg_pred, reg_target),
            "top_decile": _top_decile_summary(reg_pred, cls_target, reg_target),
        }

    importances = dict(
        sorted(
            zip(NUMERIC_FEATURES + CAT_FEATURES, lgbm_cls.feature_importances_.tolist(), strict=False),
            key=lambda x: x[1],
            reverse=True,
        )[:12]
    )
    report["models"]["lightgbm_strong_drift_10d"]["top_feature_importances"] = importances
    return report


def main() -> None:
    parser = argparse.ArgumentParser(description="Feasibility probe for full-ML predictor replacement.")
    parser.add_argument(
        "--snapshot-root",
        default=str(SNAPSHOT_ROOT),
        help="Canonical snapshot root containing train/valid/test parquet files.",
    )
    parser.add_argument(
        "--output",
        default="runs/ml_predictor_feasibility_20260401.json",
        help="Path to write JSON report.",
    )
    args = parser.parse_args()

    snapshot_root = Path(args.snapshot_root)
    report = run_probe(snapshot_root)

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(report, indent=2))
    print(output_path)


if __name__ == "__main__":
    main()