"""Устойчивость к формату облака: тот же проезд в других обличьях.

Контрольная запись может прийти не в том виде, что выданные: другой драйвер,
другая раскладка, другой порядок и типы полей. Инструмент берёт настоящие
кадры, переписывает их, прогоняет через CDR — тот же разбор, что в узле, — и
через конвейер и сравнивает решения с исходными: сколько кадров обработано,
были ли сбои, совпадает ли тревога и расстояние.

    python tools/check_formats.py
    python tools/check_formats.py --bag data/for_hackathon/doubleT_obstacle
    python tools/check_formats.py --only "без intensity,ось X вперёд"
"""
from __future__ import annotations

import argparse
import struct
import time
from pathlib import Path

import numpy as np

import _bootstrap as B
from flyguard.bag import Bag
from flyguard.cdr import PointCloud2, parse_pointcloud2
from flyguard.mbon_readout import MbonReadout
from flyguard.mushroom_body import MushroomBody
from flyguard.pipeline import FlyGuard, Params

_ROS_TYPE = {("i", 1): 1, ("u", 1): 2, ("i", 2): 3, ("u", 2): 4,
             ("i", 4): 5, ("u", 4): 6, ("f", 4): 7, ("f", 8): 8}


def to_cdr(pc: PointCloud2) -> bytes:
    """Обратно в CDR sensor_msgs/msg/PointCloud2 — в том виде, в каком пишет ROS."""
    out = bytearray(b"\x00\x01\x00\x00")          # заголовок инкапсуляции

    def align(n: int) -> None:
        out.extend(b"\x00" * ((-(len(out) - 4)) % n))

    def u32(v: int) -> None:
        align(4)
        out.extend(struct.pack("<I", v))

    def string(s: str) -> None:
        b = s.encode() + b"\x00"
        u32(len(b))
        out.extend(b)

    sec = int(np.floor(pc.stamp))
    align(4)
    out.extend(struct.pack("<i", sec))
    u32(int(round((pc.stamp - sec) * 1e9)) % 1_000_000_000)
    string(pc.frame_id)
    u32(pc.height)
    u32(pc.width)
    dt = pc.points.dtype
    named = [n for n in dt.names if not n.startswith("_pad")]
    u32(len(named))
    for name in named:
        ft, off = dt.fields[name][:2]
        base = ft.subdtype[0] if ft.subdtype else ft
        string(name)
        u32(off)
        out.append(_ROS_TYPE[(base.kind, base.itemsize)])
        u32(int(np.prod(ft.shape)) if ft.shape else 1)
    out.append(0)                                    # is_bigendian
    u32(dt.itemsize)
    u32(dt.itemsize * pc.width)
    data = np.ascontiguousarray(pc.points).tobytes()
    u32(len(data))
    out.extend(data)
    out.append(1 if pc.is_dense else 0)
    return bytes(out)


# ------------------------------------------------------------------ обличья кадра

def _like(pc: PointCloud2, pts: np.ndarray, *, height: int = 1, dense: bool | None = None,
          stamp: float | None = None) -> PointCloud2:
    return PointCloud2(stamp=pc.stamp if stamp is None else stamp, frame_id=pc.frame_id,
                       height=height, width=pts.size // height,
                       point_step=pts.dtype.itemsize,
                       is_dense=pc.is_dense if dense is None else dense, points=pts)


def _fields(pc: PointCloud2) -> list[str]:
    return [n for n in pc.points.dtype.names if not n.startswith("_pad")]


def _recast(pc: PointCloud2, spec: list[tuple[str, str, int]], itemsize: int) -> np.ndarray:
    """Новая раскладка точки: (имя, тип, смещение); значения — по имени поля."""
    names = [s[0] for s in spec if s[0] in pc.points.dtype.names]
    dt = np.dtype({"names": names,
                   "formats": [s[1] for s in spec if s[0] in names],
                   "offsets": [s[2] for s in spec if s[0] in names],
                   "itemsize": itemsize})
    out = np.zeros(pc.points.size, dt)
    for n in names:
        out[n] = pc.points[n]
    return out


def _drop(pc: PointCloud2, gone: set[str]) -> np.ndarray:
    keep = [n for n in _fields(pc) if n not in gone]
    spec, off = [], 0
    for n in keep:
        t = pc.points.dtype.fields[n][0]
        spec.append((n, t.str, off))
        off += t.itemsize
    return _recast(pc, spec, off)


def _echo(pc: PointCloud2) -> np.ndarray:
    p = pc.points
    return ((p["x"] != 0) | (p["y"] != 0) | (p["z"] != 0)) & np.isfinite(p["x"])


def as_is(pc, k):
    return pc


def nan_for_no_return(pc, k):
    pts = np.array(pc.points)
    miss = ~_echo(pc)
    for n in ("x", "y", "z"):
        pts[n][miss] = np.nan
    return _like(pc, pts, dense=False)


def header_128(pc, k):
    n = pc.points.size
    return _like(pc, np.array(pc.points), height=128) if n % 128 == 0 else pc


def ring_major(pc, k):
    if "ring" not in pc.points.dtype.names:
        return pc
    order = np.argsort(pc.points["ring"], kind="stable")
    return _like(pc, np.array(pc.points)[order])


def reordered_fields(pc, k):
    spec = [("intensity", "<f4", 0), ("ring", "<u2", 4), ("x", "<f4", 8), ("y", "<f4", 12),
            ("z", "<f4", 16), ("timestamp", "<f8", 24)]
    return _like(pc, _recast(pc, spec, 32))


def intensity_u8(pc, k):
    pts = _drop(pc, set())
    spec = [(n, ("|u1" if n == "intensity" else pts.dtype.fields[n][0].str), 0)
            for n in pts.dtype.names]
    off = 0
    for i, (n, t, _) in enumerate(spec):
        spec[i] = (n, t, off)
        off += np.dtype(t).itemsize
    out = _recast(pc, spec, off)
    out["intensity"] = np.clip(np.asarray(pc.points["intensity"], np.float64), 0, 255)
    return _like(pc, out)


def no_intensity(pc, k):
    return _like(pc, _drop(pc, {"intensity"}))


def no_ring(pc, k):
    return _like(pc, _drop(pc, {"ring", "timestamp"}))


def echo_only(pc, k):
    return _like(pc, np.array(pc.points)[_echo(pc)], dense=True)


def echo_only_no_ring(pc, k):
    pts = _drop(pc, {"ring", "timestamp"})
    return _like(pc, pts[_echo(pc)], dense=True)


def shuffled(pc, k):
    rng = np.random.default_rng(k)
    return _like(pc, np.array(pc.points)[rng.permutation(pc.points.size)])


def float64_xyz(pc, k):
    spec, off = [], 0
    for n in _fields(pc):
        t = "<f8" if n in ("x", "y", "z") else pc.points.dtype.fields[n][0].str
        spec.append((n, t, off))
        off += np.dtype(t).itemsize
    return _like(pc, _recast(pc, spec, off))


def x_forward(pc, k):
    """REP-103: x вперёд. У выданных записей вперёд −Y — поворот на 90° вокруг z."""
    pts = np.array(pc.points)
    x, y = pts["x"].copy(), pts["y"].copy()
    pts["x"], pts["y"] = -y, x
    return _like(pc, pts)


def frozen_stamp(pc, k):
    return _like(pc, pc.points, stamp=1000.0)


VARIANTS = {
    "как есть": (as_is, 1),
    "NaN вместо нулей": (nan_for_no_return, 1),
    "шапка 128 × N": (header_128, 1),
    "по кольцам": (ring_major, 1),
    "поля в другом порядке, 32 байта": (reordered_fields, 1),
    "intensity uint8": (intensity_u8, 1),
    "без intensity": (no_intensity, 1),
    "без ring и timestamp": (no_ring, 1),
    "только точки с эхом": (echo_only, 1),
    "только точки с эхом, без ring": (echo_only_no_ring, 1),
    "перемешанный порядок": (shuffled, 1),
    "координаты float64": (float64_xyz, 1),
    "5 Гц": (as_is, 2),
    "ось X вперёд": (x_forward, 1),
    "одинаковое время кадров": (frozen_stamp, 1),
}


def run(bag: Path, n: int, fn, step: int, memory, readout) -> dict:
    # кадры читаются заново для каждого обличья: 24-мегабайтные кадры
    # doubleT_obstacle целиком в память не помещаются
    fg = FlyGuard(Params(device="cpu"), memory=memory, readout=readout)
    out, errors, ms = {}, [], []
    for k, (_, pc) in enumerate(Bag(bag).frames(stop=n)):
        if k % step:
            continue
        try:
            blob = to_cdr(fn(pc, k))
            t0 = time.perf_counter()
            res = fg.process(parse_pointcloud2(blob))
            ms.append((time.perf_counter() - t0) * 1e3)
        except Exception as e:  # noqa: BLE001 — сбой и есть результат проверки
            errors.append(f"{type(e).__name__}: {e}")
            continue
        if res is not None:
            d = res.decision
            out[k] = (bool(d.detected), float(d.distance) if d.detected else np.nan)
    return {"dec": out, "errors": errors, "ms": ms}


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--bag", action="append",
                    help="запись; можно несколько (по умолчанию doubleT_obstacle и синтетика)")
    ap.add_argument("--frames", type=int, default=600, help="сколько кадров с начала")
    ap.add_argument("--only", default="", help="обличья через запятую")
    args = ap.parse_args()
    bags = args.bag or [str(B.DATA / "for_hackathon" / "doubleT_obstacle"),
                        str(B.DATA / "cloud_with_fake_obj")]
    only = [s.strip() for s in args.only.split(",") if s.strip()]
    memory = MushroomBody.load(B.ARTIFACTS / "mushroom_body.npz")
    readout = MbonReadout.load(B.ARTIFACTS / "mbon_readout.npz")

    for bag in map(Path, bags):
        first = next(pc for _, pc in Bag(bag).frames(stop=1))
        print(f"\n== {bag.name}: до {args.frames} кадров, "
              f"поля {', '.join(_fields(first))}, точек {first.points.size}")
        base = run(bag, args.frames, as_is, 1, memory, readout)["dec"]
        n_alarm = sum(v[0] for v in base.values())
        print(f"   исходно: обработано {len(base)}, с тревогой {n_alarm}")
        print(f"   {'обличье':34s} {'обраб.':>6s} {'сбоев':>5s} {'тревог':>6s} "
              f"{'совпало':>7s} {'Δд, м':>6s} {'мс':>5s}")
        for name, (fn, step) in VARIANTS.items():
            if only and name not in only:
                continue
            r = run(bag, args.frames, fn, step, memory, readout)
            dec = r["dec"]
            common = [k for k in dec if k in base]
            same = np.mean([dec[k][0] == base[k][0] for k in common]) if common else np.nan
            dd = [abs(dec[k][1] - base[k][1]) for k in common if dec[k][0] and base[k][0]]
            alarms = sum(v[0] for v in dec.values())
            print(f"   {name:34s} {len(dec):6d} {len(r['errors']):5d} {alarms:6d} "
                  f"{same:7.1%} {np.median(dd) if dd else float('nan'):6.2f} "
                  f"{np.median(r['ms']) if r['ms'] else float('nan'):5.0f}")
            if r["errors"]:
                print(f"      первый сбой: {r['errors'][0][:160]}")


if __name__ == "__main__":
    main()