#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
train_scene_score — gelernter Szenen-Score vs Max-Confidence-Baseline (Phase A §8h-Folge).

§8h hat gezeigt: der Limiter des Gates ist die Score-Aggregation, nicht Auflösung/
Vokabular. Dieser Trainer prüft, ob ein GELERNTER Score (logistische Regression
über die Output-Features aus ml/scene_gate/features.py) die Pareto-Front
gegenüber der heutigen `waste_likelihood = max(Müll-Confidence)` bewegt.

Bewusst dependency-leicht (nur numpy): eine lineare Logistik ist der richtige
erste Test — interpretierbar, und wenn sie die Baseline nicht schlägt, ist der
Zero-Shot-Score nachweislich gedeckelt (→ Feld-Daten + echter Kopf statt Serving).

Eingabe: features_scene_gate.jsonl (eine Zeile je Bild: {has_waste, <FEATURE_NAMES>})
aus dem Kaggle-Kernel (collect_features=True). Bewertung fair aus-of-fold (CV),
damit der Score nicht auf denselben Bildern gelernt und gemessen wird.

Nutzung:
    python ml/scripts/train_scene_score.py --features features_scene_gate.jsonl \
        --out scene_score_report.json
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from typing import Dict, List, Sequence, Tuple

import numpy as np

REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO))
from ml.scene_gate.features import FEATURE_NAMES  # noqa: E402


# ---------- reine, testbare Metrik ----------

def pareto_tp_at_fpkill(scores: Sequence[float], y: Sequence[int],
                        targets: Sequence[float]) -> Dict[float, float]:
    """Für jedes FP-Kill-Ziel die höchste erreichbare TP-Retention.

    Bei Schwelle t (behalte score >= t): TP-Retention = Anteil Positiver mit
    score>=t; FP-Kill = Anteil Negativer mit score<t. Gibt {target: max_tp|None}."""
    s = np.asarray(scores, dtype=float)
    yv = np.asarray(y, dtype=int)
    pos = s[yv == 1]
    neg = s[yv == 0]
    n_pos = len(pos)
    n_neg = len(neg)
    thresholds = np.unique(np.concatenate([s, [s.min() - 1, s.max() + 1]]))
    out: Dict[float, float] = {}
    for tgt in targets:
        best = None
        for t in thresholds:
            tp = float((pos >= t).mean()) if n_pos else 0.0
            fk = float((neg < t).mean()) if n_neg else 0.0
            if fk >= tgt and (best is None or tp > best):
                best = tp
        out[float(tgt)] = best
    return out


def auc(scores: Sequence[float], y: Sequence[int]) -> float:
    """ROC-AUC über die Rang-Statistik (Mann-Whitney-U). 0.5 = Zufall."""
    s = np.asarray(scores, dtype=float)
    yv = np.asarray(y, dtype=int)
    n_pos = int((yv == 1).sum())
    n_neg = int((yv == 0).sum())
    if n_pos == 0 or n_neg == 0:
        return 0.5
    order = np.argsort(s, kind="mergesort")
    ranks = np.empty(len(s), dtype=float)
    ranks[order] = np.arange(1, len(s) + 1)
    # Ties: Durchschnittsrang
    _, inv, counts = np.unique(s, return_inverse=True, return_counts=True)
    cum = np.cumsum(counts)
    avg = {}
    start = 0
    for i, c in enumerate(counts):
        avg[i] = (start + 1 + start + c) / 2.0
        start += c
    ranks = np.array([avg[i] for i in inv])
    sum_pos = ranks[yv == 1].sum()
    return float((sum_pos - n_pos * (n_pos + 1) / 2.0) / (n_pos * n_neg))


# ---------- logistische Regression (numpy) ----------

def _standardize(X: np.ndarray, mean: np.ndarray, std: np.ndarray) -> np.ndarray:
    return (X - mean) / std


def fit_logistic(X: np.ndarray, y: np.ndarray, epochs: int = 500, lr: float = 0.1,
                 l2: float = 1.0) -> np.ndarray:
    """Batch-Gradientenabstieg, mit Bias-Spalte. Gibt Gewichte (inkl. Bias)."""
    n, d = X.shape
    Xb = np.hstack([np.ones((n, 1)), X])
    w = np.zeros(d + 1)
    for _ in range(epochs):
        z = Xb @ w
        p = 1.0 / (1.0 + np.exp(-np.clip(z, -30, 30)))
        grad = Xb.T @ (p - y) / n
        grad[1:] += (l2 / n) * w[1:]          # L2, Bias ausgenommen
        w -= lr * grad
    return w


def predict_logistic(X: np.ndarray, w: np.ndarray) -> np.ndarray:
    Xb = np.hstack([np.ones((X.shape[0], 1)), X])
    return 1.0 / (1.0 + np.exp(-np.clip(Xb @ w, -30, 30)))


def stratified_folds(y: np.ndarray, k: int = 5, seed: int = 0) -> List[np.ndarray]:
    """k Index-Folds, Klassenanteile je Fold erhalten (deterministisch)."""
    rng = np.random.RandomState(seed)
    folds: List[List[int]] = [[] for _ in range(k)]
    for cls in (0, 1):
        idx = np.where(y == cls)[0]
        rng.shuffle(idx)
        for j, i in enumerate(idx):
            folds[j % k].append(int(i))
    return [np.array(sorted(f)) for f in folds]


def cv_oof_scores(X: np.ndarray, y: np.ndarray, k: int = 5, seed: int = 0,
                  **fit_kw) -> np.ndarray:
    """Out-of-fold-Wahrscheinlichkeiten — fair (nie auf demselben Bild gelernt+gemessen)."""
    oof = np.zeros(len(y))
    for f in stratified_folds(y, k, seed):
        mask = np.zeros(len(y), dtype=bool)
        mask[f] = True
        Xtr, ytr = X[~mask], y[~mask]
        mean = Xtr.mean(axis=0)
        std = Xtr.std(axis=0)
        std[std == 0] = 1.0
        w = fit_logistic(_standardize(Xtr, mean, std), ytr, **fit_kw)
        oof[mask] = predict_logistic(_standardize(X[mask], mean, std), w)
    return oof


# ---------- IO ----------

def load_features(path: str) -> Tuple[np.ndarray, np.ndarray]:
    X_rows, y_rows = [], []
    for line in Path(path).read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        r = json.loads(line)
        X_rows.append([float(r.get(k, 0.0)) for k in FEATURE_NAMES])
        y_rows.append(1 if r.get("has_waste") else 0)
    return np.array(X_rows, dtype=float), np.array(y_rows, dtype=int)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--features", required=True)
    ap.add_argument("--out", default="scene_score_report.json")
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("--save-model", default=None,
                    help="Pfad für das trainierte Modell-Artefakt (JSON) zum Serving.")
    ap.add_argument("--backbone", default="unknown",
                    help="Backbone, aus dessen Features trainiert wurde. WICHTIG: das "
                         "Serving MUSS dasselbe Backbone fahren — Features sind nicht übertragbar.")
    args = ap.parse_args()

    X, y = load_features(args.features)
    targets = [0.60, 0.70, 0.80, 0.90, 0.95]

    # Baseline = die heutige waste_likelihood (ov_waste_max-Spalte).
    base_col = FEATURE_NAMES.index("ov_waste_max")
    base_scores = X[:, base_col]
    learned = cv_oof_scores(X, y, k=args.folds)

    base_par = pareto_tp_at_fpkill(base_scores, y, targets)
    learn_par = pareto_tp_at_fpkill(learned, y, targets)

    # Gewichte auf allen Daten (nur zur Interpretation der Feature-Wichtigkeit).
    mean = X.mean(axis=0); std = X.std(axis=0); std[std == 0] = 1.0
    w_full = fit_logistic(_standardize(X, mean, std), y)
    importance = sorted(
        ({"feature": FEATURE_NAMES[i], "weight": round(float(w_full[i + 1]), 4)}
         for i in range(len(FEATURE_NAMES))),
        key=lambda d: -abs(d["weight"]),
    )

    report = {
        "n": int(len(y)), "n_pos": int(y.sum()), "n_neg": int((y == 0).sum()),
        "auc_baseline": round(auc(base_scores, y), 4),
        "auc_learned_oof": round(auc(learned, y), 4),
        "pareto": {
            str(t): {"baseline": base_par[t], "learned": learn_par[t],
                     "delta": (round(learn_par[t] - base_par[t], 4)
                               if base_par[t] is not None and learn_par[t] is not None else None)}
            for t in targets
        },
        "feature_importance": importance,
    }
    Path(args.out).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")

    if args.save_model:
        # Serving-Artefakt: standardisiertes lineares Logit. score = sigmoid(
        # bias + Σ w_i * (x_i - mean_i)/std_i). Reproduzierbar ohne sklearn.
        model = {
            "kind": "logistic_standardized",
            "feature_names": list(FEATURE_NAMES),
            "mean": [round(float(v), 6) for v in mean],
            "std": [round(float(v), 6) for v in std],
            "bias": round(float(w_full[0]), 6),
            "weights": [round(float(v), 6) for v in w_full[1:]],
            "backbone": args.backbone,
            "trained_on": {"n": report["n"], "n_pos": report["n_pos"], "n_neg": report["n_neg"]},
            "auc_oof": report["auc_learned_oof"],
            "note": ("Trainiert auf TACO-Val (Obergrenze-Indikation); auf Feld-Daten neu "
                     "trainieren. Features sind BACKBONE-SPEZIFISCH — das Serving muss "
                     f"'{args.backbone}' fahren, sonst passen die Verteilungen nicht."),
        }
        Path(args.save_model).write_text(json.dumps(model, ensure_ascii=False, indent=2), encoding="utf-8")
        print(f"-> Modell: {args.save_model}")

    print(f"n={report['n']} (pos {report['n_pos']} / neg {report['n_neg']})")
    print(f"AUC  Baseline={report['auc_baseline']}  Gelernt(OOF)={report['auc_learned_oof']}")
    print(f"{'FP-Kill≥':>9} | {'Baseline':>9} | {'Gelernt':>9} | Δ")
    for t in targets:
        p = report["pareto"][str(t)]
        b = f"{p['baseline']:.3f}" if p["baseline"] is not None else "  —  "
        l = f"{p['learned']:.3f}" if p["learned"] is not None else "  —  "
        d = f"{p['delta']:+.3f}" if p["delta"] is not None else ""
        print(f"{t:>9.2f} | {b:>9} | {l:>9} | {d}")
    print("Top-Features:", ", ".join(f"{d['feature']}({d['weight']:+.2f})" for d in importance[:5]))
    print(f"-> {args.out}")


if __name__ == "__main__":
    main()