"""Диагностические картинки одного кадра: дальностный образ, срезы, вид сверху.

    python tools/viz_frame.py --bag data/for_hackathon/doubleT_obstacle --frame 40
"""
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.retina import ScanLayout  # noqa: E402


def load(bag_path: str, frame: int, n_calib: int = 12):
    bag = Bag(bag_path)
    clouds = [pc for _, pc in bag.frames(start=2, stop=2 + n_calib)]
    layout = ScanLayout.calibrate(clouds)
    _, pc = next(bag.frames(start=frame, stop=frame + 1))
    return bag, layout, layout.project(pc)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--bag", required=True)
    ap.add_argument("--frame", type=int, default=40)
    ap.add_argument("--fov", type=float, default=30.0, help="полусектор по азимуту, °")
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    bag, layout, img = load(args.bag, args.frame)
    cols = layout.column_slice(args.fov)
    lay = layout.sub(cols)
    sub = img.crop(cols)

    xyz = sub.xyz(lay)
    x, y, z = xyz[..., 0], xyz[..., 1], xyz[..., 2]
    d = -y                                    # вперёд
    v = sub.valid

    # уровень пола по ближней зоне под сенсором
    near = v & (d > 4) & (d < 25) & (np.abs(x) < 2.0)
    z_rail = float(np.percentile(z[near], 2)) if near.sum() > 200 else np.nan
    h = z - z_rail

    fig = plt.figure(figsize=(19, 12), dpi=110)
    fig.suptitle(f"{bag.path.name}  кадр {args.frame}   "
                 f"пол z={z_rail:.2f} м   валидных {v.mean():.0%}", fontsize=13)

    # 1. дальностный образ
    ax = fig.add_subplot(3, 2, 1)
    disp = np.where(v, sub.r_near, np.nan)
    im = ax.imshow(disp, aspect="auto", cmap="turbo", vmin=0, vmax=120,
                   extent=[lay.az_grid_deg[0], lay.az_grid_deg[-1],
                           lay.el_deg[-1], lay.el_deg[0]])
    ax.set_title("дальностный образ (выпрямленный), м")
    ax.set_xlabel("азимут, °"); ax.set_ylabel("элевация, °")
    fig.colorbar(im, ax=ax, fraction=0.03)

    # 2. интенсивность
    ax = fig.add_subplot(3, 2, 2)
    im = ax.imshow(np.where(v, sub.inten, np.nan), aspect="auto", cmap="magma",
                   vmin=0, vmax=60,
                   extent=[lay.az_grid_deg[0], lay.az_grid_deg[-1],
                           lay.el_deg[-1], lay.el_deg[0]])
    ax.set_title("интенсивность")
    ax.set_xlabel("азимут, °"); ax.set_ylabel("элевация, °")
    fig.colorbar(im, ax=ax, fraction=0.03)

    # 3. вид сверху
    ax = fig.add_subplot(3, 2, 3)
    m = v & (d > 0) & (d < 200)
    ax.scatter(d[m], x[m], s=0.05, c=h[m], cmap="viridis", vmin=-1, vmax=4, linewidths=0)
    for s in (-1.5, 1.5):
        ax.axhline(s, color="r", lw=0.8, ls="--")
    ax.set_xlim(0, 200); ax.set_ylim(-8, 8)
    ax.set_title("вид сверху (цвет — высота над головкой рельса)")
    ax.set_xlabel("расстояние вперёд, м"); ax.set_ylabel("поперёк, м")

    # 4. вид сбоку
    ax = fig.add_subplot(3, 2, 4)
    ax.scatter(d[m], h[m], s=0.05, c=np.abs(x[m]), cmap="plasma", vmin=0, vmax=4, linewidths=0)
    ax.axhline(0, color="k", lw=0.8)
    ax.axhline(2.1, color="r", lw=0.8, ls="--")
    ax.set_xlim(0, 200); ax.set_ylim(-3, 6)
    ax.set_title("вид сбоку (цвет — |поперёк|)")
    ax.set_xlabel("расстояние вперёд, м"); ax.set_ylabel("высота над рельсом, м")

    # 5. поперечные сечения
    ax = fig.add_subplot(3, 2, 5)
    for lo, hi, c in [(5, 15, "tab:blue"), (30, 45, "tab:orange"),
                      (60, 80, "tab:green"), (100, 130, "tab:red")]:
        s = v & (d >= lo) & (d < hi)
        ax.scatter(x[s], h[s], s=0.3, alpha=0.4, color=c, label=f"{lo}–{hi} м", linewidths=0)
    ax.add_patch(plt.Rectangle((-1.5, 0), 3.0, 2.1, fill=False, ec="r", lw=1.5))
    ax.set_xlim(-6, 6); ax.set_ylim(-2, 6); ax.set_aspect("equal")
    ax.legend(fontsize=8, markerscale=8)
    ax.set_title("поперечные сечения (красное — габарит 3.0 × 2.1 м)")
    ax.set_xlabel("поперёк, м"); ax.set_ylabel("высота, м")

    # 6. профиль дальности по лучам вдоль пути
    ax = fig.add_subplot(3, 2, 6)
    bins = np.arange(0, 205, 2.0)
    for lo, hi, name in [(-1.0, 1.0, "|az|<1°"), (-5, 5, "|az|<5°"), (-20, 20, "|az|<20°")]:
        sel = (lay.az_grid_deg >= lo) & (lay.az_grid_deg <= hi)
        dd = d[:, sel][v[:, sel]]
        ax.hist(dd, bins=bins, histtype="step", label=name, log=True)
    ax.set_title("распределение дальностей по секторам")
    ax.set_xlabel("расстояние вперёд, м"); ax.set_ylabel("число лучей")
    ax.legend(fontsize=8)

    out = args.out or (B.FIGURES / f"frame_{bag.path.name}_{args.frame:04d}.png")
    B.FIGURES.mkdir(parents=True, exist_ok=True)
    fig.tight_layout(rect=(0, 0, 1, 0.97))
    fig.savefig(out)
    print("сохранено:", out)


if __name__ == "__main__":
    main()