"""Инвентаризация бэгов и калибровка омматидиальной решётки.

Для каждого бэга: состав топиков, раскладка скана, углы, проверка того, что
калиброванная решётка действительно описывает направления лучей.

    python tools/inspect_bags.py                 # все бэги под data/
    python tools/inspect_bags.py --bag data/for_hackathon/doubleT_obstacle
"""
from __future__ import annotations

import argparse

import numpy as np

import _bootstrap as B  # noqa: F401  (добавляет пакет в sys.path)
from flyguard.bag import Bag, find_bags
from flyguard.retina import ScanLayout

N_CALIB_FRAMES = 12


def check_layout(layout: ScanLayout, bag: Bag, n: int = 3) -> dict:
    """Насколько точно решётка предсказывает направления реальных точек.

    Сравнение делается после выпрямления: направления точек проходят ту же
    выборку `gather`, что и дальности, и сверяются с таблицей `layout.dirs`.
    """
    H, W, E = layout.n_rings, layout.n_az, layout.n_echo
    ang_err, rng_err = [], []
    for _, pc in bag.frames(start=5, stop=5 + n):
        img = layout.project(pc)
        raw = np.stack([pc.points[k].reshape(W, E, H).transpose(2, 0, 1)
                        for k in ("x", "y", "z")], axis=-1)      # (H, W, E, 3)
        r_raw = np.linalg.norm(raw, axis=-1)
        near = np.argmin(np.where(r_raw > 0, r_raw, np.inf), axis=-1)
        sel = np.take_along_axis(raw, near[..., None, None], -2)[..., 0, :]
        r_sel = np.take_along_axis(r_raw, near[..., None], -1)[..., 0]

        g = layout.gather
        sel = np.take_along_axis(sel, g[..., None], 1)
        r_sel = np.take_along_axis(r_sel, g, 1)

        m = img.valid & (r_sel > 1.0)
        if not m.any():
            continue
        unit = sel[m] / r_sel[m][:, None]
        cos = np.clip(np.einsum("ij,ij->i", unit, layout.dirs[m]), -1, 1)
        ang_err.append(np.degrees(np.arccos(cos)))
        rng_err.append(np.abs(r_sel[m] - img.r_near[m]))
    if not ang_err:
        return {}
    ang = np.concatenate(ang_err)
    return {
        "ang_p50": float(np.percentile(ang, 50)),
        "ang_p99": float(np.percentile(ang, 99)),
        "ang_max": float(ang.max()),
        "rng_max": float(np.concatenate(rng_err).max()),
        "n": int(ang.size),
    }


def summarise(bag_path) -> None:
    bag = Bag(bag_path)
    print("=" * 78)
    print(bag.describe())

    clouds = [pc for _, pc in bag.frames(start=2, stop=2 + N_CALIB_FRAMES)]
    layout = ScanLayout.calibrate(clouds)
    print(f"  раскладка: {layout}")

    pc = clouds[0]
    img = layout.project(pc)
    r = img.r_near[img.valid]
    print(f"  точек в кадре: {pc.n_points}  валидных лучей: {img.valid.mean():6.1%}")
    print(f"  дальность: p50={np.percentile(r, 50):6.1f}  p99={np.percentile(r, 99):6.1f}  "
          f"max={r.max():6.1f}")

    sep = img.r_far - img.r_near
    two = img.valid & (sep > 0.05)
    print(f"  лучей с двумя различимыми эхами: {two.mean():6.2%}"
          + (f", разнос p50={np.median(sep[two]):.2f} м p95={np.percentile(sep[two], 95):.2f} м"
             if two.any() else ""))

    sh = layout.col_shift
    print(f"  скос каналов: {sh.min()}…{sh.max()} столбцов "
          f"({sh.min() * layout.az_step_deg:+.2f}°…{sh.max() * layout.az_step_deg:+.2f}°), "
          f"остаток выпрямления |max| {np.abs(layout.az_resid_deg).max():.4f}°")

    chk = check_layout(layout, bag)
    if chk:
        print(f"  ПРОВЕРКА решётки по {chk['n']} лучам: угловая ошибка "
              f"p50={chk['ang_p50']:.4f}° p99={chk['ang_p99']:.4f}° max={chk['ang_max']:.4f}°, "
              f"ошибка дальности max={chk['rng_max']:.2e} м")

    out = B.CACHE / "layouts"
    out.mkdir(parents=True, exist_ok=True)
    layout.save(out / f"{bag.path.name}.npz")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--bag", action="append", help="путь к бэгу (можно несколько)")
    ap.add_argument("--root", default=str(B.DATA), help="корень для поиска бэгов")
    args = ap.parse_args()

    bags = [__import__("pathlib").Path(b) for b in args.bag] if args.bag else find_bags(args.root)
    if not bags:
        raise SystemExit(f"бэги не найдены под {args.root}")
    for b in bags:
        summarise(b)


if __name__ == "__main__":
    main()