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

Разметки в датасете нет, и до сих пор это определяло архитектуру: грибовидное
тело учится **без меток**, запоминая частоту обстановки. Но у нас есть
физически обоснованный генератор предметов (`flyguard.synth`), сверенный с
единственным реальным объектом: настоящий 0.67 × 1.35 м на 55 м даёт 47–69
лучей, синтетический человек 0.44 × 1.71 м на 60 м — 50. Значит, метки можно
изготовить, и изготовить достоверно.

Каждый кандидат помечается по **пересечению лучей**, а не «по дальности
примерно»: `inject` возвращает индексы лучей, в которые предмет действительно
записан, и кандидат считается предметом, если его ядро состоит из этих лучей.
Так структура тоннеля, случайно оказавшаяся на той же дальности, в
положительные не попадает.

Сценарии намеренно ставят предмет в РАЗНЫЕ точки тоннеля (`--d-starts`): замер
показал, что одна и та же дальность в разных местах перегона ведёт себя
совершенно по-разному — где-то предмет виден целиком, где-то за поворотом
(EXPERIMENTS п. 9.6). Обучаться на одной точке постановки значит выучить эту
точку.

`doubleT_obstacle` исключён целиком: там настоящий объект, и он остаётся
независимой проверкой того, что модель выучила предмет, а не «синтетику».

    python tools/make_training_set.py --out data/cache/training_set.npz
"""
from __future__ import annotations

import argparse

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 FEATURES, describe
from flyguard.pipeline import FlyGuard, Params
from flyguard.synth import IntensityEnv, Placement, catalogue, inject

HOLDOUT = "doubleT_obstacle"      # там реальный объект — только для проверки
MIN_OVERLAP = 0.5                 # доля лучей ядра, пришедших от предмета


def _work(task):
    """Одна задача — один бэг. Верхнего уровня: иначе не передать в процесс."""
    path, params, limit, d_starts, laterals, seed = task
    return collect_bag(path, params, limit, d_starts, laterals, seed)


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


def collect_bag(path, params: Params, limit: int, d_starts, laterals, seed: int):
    bag = Bag(path)
    s_track, ego_fg = ego_track(bag, params, limit)
    have = [x for x in s_track if x is not None]
    if len(have) < 30:
        return None
    s0 = have[0]
    layout = ego_fg.layout_full
    col0 = ego_fg.cols.start

    cat = catalogue()
    scen = [(n, lat, d0) for n in cat for lat in laterals for d0 in d_starts]
    guide = FlyGuard(params, memory=None, layout=layout)
    pipes = [FlyGuard(params, memory=None, layout=layout) for _ in scen]
    rng = np.random.default_rng(seed)

    X, y, dd, obj, lat_out = [], [], [], [], []
    for k, (_, pc) in enumerate(bag.frames(stop=limit)):
        g = guide.process(pc)
        if g is None or s_track[k] is None:
            continue
        s_now = s_track[k] - s0
        env = IntensityEnv(pc)          # один раз на кадр, общий для сценариев
        for i, (name, lat, d0) in enumerate(scen):
            d_true = d0 - s_now
            if d_true < 6.0:
                continue
            u = float(g.corridor.centre(np.array([d_true], np.float32))[0]) + lat
            pc2, lab = inject(pc, layout, g.plane, cat[name],
                              Placement(d=d_true, u=u), rng=rng, env=env)
            res = pipes[i].process(pc2)
            if res is None or not res.candidates:
                continue
            rr = lab.get("rays")
            if rr is None or lab["hit_rays"] == 0:
                truth = None
            else:
                # лучи предмета в координатах полной решётки → плоский индекс;
                # отсортированный массив, а не множество: проверка идёт
                # сотни тысяч раз, и `in` по множеству тут заметно дороже
                truth = np.sort((rr[0].astype(np.int64) << 20)
                                | rr[1].astype(np.int64))
            for c in res.candidates:
                ii, jj = c.extra.get("rays", (None, None))
                if ii is None:
                    continue
                lbl = 0
                if truth is not None and truth.size:
                    key = ((ii.astype(np.int64) << 20)
                           | (jj.astype(np.int64) + col0))
                    pos = np.searchsorted(truth, key)
                    np.clip(pos, 0, truth.size - 1, out=pos)
                    frac = float((truth[pos] == key).mean())
                    lbl = int(frac >= MIN_OVERLAP)
                v = describe(c)
                acc = float(c.extra.get("acc_support", 0.0)) if c.extra else 0.0
                X.append(np.append(v, acc).astype(np.float32))
                y.append(lbl)
                dd.append(c.d)
                obj.append(name if lbl else "")
                lat_out.append(lat)
    return (np.asarray(X, np.float32), np.asarray(y, np.int8),
            np.asarray(dd, np.float32), np.asarray(obj),
            np.asarray(lat_out, np.float32))


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--root", default=str(B.DATA / "for_hackathon"))
    ap.add_argument("--out", default=str(B.CACHE / "training_set.npz"))
    ap.add_argument("--limit", type=int, default=250)
    ap.add_argument("--d-starts", default="80,140,200")
    # Поперечные положения намеренно кроют ВЕСЬ габарит, а не только ось.
    # Иначе выборка вырождается: предметы у оси, обстановка у стен, и модель
    # выучивает «всё, что у оси — предмет» вместо признаков предмета. Проверено
    # на пробном прогоне с одной постановкой: AUC 1.000 по всем бэгам — цифра
    # красивая и бессмысленная, а в тоннеле у оси полно штатных конструкций.
    ap.add_argument("--laterals", default="0,-0.6,0.6,-1.2,1.2")
    ap.add_argument("--seed", type=int, default=20260921)
    P.add_argument(ap)
    args = ap.parse_args()

    d_starts = tuple(float(x) for x in args.d_starts.split(","))
    laterals = tuple(float(x) for x in args.laterals.split(","))
    params = Params()
    B.CACHE.mkdir(parents=True, exist_ok=True)

    bags = [p for p in find_bags(args.root) if p.name != HOLDOUT]
    tasks = [(p, params, args.limit, d_starts, laterals, args.seed) for p in bags]
    # Печатается по готовности, складывается по номеру задачи: порядок бэгов
    # в файле не должен зависеть от того, какой из них досчитался первым.
    slots: list = [None] * len(tasks)
    for i, task, got, secs in P.run(_work, tasks, args.jobs):
        name = task[0].name
        if got is None:
            print(f"  {name:42s} пропущен")
            continue
        slots[i] = (name, got)
        Xb, yb = got[0], got[1]
        print(f"  {name:42s} {Xb.shape[0]:7d} кандидатов, "
              f"предметов {int(yb.sum()):6d} ({yb.mean():5.1%}), "
              f"{secs:6.0f} с", flush=True)
    parts = {name: got for name, got in (s for s in slots if s is not None)}

    if not parts:
        raise SystemExit("ничего не собрано")
    out = {}
    for name, (Xb, yb, db, ob, lb) in parts.items():
        out[f"X_{name}"] = Xb
        out[f"y_{name}"] = yb
        out[f"d_{name}"] = db
        out[f"obj_{name}"] = ob
        out[f"lat_{name}"] = lb
    np.savez_compressed(args.out, names=np.array(list(parts)),
                        features=np.array(FEATURES + ("acc_support",)), **out)
    tot = sum(v[0].shape[0] for v in parts.values())
    pos = sum(int(v[1].sum()) for v in parts.values())
    print(f"\nвсего {tot} кандидатов, предметов {pos} ({pos / tot:.1%})")
    print("сохранено:", args.out)


if __name__ == "__main__":
    main()