"""RETINA — омматидиальная решётка.

Фасеточный глаз дрозофилы — регулярная решётка омматидиев, каждый смотрит в свою
фиксированную сторону. Вращающийся лидар устроен так же: пара (кольцо, столбец)
задаёт направление луча. Поэтому облако точек сразу переводится в *ретинотопический*
дальностный образ `(кольцо, азимут)`, и вся дальнейшая обработка идёт в этой
решётке — как в зрительной системе мухи, а не в неупорядоченном облаке.

Три особенности конкретного сенсора, измеренные по данным (см. `docs/ALGORITHM.md`):

1. **Раскладка различается между бэгами**: 3600 азимутов на 360° против 1200 на
   100°. Решётка калибруется по самим данным, ничего не захардкожено.
2. **Двойное эхо**: соседние столбцы делят один азимут. Когда эхо одно, оба слота
   содержат одно значение; когда два — ближнее несёт объект, дальнее фон за ним.
   Реально различаются ~3 % лучей, и это именно тонкие предметы и кромки.
3. **Скос решётки**: у каждого лазерного канала свой постоянный азимутальный сдвиг,
   разброс достигает **15.5°** (≈155 столбцов). В сыром виде «столбец» не является
   направлением: соседние кольца одного столбца смотрят в стороны, разнесённые на
   градусы. Поэтому образ **выпрямляется** целочисленным сдвигом строк; остаточная
   ошибка < половины шага азимута и учитывается в таблице направлений.
"""
from __future__ import annotations

import warnings
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path

import numpy as np

from .cdr import PointCloud2

DEG = np.pi / 180.0


@dataclass
class RangeImage:
    """Выпрямленный дальностный образ в координатах `(кольцо, азимут)`."""

    stamp: float
    r_near: np.ndarray   # (H, W) float32 — ближнее эхо, 0 = нет эха
    r_far: np.ndarray    # (H, W) float32 — дальнее эхо, 0 = нет эха
    inten: np.ndarray    # (H, W) float32 — интенсивность ближнего эха
    valid: np.ndarray    # (H, W) bool

    @property
    def shape(self) -> tuple[int, int]:
        return self.r_near.shape

    def xyz(self, layout: "ScanLayout") -> np.ndarray:
        """Декартовы координаты ближнего эха: (H, W, 3)."""
        return layout.dirs * self.r_near[..., None]

    def crop(self, cols: slice) -> "RangeImage":
        return RangeImage(self.stamp, self.r_near[:, cols], self.r_far[:, cols],
                          self.inten[:, cols], self.valid[:, cols])


class ScanLayout:
    """Калиброванная и выпрямленная решётка лучей.

    Азимут отсчитывается от направления движения (вперёд = −Y в системе сенсора),
    положительный вправо; элевация — вверх от горизонтали сенсора.
    """

    def __init__(self, el_deg: np.ndarray, az_step_deg: float, az0_deg: float,
                 col_shift: np.ndarray, az_resid_deg: np.ndarray,
                 n_az: int, n_echo: int, wrap: bool = False):
        self.el_deg = np.asarray(el_deg, np.float64)            # (H,)
        self.az_step_deg = float(az_step_deg)                   # шаг на азимутальный индекс
        self.az0_deg = float(az0_deg)
        self.col_shift = np.asarray(col_shift, np.int64)        # (H,) выпрямление
        self.az_resid_deg = np.asarray(az_resid_deg, np.float64)  # (H,) остаток < шага/2
        self.n_az = int(n_az)
        self.n_echo = int(n_echo)
        self.wrap = bool(wrap)          # круговой скан: края смыкаются
        self.n_rings = self.el_deg.size

        self.az_grid_deg = self.az0_deg + self.az_step_deg * np.arange(self.n_az)

        # карта выборки для выпрямления: out[h, j] = raw[h, j + shift[h]]
        j = np.arange(self.n_az)[None, :]
        src = j + self.col_shift[:, None]
        if self.wrap:
            # у кругового скана «выехавшие» столбцы приходят с другого края
            self.gather = np.mod(src, self.n_az).astype(np.intp)
            self.gather_ok = np.ones_like(self.gather, bool)
        else:
            self.gather_ok = (src >= 0) & (src < self.n_az)
            self.gather = np.clip(src, 0, self.n_az - 1).astype(np.intp)

        self.az_full_deg = self.az_grid_deg[None, :] + self.az_resid_deg[:, None]
        self.dirs = self._unit_dirs().astype(np.float32)        # (H, W, 3)
        # угловой шаг по кольцам: у решётки из одного кольца градиента нет
        self.el_step_deg = (np.abs(np.gradient(self.el_deg)) if self.n_rings > 1
                            else np.full(self.n_rings, 0.125))

    # ------------------------------------------------------------------ геометрия

    def _unit_dirs(self) -> np.ndarray:
        az = self.az_full_deg * DEG
        el = (self.el_deg[:, None] * DEG) * np.ones_like(az)
        c = np.cos(el)
        return np.stack([c * np.sin(az), -c * np.cos(az), np.sin(el)], axis=-1)

    def column_slice(self, half_fov_deg: float) -> slice:
        """Непрерывный диапазон столбцов внутри ±half_fov по азимуту."""
        inside = np.flatnonzero(np.abs(self.az_grid_deg) <= half_fov_deg)
        if inside.size == 0:
            return slice(0, self.n_az)
        return slice(int(inside[0]), int(inside[-1]) + 1)

    def sub(self, cols: slice) -> "ScanLayout":
        """Урезанная по азимуту копия решётки (для обработки только переднего сектора)."""
        start = cols.start or 0
        stop = cols.stop if cols.stop is not None else self.n_az
        out = ScanLayout.__new__(ScanLayout)
        out.el_deg = self.el_deg
        out.az_step_deg = self.az_step_deg
        out.az0_deg = float(self.az_grid_deg[start])
        out.col_shift = self.col_shift
        out.az_resid_deg = self.az_resid_deg
        out.n_az = stop - start
        out.n_echo = self.n_echo
        out.wrap = False          # вырезанный сектор больше не смыкается
        out.n_rings = self.n_rings
        out.az_grid_deg = self.az_grid_deg[start:stop]
        out.gather_ok = self.gather_ok[:, start:stop]
        out.gather = self.gather[:, start:stop]
        out.az_full_deg = self.az_full_deg[:, start:stop]
        out.dirs = np.ascontiguousarray(self.dirs[:, start:stop])
        out.el_step_deg = self.el_step_deg
        return out

    # ------------------------------------------------------------------ проекция

    def project(self, pc: PointCloud2, cols: slice | None = None) -> RangeImage:
        """Облако точек → выпрямленный дальностный образ.

        `cols` задаёт нужный диапазон **выходных** столбцов. Тяжёлая
        арифметика (корень по 900 тыс. точек) выполняется только над теми
        сырыми столбцами, которые в этот диапазон попадут с учётом скоса
        каналов, — на круговом скане это экономит почти всё время стадии.
        """
        n, w, e = self.n_rings, self.n_az, self.n_echo
        start = 0 if cols is None else (cols.start or 0)
        stop = w if cols is None else (cols.stop if cols.stop is not None else w)

        lo = start + int(self.col_shift.min())
        hi = stop + int(self.col_shift.max())
        if self.wrap:
            raw_cols = np.arange(lo, hi) % w
        else:
            lo = max(lo, 0)
            hi = min(hi, w)
            raw_cols = None

        pts = pc.points

        def cube(name: str) -> np.ndarray:
            a = pts[name].reshape(w, e, n)
            a = a[raw_cols] if raw_cols is not None else a[lo:hi]
            return a.transpose(2, 0, 1)

        x, y, z = cube("x"), cube("y"), cube("z")
        good = (x != 0) | (y != 0) | (z != 0)
        r = np.sqrt(x * x + y * y + z * z, dtype=np.float32)
        r *= good

        if e == 1:
            r_near = r[..., 0]
            r_far = r[..., 0]
            it = cube("intensity")[..., 0]
            valid = good[..., 0]
        else:
            inten = cube("intensity")
            near_i = np.argmin(np.where(good, r, np.float32(np.inf)), axis=-1)[..., None]
            far_i = np.argmax(r, axis=-1)[..., None]
            r_near = np.take_along_axis(r, near_i, -1)[..., 0]
            r_far = np.take_along_axis(r, far_i, -1)[..., 0]
            it = np.take_along_axis(inten, near_i, -1)[..., 0]
            valid = good.any(axis=-1)

        # выпрямление скоса каналов, с поправкой на смещение окна
        g = self.gather[:, start:stop]
        ok = self.gather_ok[:, start:stop]
        if raw_cols is not None:
            g = (g - lo) % w
        else:
            g = g - lo
            ok = ok & (g >= 0) & (g < (hi - lo))
            np.clip(g, 0, hi - lo - 1, out=g)

        r_near = np.take_along_axis(r_near, g, 1)
        r_far = np.take_along_axis(r_far, g, 1)
        it = np.take_along_axis(it, g, 1)
        valid = np.take_along_axis(valid, g, 1) & ok

        r_near = np.where(valid, r_near, np.float32(0.0))
        r_far = np.where(valid, r_far, np.float32(0.0))
        return RangeImage(pc.stamp, np.ascontiguousarray(r_near),
                          np.ascontiguousarray(r_far), np.ascontiguousarray(it),
                          np.ascontiguousarray(valid))

    # ------------------------------------------------------------------ калибровка

    @staticmethod
    def calibrate(clouds: list[PointCloud2], n_rings: int | None = None) -> "ScanLayout":
        """Восстановить решётку по нескольким кадрам.

        Определяются: число колец и эх, элевация каждого кольца, шаг развёртки,
        азимутальный сдвиг каждого канала и целочисленное выпрямление образа.
        """
        if not clouds:
            raise ValueError("нужен хотя бы один кадр для калибровки")
        pc0 = clouds[0]
        if n_rings is None:
            n_rings = int(pc0.points["ring"].max()) + 1
        n_cols, rem = divmod(pc0.n_points, n_rings)
        if rem:
            raise ValueError(f"{pc0.n_points} точек не делится на {n_rings} колец")

        # Направление луча задано сенсором и в каждом кадре одно и то же:
        # кадры нужны только чтобы закрыть лучи, не вернувшие эхо. Поэтому
        # берётся первое конечное значение, а не медиана по стопке кадров —
        # та стоила 1.4 с из 2.8 с всей калибровки и ничего не уточняла:
        # ниже и шаг развёртки, и сдвиг канала берутся медианой по тысячам
        # столбцов, где шум одного отсчёта всё равно усредняется.
        az = el = None
        for pc in clouds:
            x = pc.points["x"].reshape(n_cols, n_rings).T.astype(np.float32)
            y = pc.points["y"].reshape(n_cols, n_rings).T.astype(np.float32)
            z = pc.points["z"].reshape(n_cols, n_rings).T.astype(np.float32)
            ok = (x != 0) | (y != 0) | (z != 0)
            with _quiet():
                a = np.where(ok, np.degrees(np.arctan2(x, -y)), np.nan)
                e = np.where(ok, np.degrees(np.arctan2(z, np.hypot(x, y))), np.nan)
            if az is None:
                az, el = a, e
                continue
            gap = np.isnan(az)
            if not gap.any():
                break
            az[gap] = a[gap]
            el[gap] = e[gap]
        az = az.astype(np.float64)
        el = el.astype(np.float64)

        n_echo = _detect_echoes(az)
        n_az = n_cols // n_echo
        if n_echo > 1:
            with _quiet():
                az = np.nanmean(az.reshape(n_rings, n_az, n_echo), axis=2)
                el = np.nanmean(el.reshape(n_rings, n_az, n_echo), axis=2)

        # общий шаг развёртки: медиана по кольцам от робастной оценки наклона
        slopes = []
        for h in range(n_rings):
            row = az[h]
            idx = np.flatnonzero(np.isfinite(row))
            if idx.size < 50:
                continue
            d = np.diff(np.unwrap(np.radians(row[idx]))) / np.diff(idx)
            slopes.append(np.median(np.degrees(d)))
        if not slopes:
            raise ValueError("недостаточно валидных лучей для калибровки развёртки")
        step = float(np.median(slopes))

        # смещение каждого канала относительно общей развёртки
        j = np.arange(n_az, dtype=np.float64)
        base = step * j
        with _quiet():
            c_ring = np.nanmedian(_wrap180(az - base[None, :]), axis=1)   # (H,)
        c_ring = _fill_linear(c_ring)
        c0 = float(np.median(c_ring))

        # круговой скан: развёртка покрывает полные 360°
        wrap = abs(step) * n_az > 350.0

        # начало отсчёта выбирается так, чтобы «вперёд» (азимут 0) был в середине
        # образа — иначе шов ±180° разрезал бы рабочий сектор пополам
        if wrap:
            k = int(np.rint(-c0 / step)) - n_az // 2
            c0 = _wrap180(c0 + step * k)
        # разница берётся по кратчайшей дуге: иначе шов ±180° даёт сдвиг в пол-оборота
        shift = np.rint(_wrap180(c0 - c_ring) / step).astype(np.int64)
        resid = _wrap180(c_ring + step * shift - c0)

        with _quiet():
            el_ring = _fill_linear(np.nanmedian(el, axis=1))
        return ScanLayout(el_ring, step, float(c0), shift, resid, n_az, n_echo, wrap)

    # ------------------------------------------------------------------ сериализация

    def save(self, path: str | Path) -> None:
        np.savez_compressed(path, el_deg=self.el_deg, az_step_deg=self.az_step_deg,
                            az0_deg=self.az0_deg, col_shift=self.col_shift,
                            az_resid_deg=self.az_resid_deg, n_az=self.n_az,
                            n_echo=self.n_echo, wrap=self.wrap)

    @staticmethod
    def load(path: str | Path) -> "ScanLayout":
        d = np.load(path)
        return ScanLayout(d["el_deg"], float(d["az_step_deg"]), float(d["az0_deg"]),
                          d["col_shift"], d["az_resid_deg"], int(d["n_az"]),
                          int(d["n_echo"]), bool(d["wrap"]) if "wrap" in d else False)

    def __repr__(self) -> str:
        return (f"ScanLayout(колец={self.n_rings}, азимутов={self.n_az}, эх={self.n_echo}, "
                f"сектор {self.az_grid_deg.min():.1f}°…{self.az_grid_deg.max():.1f}°, "
                f"шаг {abs(self.az_step_deg):.3f}°, "
                f"элевация {self.el_deg.min():.1f}°…{self.el_deg.max():.1f}°, "
                f"скос каналов {np.ptp(self.col_shift)} стлб)")


# ---------------------------------------------------------------------- вспомогательное

@contextmanager
def _quiet():
    """Пустые срезы и деление на ноль при калибровке — норма, а не ошибка."""
    with warnings.catch_warnings(), np.errstate(all="ignore"):
        warnings.simplefilter("ignore", RuntimeWarning)
        yield


def _wrap180(a):
    """Привести угол(ы) в градусах к полуинтервалу (−180, 180]."""
    return -((-np.asarray(a, np.float64) + 180.0) % 360.0 - 180.0)


def _detect_echoes(az: np.ndarray) -> int:
    """Двойное эхо: соседние столбцы делят азимут (проверяется по каждому кольцу)."""
    if az.shape[1] < 4:
        return 1
    a, b = az[:, 0::2], az[:, 1::2]
    m = np.isfinite(a) & np.isfinite(b)
    if m.sum() < 100:
        return 1
    return 2 if np.mean(np.abs(a[m] - b[m]) < 1e-3) > 0.8 else 1


def _fill_linear(a: np.ndarray) -> np.ndarray:
    """Линейно достроить NaN-пропуски по индексу (сетка равномерная)."""
    a = np.asarray(a, np.float64).copy()
    bad = ~np.isfinite(a)
    if bad.all():
        raise ValueError("нет ни одного валидного угла для калибровки")
    if bad.any():
        idx = np.arange(a.size)
        a[bad] = np.interp(idx[bad], idx[~bad], a[~bad])
    return a