"""Проверка стабилизации и осевой линии пути на всех бэгах.

Печатает устойчивость плоскости рельсов, оценку радиуса кривой и реальную
дальность прямой видимости; рисует профили осевой линии.

    python tools/analyze_corridor.py --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, find_bags  # noqa: E402
from flyguard.geometry import TrackFrame, fit_corridor, fit_rail_plane  # noqa: E402
from flyguard.retina import ScanLayout  # noqa: E402


def run(bag_path, n_frames: int, stride: int, fov: float, ax):
    bag = Bag(bag_path)
    layout = ScanLayout.calibrate([pc for _, pc in bag.frames(start=2, stop=14)])
    cols = layout.column_slice(fov)
    lay = layout.sub(cols)

    plane = corridor = None
    heights, pitches, rolls, rmss, radii, reach = [], [], [], [], [], []
    curves = []
    for k, (_, pc) in enumerate(bag.frames(stop=n_frames * stride, stride=stride)):
        img = layout.project(pc).crop(cols)
        plane = fit_rail_plane(img, lay, prev=plane)
        tf = TrackFrame(img, lay, plane)
        corridor = fit_corridor(tf, prev=corridor)

        heights.append(plane.height); pitches.append(plane.pitch_deg)
        rolls.append(plane.roll_deg); rmss.append(plane.rms)
        radii.append(corridor.radius)
        # дальность прямой видимости: дальше какой дистанции возвраты иссякают
        d = tf.d[tf.valid]
        reach.append(np.percentile(d, 99.9) if d.size else 0.0)
        if k % max(1, n_frames // 6) == 0:
            dd = np.linspace(0, 200, 100)
            curves.append(corridor.centre(dd))

    name = bag.path.name
    print(f"=== {name}  ({len(heights)} кадров)")
    print(f"  высота сенсора: {np.mean(heights):.3f} ± {np.std(heights):.3f} м")
    print(f"  тангаж {np.mean(pitches):+.3f}° ± {np.std(pitches):.3f}°   "
          f"крен {np.mean(rolls):+.3f}° ± {np.std(rolls):.3f}°   "
          f"rms {np.mean(rmss):.3f} м")
    r = np.array(radii)
    fin = np.isfinite(r)
    print(f"  радиус кривой: медиана {np.median(r[fin]) if fin.any() else float('inf'):.0f} м, "
          f"прямых кадров {np.mean(~fin):.0%}")
    print(f"  дальность прямой видимости p99.9: медиана {np.median(reach):.0f} м, "
          f"макс {np.max(reach):.0f} м")

    dd = np.linspace(0, 200, 100)
    for c in curves:
        ax.plot(dd, c, lw=1, alpha=0.7)
    ax.axhline(1.7, color="r", ls="--", lw=0.8)
    ax.axhline(-1.7, color="r", ls="--", lw=0.8)
    ax.set_title(f"{name}\nR≈{np.median(r[fin]) if fin.any() else float('inf'):.0f} м, "
                 f"видимость {np.median(reach):.0f} м", fontsize=9)
    ax.set_xlabel("вперёд, м"); ax.set_ylabel("ось пути, м")
    ax.set_ylim(-12, 12)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--root", default=str(B.DATA / "for_hackathon"))
    ap.add_argument("--frames", type=int, default=60)
    ap.add_argument("--stride", type=int, default=4)
    ap.add_argument("--fov", type=float, default=35.0)
    args = ap.parse_args()

    bags = find_bags(args.root)
    fig, axes = plt.subplots(1, len(bags), figsize=(3.6 * len(bags), 3.6), dpi=110, squeeze=False)
    for ax, b in zip(axes[0], bags):
        run(b, args.frames, args.stride, args.fov, ax)
    B.FIGURES.mkdir(parents=True, exist_ok=True)
    out = B.FIGURES / "corridor.png"
    fig.tight_layout(); fig.savefig(out)
    print("сохранено:", out)


if __name__ == "__main__":
    main()