"""Что попадает внутрь габарита в нормальном тоннеле.

Накапливает по многим кадрам занятость плоскости (поперёк, высота) в разных
поясах дальности и устойчивость плоскости пути. Нужно, чтобы выбрать границы
габарита и пороги, а не угадывать их.

    python tools/analyze_gauge.py --bag data/for_hackathon/roundT_doubleT --frames 60
"""
from __future__ import annotations

import argparse

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

import _bootstrap as B  # noqa: F401,E402
from flyguard.bag import Bag  # noqa: E402
from flyguard.geometry import RailPlane, TrackFrame, fit_rail_plane  # noqa: E402
from flyguard.retina import ScanLayout  # noqa: E402

BANDS = [(5, 20), (20, 50), (50, 90), (90, 150), (150, 250)]
U_EDGES = np.arange(-4.0, 4.01, 0.05)
H_EDGES = np.arange(-1.0, 4.01, 0.05)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--bag", required=True)
    ap.add_argument("--frames", type=int, default=60)
    ap.add_argument("--stride", type=int, default=3)
    ap.add_argument("--fov", type=float, default=25.0)
    args = ap.parse_args()

    bag = Bag(args.bag)
    layout = ScanLayout.calibrate([pc for _, pc in bag.frames(start=2, stop=14)])
    cols = layout.column_slice(args.fov)
    lay = layout.sub(cols)

    hist = [np.zeros((U_EDGES.size - 1, H_EDGES.size - 1), np.int64) for _ in BANDS]
    planes: list[RailPlane] = []
    deficits: list[np.ndarray] = []
    prev = None

    for _, pc in bag.frames(stop=args.frames * args.stride, stride=args.stride):
        img = layout.project(pc).crop(cols)
        plane = fit_rail_plane(img, lay, prev=prev)
        prev = plane
        planes.append(plane)
        tf = TrackFrame(img, lay, plane)

        for k, (lo, hi) in enumerate(BANDS):
            m = tf.valid & (tf.d >= lo) & (tf.d < hi)
            if m.any():
                hh, _, _ = np.histogram2d(tf.u[m], tf.h[m], bins=(U_EDGES, H_EDGES))
                hist[k] += hh.astype(np.int64)

        # дефицит до пола: луч должен был дойти до полотна внутри коридора
        fr = tf.floor_r
        uf = fr * lay.dirs[..., 0]
        hit_in = np.isfinite(fr) & (np.abs(uf) < 1.7) & (fr > 10) & (fr < 200)
        m = hit_in & tf.valid
        if m.any():
            deficits.append((fr[m] - tf.r[m]).astype(np.float32))

    a = np.array([p.a for p in planes]); b = np.array([p.b for p in planes])
    hgt = np.array([p.height for p in planes]); rms = np.array([p.rms for p in planes])
    print(f"=== {bag.path.name}: {len(planes)} кадров")
    print(f"  высота сенсора над путём: {hgt.mean():.3f} ± {hgt.std():.3f} м "
          f"(разброс {hgt.min():.3f}…{hgt.max():.3f})")
    print(f"  тангаж: {np.degrees(np.arctan(a)).mean():+.3f}° ± {np.degrees(np.arctan(a)).std():.3f}°")
    print(f"  крен:   {np.degrees(np.arctan(b)).mean():+.3f}° ± {np.degrees(np.arctan(b)).std():.3f}°")
    print(f"  невязка плоскости rms: {rms.mean():.3f} м, инлайеров {np.mean([p.inliers for p in planes]):.0f}")

    if deficits:
        dd = np.concatenate(deficits)
        print(f"  дефицит до пола внутри коридора, м: "
              + "  ".join(f"p{q}={np.percentile(dd, q):+.2f}" for q in (1, 50, 90, 99, 99.9))
              + f"  доля > 1 м: {np.mean(dd > 1.0):.3%}")

    fig, axes = plt.subplots(2, len(BANDS), figsize=(4 * len(BANDS), 8), dpi=110)
    fig.suptitle(f"{bag.path.name}: занятость (поперёк, высота) по поясам дальности; "
                 f"красное — габарит 3.4 × 2.2 м", fontsize=12)
    for k, (lo, hi) in enumerate(BANDS):
        hh = hist[k].T
        tot = hh.sum()
        for row, norm in enumerate(("log", "col")):
            ax = axes[row, k]
            if tot == 0:
                ax.set_title(f"{lo}–{hi} м: пусто"); continue
            show = np.log10(1 + hh) if norm == "log" else hh / np.maximum(hh.sum(0, keepdims=True), 1)
            ax.imshow(show, origin="lower", aspect="auto", cmap="inferno",
                      extent=[U_EDGES[0], U_EDGES[-1], H_EDGES[0], H_EDGES[-1]])
            ax.add_patch(plt.Rectangle((-1.7, 0.05), 3.4, 2.15, fill=False, ec="r", lw=1.2))
            ax.axhline(0, color="c", lw=0.6)
            ax.set_title(f"{lo}–{hi} м, n={tot/1e3:.0f}k" + ("" if row == 0 else " (норм. по столбцу)"),
                         fontsize=9)
            ax.set_xlabel("поперёк, м")
            if k == 0:
                ax.set_ylabel("высота над рельсом, м")

    B.FIGURES.mkdir(parents=True, exist_ok=True)
    out = B.FIGURES / f"gauge_{bag.path.name}.png"
    fig.tight_layout(rect=(0, 0, 1, 0.95))
    fig.savefig(out)
    print("  сохранено:", out)


if __name__ == "__main__":
    main()