"""Крупный план участка дальностного образа вокруг заданной точки пространства.

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

    python tools/viz_object.py --bag data/for_hackathon/doubleT_obstacle \
        --frame 40 --dist 55 --lat -0.8
"""
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 main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--bag", required=True)
    ap.add_argument("--frame", type=int, default=40)
    ap.add_argument("--dist", type=float, required=True, help="расстояние вперёд, м")
    ap.add_argument("--lat", type=float, default=0.0, help="смещение поперёк, м")
    ap.add_argument("--pad-az", type=float, default=4.0, help="окно по азимуту, °")
    ap.add_argument("--pad-el", type=float, default=4.0, help="окно по элевации, °")
    args = ap.parse_args()

    bag = Bag(args.bag)
    layout = ScanLayout.calibrate([pc for _, pc in bag.frames(start=2, stop=14)])
    _, pc = next(bag.frames(start=args.frame, stop=args.frame + 1))
    img = layout.project(pc)

    az_c = np.degrees(np.arctan2(args.lat, args.dist))
    el_c = 0.0
    ja = np.flatnonzero(np.abs(layout.az_grid_deg - az_c) <= args.pad_az)
    ie = np.flatnonzero(np.abs(layout.el_deg - el_c) <= args.pad_el)
    js, je = slice(ja[0], ja[-1] + 1), slice(ie[0], ie[-1] + 1)

    r = np.where(img.valid, img.r_near, np.nan)[je, js]
    it = np.where(img.valid, img.inten, np.nan)[je, js]
    gap = np.where(img.valid, img.r_far - img.r_near, np.nan)[je, js]
    az = layout.az_grid_deg[js]
    el = layout.el_deg[je]
    ext = [az[0], az[-1], el[-1], el[0]]

    print(f"окно: азимут {az.min():.2f}…{az.max():.2f}°  ({az.size} стлб), "
          f"элевация {el.min():.2f}…{el.max():.2f}°  ({el.size} колец)")
    band = r[np.isfinite(r)]
    print(f"дальности в окне: p5={np.percentile(band, 5):.1f} p50={np.percentile(band, 50):.1f} "
          f"p95={np.percentile(band, 95):.1f} max={band.max():.1f}")
    near = np.isfinite(r) & (np.abs(r - args.dist) < 5)
    print(f"лучей на дистанции {args.dist}±5 м: {near.sum()}, "
          f"доля окна {near.mean():.1%}")
    print(f"пустых лучей в окне: {np.mean(~np.isfinite(r)):.1%}")

    fig, axes = plt.subplots(2, 2, figsize=(15, 9), dpi=110)
    fig.suptitle(f"{bag.path.name} кадр {args.frame}: окрестность точки "
                 f"({args.dist:.0f} м вперёд, {args.lat:+.1f} м поперёк)")

    ax = axes[0, 0]
    im = ax.imshow(r, extent=ext, aspect="auto", cmap="turbo", interpolation="nearest")
    ax.set_title("дальность, м"); fig.colorbar(im, ax=ax, fraction=0.04)

    ax = axes[0, 1]
    lo, hi = np.nanpercentile(r, [2, 98])
    im = ax.imshow(r, extent=ext, aspect="auto", cmap="turbo", vmin=lo, vmax=hi,
                   interpolation="nearest")
    ax.set_title(f"дальность, растянутая шкала {lo:.0f}…{hi:.0f} м")
    fig.colorbar(im, ax=ax, fraction=0.04)

    ax = axes[1, 0]
    im = ax.imshow(it, extent=ext, aspect="auto", cmap="magma", interpolation="nearest")
    ax.set_title("интенсивность"); fig.colorbar(im, ax=ax, fraction=0.04)

    ax = axes[1, 1]
    im = ax.imshow(gap, extent=ext, aspect="auto", cmap="cividis", vmin=0, vmax=10,
                   interpolation="nearest")
    ax.set_title("разнос эх (дальнее − ближнее), м")
    fig.colorbar(im, ax=ax, fraction=0.04)

    for ax in axes.ravel():
        ax.set_xlabel("азимут, °"); ax.set_ylabel("элевация, °")
        ax.invert_xaxis()
        ax.axvline(az_c, color="w", lw=0.6, ls=":")
        ax.axhline(el_c, color="w", lw=0.6, ls=":")

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


if __name__ == "__main__":
    main()