"""Обучение считывания по ТРЕКАМ и честная проверка по бэгам.

Покадровое считывание (`train_mbon.py`) отвечает на вопрос «похоже ли это
пятно на предмет». Здесь вопрос другой: «стоит ли по этому треку тормозить» —
и отвечать на него по одному кадру нечем. Воронка потерь (EXPERIMENTS п. 12)
показывает, где это болит: на 120…160 м трек заводится в 62 % наблюдений, а до
решения доходит 38 %. Улика есть, порога не набирает.

Главная таблица внизу — не AUC модели сам по себе, а **сравнение с уликой**.
Улика уже сейчас стоит в пороге, и если модель её не обгоняет, весь замысел
не стоит ни строчки кода: значит, в истории трека нет ничего сверх того, что
накопитель и так посчитал.

Проверка — leave-one-bag-out, как и везде: модель обучается на всех записях,
кроме проверяемой.

    python tools/make_benchmark.py --memory artifacts/mushroom_body.npz \
        --mbon-dir artifacts/mbon_folds --tracks-out data/cache/track_set.npz
    python tools/train_track.py --device cuda
"""
from __future__ import annotations

import argparse
import time
from pathlib import Path

import numpy as np

import _bootstrap as B  # noqa: F401
from _metrics import auc, fpr_at_tpr
from flyguard.mbon_readout import MbonConfig
from flyguard.track_readout import TRACK_FEATURES, TrackReadout

BANDS = ((0, 30), (30, 55), (55, 80), (80, 110), (110, 160), (160, 230))


def band_table(d, y, cols: dict[str, np.ndarray]) -> None:
    """AUC по полосам дальности для нескольких оценок разом."""
    head = "".join(f"{k:>16}" for k in cols)
    print(f"{'полоса':>12}{'наблюдений':>12}{'предметов':>11}{head}")
    for lo, hi in BANDS:
        m = (d >= lo) & (d < hi)
        if m.sum() < 30 or len(set(y[m].tolist())) < 2:
            continue
        cells = "".join(f"{auc(v[m], y[m]):16.3f}" for v in cols.values())
        print(f"{f'{lo}-{hi} м':>12}{int(m.sum()):12d}{int(y[m].sum()):11d}{cells}")


def loo(names, X, Y, cfg, args):
    """Обучение на всех записях, кроме проверяемой."""
    out = []
    for held in names:
        tr = [n for n in names if n != held]
        Xtr = np.concatenate([X[n] for n in tr])
        ytr = np.concatenate([Y[n] for n in tr])
        m = TrackReadout(cfg, n_pn=Xtr.shape[1])
        m.fit_normalizer(Xtr)
        m.learn(Xtr, ytr, epochs=args.epochs, lr=args.lr, l2=args.l2,
                device=args.device)
        out.append((held, m.score(X[held]), m))
    return out


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--data", default=str(B.CACHE / "track_set.npz"))
    ap.add_argument("--out", default=str(B.ARTIFACTS / "track_readout.npz"))
    ap.add_argument("--save-folds", default=str(B.ARTIFACTS / "track_folds"))
    ap.add_argument("--n-kc", type=int, default=0,
                    help="ёмкость итоговой модели; 0 — лучшая из развёртки")
    ap.add_argument("--sweep-kc", default="500,1000,2000,4000")
    ap.add_argument("--epochs", type=int, default=200)
    ap.add_argument("--lr", type=float, default=2.0)
    ap.add_argument("--device", default="auto",
                    help="устройство обучения ('auto', 'cuda', 'cpu')")
    ap.add_argument("--baseline", action="store_true",
                    help="сравнить с GPU/CPU градиентным бустингом (LightGBM)")
    ap.add_argument("--drop", default="",
                    help="признаки, выброшенные из обучения — это РЕЖИМ ЗАМЕРА, "
                         "модель при этом не сохраняется: конвейер считает "
                         "описание целиком и укороченную модель не примет. "
                         "Смысл — проверить вклад признака. Так, оба разброса "
                         "на полигоне выглядят "
                         "сильными (-s_std даёт AUC 0.869 за 160 м), а на "
                         "настоящем объекте ПЕРЕВЁРНУТЫ — 0.365 м у предмета "
                         "против 0.013 м у ложных треков. Вставленный предмет "
                         "стоит на дальности, посчитанной из той же оценки "
                         "собственного движения, что двигает мировую "
                         "координату, поэтому разброса у него нет по "
                         "построению. EXPERIMENTS п. 15.1")
    ap.add_argument("--min-hits", type=int, default=2,
                    help="как в descending.py: трек с меньшим числом "
                         "подтверждений всё равно не рассматривается")
    args = ap.parse_args()

    d = np.load(args.data, allow_pickle=True)
    names = [str(n) for n in d["names"]]
    feats = [str(f) for f in d["features"]]
    assert feats == list(TRACK_FEATURES), "выборка собрана другим дескриптором"
    # Модель, обученная без части признаков, не подойдёт конвейеру: он
    # считает описание целиком. Такая модель — инструмент замера, и в
    # `descending.py` она включается только явным `--track-score model`.
    i_hits = feats.index("hits")
    i_ev = feats.index("evidence")

    X, Y, D = {}, {}, {}
    for n in names:
        Xb, yb, db = d[f"X_{n}"], d[f"y_{n}"].astype(np.int8), d[f"d_{n}"]
        keep = Xb[:, i_hits] >= args.min_hits
        X[n], Y[n], D[n] = Xb[keep].astype(np.float32), yb[keep], db[keep]

    drop = [s.strip() for s in args.drop.split(",") if s.strip()]
    if drop:
        bad = [s for s in drop if s not in feats]
        if bad:
            raise SystemExit(f"нечего выбрасывать: {bad}")
        cols = [k for k, f in enumerate(feats) if f not in drop]
        feats = [feats[k] for k in cols]
        i_hits, i_ev = feats.index("hits"), feats.index("evidence")
        for n in names:
            X[n] = X[n][:, cols]
        print("выброшено из обучения:", ", ".join(drop))

    tot = sum(v.shape[0] for v in X.values())
    pos = sum(int(v.sum()) for v in Y.values())
    print(f"выборка: {tot} наблюдений трека, предметных {pos} "
          f"({pos / max(tot, 1):.1%}), признаков {len(feats)}, "
          f"подтверждений не меньше {args.min_hits}")
    for n in names:
        print(f"  {n:<42}{X[n].shape[0]:7d}  предметных {int(Y[n].sum()):6d} "
              f"({Y[n].mean():5.1%})")

    Xall = np.concatenate([X[n] for n in names])
    yall = np.concatenate([Y[n] for n in names])
    dall = np.concatenate([D[n] for n in names])

    # ------------------------------------------- что даёт каждый признак сам
    # Ловля артефактов: если наверх вылезает признак, который не может нести
    # физику (возраст трека, ширина), — выборка перекошена, а не модель умна.
    print("\nодин признак — какой AUC, вся выборка:")
    rank = sorted(((auc(Xall[:, i], yall), f) for i, f in enumerate(feats)),
                  key=lambda p: -abs(p[0] - 0.5))
    for a, f in rank:
        print(f"  {f:<12}{a:6.3f}")

    # ------------------------------------------------------------- развёртка
    sweep = [int(x) for x in args.sweep_kc.split(",") if x.strip()]
    best, best_auc = args.n_kc or 2000, -1.0
    if sweep and not args.n_kc:
        print("\nразвёртка по ёмкости (leave-one-bag-out):")
        for n_kc in sweep:
            cfg = MbonConfig(n_kc=n_kc, sparsity=100.0 / n_kc, seed=20260922)
            t0 = time.time()
            res = loo(names, X, Y, cfg, args)
            s = np.concatenate([r[1] for r in res])
            yy = np.concatenate([Y[r[0]] for r in res])
            a = auc(s, yy)
            print(f"  {n_kc:6d} клеток: AUC {a:.4f}, "
                  f"ложных при 95 % пойманных {fpr_at_tpr(s, yy):.3f}, "
                  f"{time.time() - t0:5.1f} с")
            if a > best_auc:
                best, best_auc = n_kc, a

    cfg = MbonConfig(n_kc=best, sparsity=100.0 / best, seed=20260922)
    print(f"\nитоговая ёмкость: {best} клеток")
    res = loo(names, X, Y, cfg, args)
    s_loo = np.concatenate([r[1] for r in res])
    y_loo = np.concatenate([Y[r[0]] for r in res])
    d_loo = np.concatenate([D[r[0]] for r in res])
    ev_loo = np.concatenate([X[r[0]][:, i_ev] for r in res])

    print(f"\nleave-one-bag-out: AUC {auc(s_loo, y_loo):.4f}, "
          f"ложных при 95 % пойманных {fpr_at_tpr(s_loo, y_loo):.3f}")
    for held, s, _ in res:
        print(f"  {held:<42}AUC {auc(s, Y[held]):.4f}")

    table_cols = {"улика: AUC": ev_loo, "модель: AUC": s_loo}
    if args.baseline:
        import lightgbm as lgb
        s_lgb = []
        for held in names:
            tr = [n for n in names if n != held]
            Xtr = np.concatenate([X[n] for n in tr])
            ytr = np.concatenate([Y[n] for n in tr])
            dev_type = "gpu" if (args.device == "cuda" or (args.device == "auto" and B.Path("/dev/nvidia0").exists())) else "cpu"
            try:
                g = lgb.LGBMClassifier(n_estimators=300, learning_rate=0.05, num_leaves=31,
                                       device=dev_type, verbose=-1)
                g.fit(Xtr, ytr)
            except Exception:
                g = lgb.LGBMClassifier(n_estimators=300, learning_rate=0.05, num_leaves=31,
                                       verbose=-1)
                g.fit(Xtr, ytr)
            s_lgb.append(g.predict_proba(X[held])[:, 1])
        s_lgb_arr = np.concatenate(s_lgb)
        print(f"\nбустинг leave-one-bag-out: AUC {auc(s_lgb_arr, y_loo):.4f}, "
              f"ложных при 95 % пойманных {fpr_at_tpr(s_lgb_arr, y_loo):.3f}")
        table_cols["бустинг: AUC"] = s_lgb_arr

    # ------------------------------- главное: обгоняет ли модель саму улику
    print("\nмодель против улики, по полосам дальности:")
    band_table(d_loo, y_loo, table_cols)

    # --------------------------------------------------------------- хранение
    if drop:
        print()
        print("режим замера (--drop): модель не сохраняется")
        return
    folds = Path(args.save_folds)
    folds.mkdir(parents=True, exist_ok=True)
    for held, _, m in res:
        m.save(folds / f"track_{held}.npz")
    print(f"\nмодели по складкам: {folds} ({len(res)} шт.)")

    final = TrackReadout(cfg, n_pn=Xall.shape[1])
    final.fit_normalizer(Xall)
    final.learn(Xall, yall, epochs=args.epochs, lr=args.lr, l2=args.l2,
                device=args.device)
    B.ARTIFACTS.mkdir(parents=True, exist_ok=True)
    final.save(args.out)
    print("сохранено:", args.out)


if __name__ == "__main__":
    main()