"""Размеченный полигон: сценарии сближения с синтетическим препятствием.

Для каждого бэга и каждого типа предмета строится сценарий: предмет ставится
в фиксированную точку тоннеля далеко впереди, поезд к нему подъезжает, и на
каждом кадре известна истинная дистанция. Отсюда получаются именно те цифры,
которые просит ТЗ: с какой дальности предмет уверенно виден, сколько ложных
тревог и как это зависит от размера.

Первый проход считает собственное движение по чистым данным (это и есть
разметка по дистанции), второй — гоняет конвейер по кадрам со вставленным
предметом. Все сценарии одного бэга обрабатываются в одном проходе по файлу:
чтение данных дороже самой обработки.

    python tools/make_benchmark.py --out artifacts/benchmark.npz --memory artifacts/mushroom_body.npz
"""
from __future__ import annotations

import argparse
import json

import numpy as np

import _bootstrap as B  # noqa: F401
import _parallel as P
from flyguard.bag import Bag, find_bags
from flyguard.mushroom_body import MushroomBody
from flyguard.pipeline import FlyGuard, Params
from flyguard.synth import IntensityEnv, Placement, catalogue, inject
from flyguard.track_readout import describe_track

HOLDOUT = "doubleT_obstacle"        # там уже есть настоящий объект


def _work(task):
    """Одна задача — один бэг.

    Обученное грузится путями и уже внутри процесса: передавать модели через
    межпроцессную границу незачем, а свою складку каждый воркер берёт сам.
    """
    (path, params, limit, d_start, laterals, seed, mem_path, rd_path,
     want_tracks, trk_path) = task
    memory = MushroomBody.load(mem_path) if mem_path else None
    readout = None
    if rd_path:
        from flyguard.mbon_readout import MbonReadout
        readout = MbonReadout.load(rd_path)
    track_rd = None
    if trk_path:
        from flyguard.track_readout import TrackReadout
        track_rd = TrackReadout.load(trk_path)
    return run_bag(path, params, memory, limit, d_start, laterals, seed,
                   readout=readout, want_tracks=want_tracks,
                   track_readout=track_rd)


def ego_track(bag: Bag, params: Params, limit: int | None):
    """Первый проход: пройденный путь на каждом кадре (разметка по дистанции)."""
    fg = FlyGuard(params, memory=None)
    s, stamps = [], []
    total = 0.0
    for _, pc in bag.frames(stop=limit):
        res = fg.process(pc)
        if res is None:
            s.append(None); stamps.append(pc.stamp); continue
        total += res.ego.ds if res.ego else 0.0
        s.append(total); stamps.append(pc.stamp)
    return s, stamps


def run_bag(bag_path, params: Params, memory, limit: int, d_start: float,
            laterals: tuple[float, ...], seed: int, readout=None,
            want_tracks: bool = False, track_readout=None):
    bag = Bag(bag_path)
    s_track, _ = ego_track(bag, params, limit)
    have = [x for x in s_track if x is not None]
    if len(have) < 20:
        return []
    travel = have[-1] - have[0]

    cat = catalogue()
    scen = [(name, lat) for name in cat for lat in laterals]
    pipes = [FlyGuard(params, memory=memory, readout=readout,
                      track_readout=track_readout) for _ in scen]
    rng = np.random.default_rng(seed)
    records = [[] for _ in scen]
    # Выборка для считывания по трекам: описание каждого живого трека на
    # каждом кадре и метка «это вставленный предмет». Собирается здесь, а не
    # отдельным инструментом, ровно затем, чтобы распределение совпадало с
    # рабочим: та же память тоннеля, то же покадровое считывание, те же
    # пороги. Обученное на другой обстановке считывание нечего и мерить.
    trk_rows: list = []

    # решётка и поза нужны для вставки — берутся из отдельного «чистого» конвейера
    guide = FlyGuard(params, memory=None)

    for k, (_, pc) in enumerate(bag.frames(stop=limit)):
        gres = guide.process(pc)
        if gres is None or s_track[k] is None:
            continue
        s_now = s_track[k] - have[0]
        env = IntensityEnv(pc)          # один раз на кадр, общий для сценариев
        for i, (name, lat) in enumerate(scen):
            d_true = d_start - s_now
            if d_true < 6.0:
                continue
            # Предмет лежит НА ПУТИ, а путь в кривой уходит вбок: на 150 м при
            # радиусе 1300 м это 8.6 м. Если ставить его в поперечных координатах
            # сенсора, он окажется в стене, а не в габарите.
            u_obj = float(gres.corridor.centre(np.array([d_true], np.float32))[0]) + lat
            pc2, lab = inject(pc, guide.layout_full, gres.plane, cat[name],
                              Placement(d=d_true, u=u_obj), rng=rng, env=env)
            res = pipes[i].process(pc2)
            if res is None:
                continue
            tol = max(3.0, 0.12 * d_true)
            hit = any(abs(o.distance - d_true) < tol for o in res.decision.objects)
            fp = sum(1 for o in res.decision.objects if abs(o.distance - d_true) >= tol)
            # Воронка потерь. Лучи в предмет попали — а дальше он может
            # пропасть на любой из трёх ступеней, и лечатся они по-разному:
            # нет кандидата — вопрос к кластеризации и разделению фигуры и
            # фона, нет трека — к сопоставлению по кадрам, нет решения —
            # к порогу. Без этого разбиения улучшать нечего, кроме удачи.
            cand = any(abs(c.d - d_true) < tol for c in res.candidates)
            cx = pipes[i].cx
            trk = any(abs(tr.distance(cx.s_world) - d_true) < tol
                      for tr in cx.tracks)
            records[i].append((d_true, int(hit), fp, lab["hit_rays"],
                               int(cand), int(trk)))
            if want_tracks:
                for tr in cx.tracks:
                    dd = tr.distance(cx.s_world)
                    if not (0.0 < dd <= params.d_max):
                        continue
                    trk_rows.append((describe_track(tr, cx.s_world),
                                     int(abs(dd - d_true) < tol), dd))

    out = []
    for (name, lat), rec in zip(scen, records):
        if not rec:
            continue
        a = np.array(rec, np.float32)
        out.append(dict(bag=bag.path.name, obj=name, lateral=lat,
                        d=a[:, 0].tolist(), hit=a[:, 1].tolist(),
                        fp=a[:, 2].tolist(), rays=a[:, 3].tolist(),
                        cand=a[:, 4].tolist(), trk=a[:, 5].tolist(),
                        travel=float(travel)))
    tracks = None
    if want_tracks and trk_rows:
        tracks = (np.stack([r[0] for r in trk_rows]).astype(np.float32),
                  np.array([r[1] for r in trk_rows], np.int8),
                  np.array([r[2] for r in trk_rows], np.float32))
    return out, tracks


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--root", default=str(B.DATA / "for_hackathon"))
    ap.add_argument("--memory")
    ap.add_argument("--mbon", default="",
                    help="обученное считывание MBON одной моделью; она видела "
                         "эти бэги — годится только для отладки")
    ap.add_argument("--mbon-dir", default="",
                    help="каталог с моделями по складкам (mbon_<бэг>.npz): для "
                         "каждого бэга берётся модель, его НЕ видевшая. Иначе "
                         "дальность завышена: считывание обучалось ровно на "
                         "таких же вставках в этот же тоннель")
    ap.add_argument("--out", default=str(B.ARTIFACTS / "benchmark.json"))
    ap.add_argument("--tracks-out", default="",
                    help="куда сложить выборку по трекам (npz); пусто — не собирать")
    ap.add_argument("--track", default="",
                    help="считывание по трекам одной моделью; она видела эти "
                         "записи — годится только для отладки")
    ap.add_argument("--track-dir", default="",
                    help="каталог с моделями по складкам (track_<бэг>.npz)")
    ap.add_argument("--track-blend", type=float, default=None,
                    help="доля считывания по треку в решении; 0 — только улика")
    ap.add_argument("--track-score", default=None,
                    choices=("w_mean", "w_ema", "p_mean", "model"),
                    help="чем мерить трек: средний вес наблюдения, средний "
                         "отсчёт считывания или обученная модель")
    ap.add_argument("--limit", type=int, default=250)
    ap.add_argument("--d-start", type=float, default=200.0)
    ap.add_argument("--laterals", default="0.0,0.9")
    ap.add_argument("--seed", type=int, default=12345)
    ap.add_argument("--mbon-blend", type=float, default=None,
                    help="1 — только модель, 0 — только ручная формула")
    ap.add_argument("--warn", type=float, default=None,
                    help="порог улики для тревоги (по умолчанию 0.5)")
    ap.add_argument("--clear", type=float, default=None)
    ap.add_argument("--min-hits", type=int, default=None)
    ap.add_argument("--warn-far", type=float, default=None,
                    help="порог тревоги на дальнем краю; ниже обычного — послабление далёким трекам")
    ap.add_argument("--warn-far-from", type=float, default=None,
                    help="с какой дальности порог начинает падать, м")
    ap.add_argument("--novelty-floor", type=float, default=None,
                    help="ниже этой новизны трек не считается")
    ap.add_argument("--near-long", type=float, default=None,
                    help="разброс дальности, выше которого компоненту режут и вблизи; 0 — не резать вблизи")
    ap.add_argument("--min-rays", type=int, default=None,
                    help="сколько лучей минимум образуют кандидата")
    ap.add_argument("--min-rays-far", type=int, default=None,
                    help="порог по лучам за `--min-rays-far-from`; 0 — не различать")
    ap.add_argument("--min-rays-far-from", type=float, default=None)
    ap.add_argument("--leak-far", type=float, default=None,
                    help="утечка улики на дальнем краю; ниже обычной — далёкий трек прощает промахи")
    ap.add_argument("--leak-far-from", type=float, default=None,
                    help="с какой дальности утечка начинает падать, м")
    ap.add_argument("--nov-fade-from", type=float, default=None,
                    help="с какой дальности гасить вклад знакомости; 0 — не гасить")
    ap.add_argument("--nov-fade-to", type=float, default=None,
                    help="к какой дальности вклад знакомости обнуляется")
    ap.add_argument("--mbon-prior-from", type=float, default=None,
                    help="с какой дальности поправлять оценку модели на распространённость предметов; 0 — не поправлять")
    P.add_argument(ap)
    args = ap.parse_args()

    B.ARTIFACTS.mkdir(parents=True, exist_ok=True)
    fold_paths: dict[str, str] = {}
    if args.mbon_dir:
        from pathlib import Path as _P
        for f in _P(args.mbon_dir).glob("mbon_*.npz"):
            fold_paths[f.stem[len("mbon_"):]] = str(f)
        print(f"считывание MBON по складкам: {args.mbon_dir} "
              f"({len(fold_paths)} моделей)")
    elif args.mbon:
        print(f"считывание MBON: {args.mbon}")
    track_paths: dict[str, str] = {}
    if args.track_dir:
        from pathlib import Path as _P2
        for f in _P2(args.track_dir).glob("track_*.npz"):
            track_paths[f.stem[len("track_"):]] = str(f)
        print(f"считывание по трекам по складкам: {args.track_dir} "
              f"({len(track_paths)} моделей)")
    elif args.track:
        print(f"считывание по трекам: {args.track}")
    over = {}
    if args.mbon_blend is not None:
        over["mbon_blend"] = args.mbon_blend
    if args.warn is not None:
        over["warn_evidence"] = args.warn
        over["clear_evidence"] = args.warn * 0.6 if args.clear is None else args.clear
    elif args.clear is not None:
        over["clear_evidence"] = args.clear
    if args.min_hits is not None:
        over["min_hits"] = args.min_hits
    if args.warn_far is not None:
        over["warn_far"] = args.warn_far
    if args.warn_far_from is not None:
        over["warn_far_from"] = args.warn_far_from
    if args.novelty_floor is not None:
        over["novelty_floor"] = args.novelty_floor
    if args.near_long is not None:
        over["near_long"] = args.near_long
    if args.min_rays is not None:
        over["min_rays"] = args.min_rays
    if args.min_rays_far is not None:
        over["min_rays_far"] = args.min_rays_far
    if args.min_rays_far_from is not None:
        over["min_rays_far_from"] = args.min_rays_far_from
    if args.leak_far is not None:
        over["leak_far"] = args.leak_far
    if args.leak_far_from is not None:
        over["leak_far_from"] = args.leak_far_from
    if args.nov_fade_from is not None:
        over["nov_fade_from"] = args.nov_fade_from
    if args.nov_fade_to is not None:
        over["nov_fade_to"] = args.nov_fade_to
    if args.mbon_prior_from is not None:
        over["mbon_prior_from"] = args.mbon_prior_from
    if args.track_blend is not None:
        over["track_blend"] = args.track_blend
    if args.track_score is not None:
        over["track_score"] = args.track_score
    params = Params(**over)
    laterals = tuple(float(x) for x in args.laterals.split(","))

    tasks = []
    for p in find_bags(args.root):
        if p.name == HOLDOUT:
            continue
        rd = args.mbon
        if fold_paths:
            rd = fold_paths.get(p.name, "")
            if not rd:
                print(f"  внимание: для {p.name} нет своей складки — пропуск")
                continue
        trk = track_paths.get(p.name, args.track) if track_paths else args.track
        tasks.append((p, params, args.limit, args.d_start, laterals, args.seed,
                      args.memory or "", rd, bool(args.tracks_out), trk))

    # Печатается по готовности, собирается по номеру задачи: порядок сценариев
    # в файле не должен зависеть от того, какой бэг досчитался первым.
    slots: list = [None] * len(tasks)
    trk_slots: list = [None] * len(tasks)
    for i, task, got, secs in P.run(_work, tasks, args.jobs):
        rec, tracks = got
        slots[i] = rec
        trk_slots[i] = (task[0].name, tracks)
        n = sum(len(r["d"]) for r in rec)
        extra = f", треков {tracks[0].shape[0]:7d}" if tracks is not None else ""
        print(f"  {task[0].name:42s} сценариев {len(rec):3d}, наблюдений {n:6d}"
              f"{extra}, {secs:6.1f} с", flush=True)
    all_rec = [r for rec in slots if rec for r in rec]

    if args.tracks_out:
        from flyguard.track_readout import TRACK_FEATURES
        out_arr, names = {}, []
        for got in trk_slots:
            if got is None or got[1] is None:
                continue
            name, (Xb, yb, db) = got
            names.append(name)
            out_arr[f"X_{name}"] = Xb
            out_arr[f"y_{name}"] = yb
            out_arr[f"d_{name}"] = db
        np.savez_compressed(args.tracks_out, names=np.array(names),
                            features=np.array(TRACK_FEATURES), **out_arr)
        tot = sum(v.shape[0] for k, v in out_arr.items() if k.startswith("X_"))
        pos = sum(int(v.sum()) for k, v in out_arr.items() if k.startswith("y_"))
        print(f"выборка по трекам: {tot} наблюдений, предметных {pos} "
              f"({pos / max(tot, 1):.1%}) -> {args.tracks_out}")

    with open(args.out, "w", encoding="utf-8") as f:
        json.dump(all_rec, f, ensure_ascii=False)
    print(f"сохранено: {args.out} ({len(all_rec)} сценариев)")


if __name__ == "__main__":
    main()