"""Сравнение способов оценить осевую линию пути: свод тоннеля против полотна.

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

    python tools/compare_centreline.py --bag data/for_hackathon/roundT_doubleT
"""
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 TrackFrame, fit_rail_plane  # noqa: E402
from flyguard.retina import ScanLayout  # noqa: E402


def slices(d_lo=8.0, d_hi=200.0, n=30):
    e = np.geomspace(d_lo, d_hi, n + 1)
    return list(zip(e[:-1], e[1:]))


def bore_centre(tf):
    out = []
    for lo, hi in slices():
        m = tf.valid & (tf.d >= lo) & (tf.d < hi) & (tf.h > 0.6) & (tf.h < 3.2)
        n = int(m.sum())
        if n < 60:
            continue
        u = tf.u[m]
        a, b = np.percentile(u, (3, 97))
        if b - a < 1.5:
            continue
        out.append((0.5 * (lo + hi), 0.5 * (a + b), n))
    return np.array(out) if out else np.zeros((0, 3))


def bed_centre(tf, track=0.0):
    """Центр полотна: полоса вокруг уровня рельсов, узкая по ширине."""
    out = []
    ref = track
    for lo, hi in slices(8.0, 140.0, 26):
        m = (tf.valid & (tf.d >= lo) & (tf.d < hi)
             & (tf.h > -0.20) & (tf.h < 0.45) & (np.abs(tf.u - ref) < 2.6))
        n = int(m.sum())
        if n < 40:
            continue
        u = tf.u[m]
        a, b = np.percentile(u, (5, 95))
        if not (0.8 < b - a < 5.0):
            continue
        c = 0.5 * (a + b)
        out.append((0.5 * (lo + hi), c, n))
        ref = c                       # следующий срез ищем рядом с найденным
    return np.array(out) if out else np.zeros((0, 3))


def fit_arc(d, u, w):
    """u ≈ c1·d + c2·d² (без свободного члена: под поездом ось проходит через 0)."""
    A = np.stack([d, d * d], axis=1) * w[:, None]
    c, *_ = np.linalg.lstsq(A, u * w, rcond=None)
    return c


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--bag", required=True)
    ap.add_argument("--frames", type=int, default=8)
    ap.add_argument("--stride", type=int, default=20)
    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(40.0)
    lay = layout.sub(cols)

    fig, axes = plt.subplots(1, 2, figsize=(13, 5), dpi=110)
    plane = None
    radii_b, radii_d = [], []
    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=plane)
        tf = TrackFrame(img, lay, plane)

        bo = bore_centre(tf)
        be = bed_centre(tf)
        dd = np.linspace(0, 200, 100)
        if bo.shape[0] >= 4:
            u0 = bo[0, 1]
            c = fit_arc(bo[:, 0], bo[:, 1] - u0, np.sqrt(bo[:, 2]))
            axes[0].plot(bo[:, 0], bo[:, 1] - u0, ".", ms=4)
            axes[0].plot(dd, c[0] * dd + c[1] * dd ** 2, lw=1, alpha=0.6)
            radii_b.append(1 / (2 * c[1]) if c[1] else np.inf)
        if be.shape[0] >= 4:
            c = fit_arc(be[:, 0], be[:, 1], np.sqrt(be[:, 2]))
            axes[1].plot(be[:, 0], be[:, 1], ".", ms=4)
            axes[1].plot(dd, c[0] * dd + c[1] * dd ** 2, lw=1, alpha=0.6)
            radii_d.append(1 / (2 * c[1]) if c[1] else np.inf)

    for ax, t, r in ((axes[0], "центр свода тоннеля", radii_b),
                     (axes[1], "центр полотна пути", radii_d)):
        ax.axhline(1.7, color="r", ls="--", lw=0.8); ax.axhline(-1.7, color="r", ls="--", lw=0.8)
        med = np.median(r) if r else np.nan
        spread = np.std(r) if len(r) > 1 else np.nan
        ax.set_title(f"{t}\nR медиана {med:.0f} м, разброс {spread:.0f} м, оценок {len(r)}")
        ax.set_xlabel("вперёд, м"); ax.set_ylabel("смещение оси, м")
        ax.set_ylim(-14, 14); ax.set_xlim(0, 200)
        print(f"{t}: R = " + ", ".join(f"{x:.0f}" for x in r))

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


if __name__ == "__main__":
    main()