"""
Train the stress head, and measure it honestly.

Two algorithms are fitted on identical folds — LightGBM and XGBoost — so the
choice between them is a recorded number rather than a preference. Validation is
an expanding-window time-series split: fold *k* trains on every as-of date before
a cut and tests on the dates after it, which is the only split that means
anything for a model whose job is to see forward.

Nothing here is calibrated against real defaults. The labels are generated DPD
history, and every metric this writes is stamped with that fact.

Outputs
-------
``models/stress_<algo>.pkl``   the fitted pipeline plus its feature order
``models/metrics.json``        fold-by-fold and pooled metrics
"""

from __future__ import annotations

import argparse
import json
import pickle
from datetime import datetime
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.metrics import (
    average_precision_score,
    brier_score_loss,
    precision_score,
    recall_score,
    roc_auc_score,
)

import lightgbm as lgb
import xgboost as xgb

# Mirrors udyam.risk.red-threshold / watch-threshold in application.properties.
# The console, the apex suppression rules and the Hindi PWA all speak these three
# bands; the model maps into them rather than inventing a fourth.
RED_THRESHOLD = 65
WATCH_THRESHOLD = 35


def load(data_dir: Path) -> tuple[pd.DataFrame, dict]:
    df = pd.read_parquet(data_dir / "training.parquet")
    manifest = json.loads((data_dir / "manifest.json").read_text(encoding="utf-8"))
    return df, manifest


def feature_columns(df: pd.DataFrame, manifest: dict) -> list[str]:
    """
    Feature order is fixed here and written into the artifact. A model served its
    columns in a different order does not error — it returns confident nonsense.
    """
    cols = list(manifest["features"])
    cols += [f"{c}__missing" for c in manifest["features"]]
    cols += ["dpd_at_t", "outstanding_at_t", "emi_at_t", "sector_id", "district_id"]
    return [c for c in cols if c in df.columns]


def expanding_folds(dates: np.ndarray, n_folds: int = 4, min_train: int = 8):
    """
    Expanding-window folds over the distinct as-of dates.

    Never a random split. A random split over a panel puts March and June of the
    same enterprise on opposite sides of the fence, and the model learns that
    enterprise rather than the pattern — which reports as excellent accuracy and
    reproduces as nothing.
    """
    uniq = np.array(sorted(pd.unique(dates)))
    if len(uniq) < min_train + n_folds:
        n_folds = max(1, len(uniq) - min_train)
    step = max(1, (len(uniq) - min_train) // max(1, n_folds))
    for k in range(n_folds):
        cut = min_train + k * step
        if cut >= len(uniq):
            break
        yield uniq[:cut], uniq[cut : cut + step] if k < n_folds - 1 else uniq[cut:]


def fit_lightgbm(x: pd.DataFrame, y: np.ndarray, pos_weight: float):
    model = lgb.LGBMClassifier(
        objective="binary",
        n_estimators=250,
        learning_rate=0.04,
        num_leaves=7,
        min_child_samples=45,
        subsample=0.85,
        subsample_freq=1,
        colsample_bytree=0.85,
        reg_lambda=5.0,
        scale_pos_weight=pos_weight,
        verbose=-1,
        random_state=42,
    )
    model.fit(x, y)
    return model


def fit_xgboost(x: pd.DataFrame, y: np.ndarray, pos_weight: float):
    model = xgb.XGBClassifier(
        objective="binary:logistic",
        n_estimators=400,
        learning_rate=0.04,
        max_depth=4,
        min_child_weight=5,
        subsample=0.85,
        colsample_bytree=0.85,
        reg_lambda=1.0,
        scale_pos_weight=pos_weight,
        tree_method="hist",
        enable_categorical=True,
        eval_metric="aucpr",
        random_state=42,
    )
    model.fit(x, y)
    return model


FITTERS = {"lightgbm": fit_lightgbm, "xgboost": fit_xgboost}


def score_to_band(score: int) -> str:
    if score >= RED_THRESHOLD:
        return "RED"
    if score >= WATCH_THRESHOLD:
        return "WATCH"
    return "STANDARD"


def evaluate_fold(y_true: np.ndarray, proba: np.ndarray) -> dict:
    """
    Precision at the Red band is the number that matters operationally: the cost
    of a false positive there is one field visit, and an officer who is sent on
    three wasted visits stops believing the fourth alert.
    """
    scores = np.clip(np.round(proba * 100), 0, 100)
    red = scores >= RED_THRESHOLD
    flagged = red | (scores >= WATCH_THRESHOLD)

    out = {
        "n": int(len(y_true)),
        "positives": int(y_true.sum()),
        "auc_roc": float(roc_auc_score(y_true, proba)) if y_true.sum() else None,
        "pr_auc": float(average_precision_score(y_true, proba)) if y_true.sum() else None,
        "brier": float(brier_score_loss(y_true, proba)),
        "red_flagged": int(red.sum()),
        "precision_red": float(precision_score(y_true, red, zero_division=0)),
        "recall_red": float(recall_score(y_true, red, zero_division=0)),
        "recall_flagged": float(recall_score(y_true, flagged, zero_division=0)),
    }
    return out


def run(data_dir: Path, out_dir: Path) -> dict:
    df, manifest = load(data_dir)
    cols = feature_columns(df, manifest)
    dates = df["as_of_date"].values

    results: dict = {
        "dataset": {
            "rows": manifest["rows"],
            "positives": manifest["positives"],
            "positive_rate": manifest["positive_rate"],
            "as_of_dates": manifest["as_of_dates"],
            "synthetic_label_share": manifest["synthetic_label_share"],
        },
        "validation": "expanding-window time-series split over as-of dates",
        "feature_order": cols,
        "algorithms": {},
        "caveats": manifest["caveats"],
    }

    for name, fitter in FITTERS.items():
        folds = []
        for train_dates, test_dates in expanding_folds(dates):
            tr = df[df["as_of_date"].isin(train_dates)]
            te = df[df["as_of_date"].isin(test_dates)]
            if te.empty or tr["label"].sum() == 0:
                continue

            pos_weight = float((tr["label"] == 0).sum() / max(1, (tr["label"] == 1).sum()))
            model = fitter(tr[cols], tr["label"].values, pos_weight)
            proba = model.predict_proba(te[cols])[:, 1]

            fold = evaluate_fold(te["label"].values, proba)
            fold["train_to"] = str(max(train_dates))
            fold["test_from"] = str(min(test_dates))
            folds.append(fold)

        pooled = {
            k: round(float(np.mean([f[k] for f in folds if f[k] is not None])), 4)
            for k in ("auc_roc", "pr_auc", "brier", "precision_red", "recall_red", "recall_flagged")
            if any(f[k] is not None for f in folds)
        }
        results["algorithms"][name] = {"folds": folds, "pooled": pooled}
        print(f"{name:10s} AUC {pooled.get('auc_roc', float('nan')):.3f}  "
              f"PR-AUC {pooled.get('pr_auc', float('nan')):.3f}  "
              f"precision@RED {pooled.get('precision_red', float('nan')):.3f}")

    # Baselines, on the identical folds.
    #
    # Without these a metric is unreadable: AUC 0.74 is either a real result or a
    # restatement of "a borrower who is late today is late in three months",
    # and only the comparison tells you which. The Reports page charts the model
    # against these rather than against an invented bureau score.
    results["baselines"] = {}
    for label, column in (("repayment_history_only", "dpd_at_t"),
                          ("obligation_ratio_only", "emi_to_inflow")):
        folds = []
        for train_dates, test_dates in expanding_folds(dates):
            te = df[df["as_of_date"].isin(test_dates)]
            if te.empty or te["label"].sum() == 0:
                continue
            raw = te[column].fillna(te[column].median()).values.astype(float)
            # Rank-scale to [0,1] so it is comparable to a probability.
            proba = (pd.Series(raw).rank(pct=True)).values
            folds.append(evaluate_fold(te["label"].values, proba))
        if folds:
            results["baselines"][label] = {
                k: round(float(np.mean([f[k] for f in folds if f[k] is not None])), 4)
                for k in ("auc_roc", "pr_auc", "precision_red")
                if any(f[k] is not None for f in folds)
            }
            print(f"{label:24s} AUC {results['baselines'][label].get('auc_roc', float('nan')):.3f}  "
                  f"PR-AUC {results['baselines'][label].get('pr_auc', float('nan')):.3f}")

    # The winner is chosen by PR-AUC, not AUC-ROC. At a 4% positive rate ROC is
    # flattered by the majority class; precision-recall is the honest read.
    winner = max(results["algorithms"], key=lambda a: results["algorithms"][a]["pooled"].get("pr_auc", 0))
    results["winner"] = winner
    results["base_rate"] = round(float(df["label"].mean()), 4)
    print(f"\nwinner: {winner} (by PR-AUC)")

    # Refit the winner on everything for serving.
    pos_weight = float((df["label"] == 0).sum() / max(1, (df["label"] == 1).sum()))
    final = FITTERS[winner](df[cols], df["label"].values, pos_weight)

    out_dir.mkdir(parents=True, exist_ok=True)
    with open(out_dir / "stress_model.pkl", "wb") as fh:
        pickle.dump({"model": final, "algorithm": winner, "features": cols,
                     "red_threshold": RED_THRESHOLD, "watch_threshold": WATCH_THRESHOLD}, fh)
    (out_dir / "metrics.json").write_text(json.dumps(results, indent=2), encoding="utf-8")
    print(f"written: {out_dir / 'stress_model.pkl'}")
    return results


DSN = "postgresql://postgres:root@localhost:5432/UdyamLinc_DEMO"

METRICS = ("auc_roc", "pr_auc", "precision_red", "recall_red", "brier")


def publish(results: dict, dsn: str = DSN) -> int:
    """
    Write the backtest to `model_metric`, where the console reads it.

    The Reports page used to render 0.87 from a constant nobody computed. These
    rows are the replacement, and each one carries the share of synthetic labels
    it was measured against — so the accuracy cannot be quoted without the
    caveat, because they are the same record.

    Baselines are written as rows too. A model AUC means nothing on its own; the
    single-feature arms are what say whether it is a result or a restatement.
    """
    import psycopg

    stamp = datetime.now()
    ds = results["dataset"]
    version = f"{results['winner']}-stress-v1"
    rows = []

    for arm, payload in results["algorithms"].items():
        for metric, value in payload["pooled"].items():
            if metric in METRICS:
                rows.append((version, arm, False, metric, value))

    for arm, payload in results.get("baselines", {}).items():
        for metric, value in payload.items():
            if metric in METRICS:
                rows.append((version, arm, True, metric, value))

    with psycopg.connect(dsn) as conn, conn.cursor() as cur:
        for version_, arm, is_baseline, metric, value in rows:
            cur.execute(
                """
                insert into model_metric
                  (model_version, arm, baseline, metric, value, dataset_rows,
                   dataset_positives, synthetic_label_share, validation, computed_at)
                values (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                """,
                (version_, arm, is_baseline, metric, round(float(value), 5),
                 ds["rows"], ds["positives"], ds["synthetic_label_share"],
                 results["validation"][:120], stamp),
            )
        conn.commit()

    print(f"published {len(rows)} metric rows at {stamp:%Y-%m-%d %H:%M}")
    return len(rows)


def main() -> None:
    ap = argparse.ArgumentParser(description="Train the UdyamLinc stress head")
    ap.add_argument("--data", default="data")
    ap.add_argument("--out", default="models")
    ap.add_argument("--no-publish", action="store_true",
                    help="train and evaluate without writing to model_metric")
    args = ap.parse_args()
    results = run(Path(args.data), Path(args.out))
    if not args.no_publish:
        publish(results)


if __name__ == "__main__":
    main()