"""MEDULLA и LOBULA PLATE — движение: T4/T5, LPTC и LPLC2.

Три схемы из коннектома, работающие подряд:

* **T4/T5** — элементарные детекторы движения. Каждый тип существует в четырёх
  подтипах, настроенных на четыре стороны света в поле зрения; T4 читает ON-канал,
  T5 — OFF. Вычислительно это коррелятор Хассенштайна–Райхардта: сигнал одного
  омматидия задерживается и умножается на сигнал соседнего, разность двух таких
  произведений даёт направленный отклик.
* **LPTC (HS/VS)** — широкопольные тангенциальные клетки лобулярной пластинки.
  Каждая суммирует выход тысяч T4/T5 по своему рецептивному полю и тем самым
  измеряет **собственное движение**. Здесь это критично: колёсной одометрии в
  задаче нет, и скорость поезда неоткуда взять, кроме как из самого потока.
* **LPLC2** — детектор надвигания. Его дендриты разложены на четыре слоя так,
  что клетка отвечает только на поток, расходящийся **из центра её рецептивного
  поля**, и подавляется однородным широкопольным потоком. То есть она отделяет
  «на меня что-то летит» от «я сам двигаюсь». Выход идёт на гигантское волокно.

Практически: продвижение поезда за кадр оценивается корреляцией продольного
профиля тоннеля (устойчиво и дёшево), а T4/T5 и LPLC2 дают карту остаточного
движения — того, что не объясняется собственным ходом.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np

# шаг гистограммы продольного профиля, м
PROFILE_BIN = 0.5
PROFILE_MAX = 250.0
MAX_SPEED = 35.0      # м/с, заведомо выше любого метропоезда
MAX_ACCEL = 3.0       # м/с², предел разгона и экстренного торможения состава
WARMUP_FRAMES = 4     # столько кадров оценка принимается как есть, без фильтра


@dataclass
class EgoMotion:
    """Собственное движение за один кадр."""

    ds: float          # продвижение вперёд, м
    speed: float       # м/с
    yaw_deg: float     # поворот за кадр, °
    conf: float        # 0…1, качество корреляционного пика
    dt: float

    @property
    def kmh(self) -> float:
        return self.speed * 3.6


PROFILE_MIN = 14.0     # ближняя зона в профиль не идёт: там сдвиг кадра не читается


def longitudinal_profile(d: np.ndarray, valid: np.ndarray) -> np.ndarray:
    """Продольная «подпись» тоннеля: сколько лучей оборвалось на каждой дальности.

    Тюбинговые кольца, лотки, ниши и стыки дают ей богатый рисунок, поэтому
    сдвиг профиля между кадрами читается как пройденный путь.

    Две поправки, без которых корреляция залипает на нулевом сдвиге:
    ближняя зона исключается (там на один метр пути приходятся тысячи лучей,
    и её вклад подавляет всё остальное), а от профиля отнимается скользящее
    среднее — остаётся только рисунок структур, без общей огибающей.
    """
    n = int(PROFILE_MAX / PROFILE_BIN)
    dd = d[valid]
    dd = dd[(dd > PROFILE_MIN) & (dd < PROFILE_MAX)]
    if dd.size < 50:
        return np.zeros(n, np.float32)
    p = np.bincount((dd / PROFILE_BIN).astype(np.int32), minlength=n)[:n].astype(np.float32)
    p = np.log1p(p)
    k = 11                                     # ≈5 м — крупнее шага структур
    kern = np.ones(k, np.float32) / k
    env = np.convolve(p, kern, mode="same")
    return (p - env).astype(np.float32)


def _norm(x: np.ndarray) -> np.ndarray:
    x = x - x.mean()
    s = np.linalg.norm(x)
    return x / s if s > 1e-6 else x


def match_shift(prev: np.ndarray, cur: np.ndarray, max_bins: int,
                prior_bins: float | None = None, prior_w: float = 0.03) -> tuple[float, float]:
    """Сдвиг `cur` относительно `prev` по максимуму нормированной корреляции.

    Профиль текущего кадра смещён к меньшим дальностям на пройденный путь,
    поэтому ищется такой сдвиг k, при котором cur[i] ≈ prev[i + k], k ≥ 0.
    Слабый приор по предыдущей скорости снимает неоднозначность на периодических
    структурах вроде тюбинговых колец; его вес нормирован на диапазон поиска,
    чтобы он подправлял выбор между близкими пиками, а не диктовал ответ.

    Корреляция считается по общей части профилей, поэтому при больших сдвигах
    выборка короче — нормировка на длину не даёт этому создать ложный уклон.
    """
    n = prev.size
    ks = np.arange(0, max_bins + 1)
    scores = np.empty(ks.size, np.float32)
    for i, k in enumerate(ks):
        m = n - k
        scores[i] = float(np.dot(_norm(cur[:m]), _norm(prev[k:k + m])))
    if prior_bins is not None and max_bins > 0:
        scores = scores - prior_w * ((ks - prior_bins) / max_bins) ** 2

    i = int(np.argmax(scores))
    peak = float(scores[i])
    if 0 < i < ks.size - 1:               # уточнение параболой по трём точкам
        y0, y1, y2 = scores[i - 1], scores[i], scores[i + 1]
        den = y0 - 2 * y1 + y2
        sub = 0.5 * (y0 - y2) / den if abs(den) > 1e-9 else 0.0
    else:
        sub = 0.0
    return float(ks[i] + np.clip(sub, -1, 1)), peak


class RayIndexer:
    """Обратный поиск по решётке: направление → (кольцо, столбец).

    Нужен, чтобы перепроецировать точки предыдущего кадра в текущую решётку.
    Элевации каналов заданы убывающей таблицей, поэтому индекс кольца берётся
    линейной интерполяцией по ней, а не делением на постоянный шаг.
    """

    def __init__(self, layout):
        el = np.asarray(layout.el_deg, np.float64)
        order = np.argsort(el)
        self.el_sorted = el[order]
        self.ring_sorted = order.astype(np.float64)
        self.az0 = float(layout.az_grid_deg[0])
        self.step = float(layout.az_step_deg)
        self.n_az = int(layout.n_az)
        self.n_rings = int(layout.n_rings)

    def __call__(self, az_deg: np.ndarray, el_deg: np.ndarray):
        col = np.rint((az_deg - self.az0) / self.step).astype(np.int32)
        ring = np.rint(np.interp(el_deg, self.el_sorted, self.ring_sorted)).astype(np.int32)
        ok = (col >= 0) & (col < self.n_az) & (ring >= 0) & (ring < self.n_rings)
        np.clip(col, 0, self.n_az - 1, out=col)
        np.clip(ring, 0, self.n_rings - 1, out=ring)
        return ring, col, ok


def advance_score(prev_pts: np.ndarray, r_cur: np.ndarray, valid_cur: np.ndarray,
                  indexer: RayIndexer, ds: float) -> float:
    """Доля точек прошлого кадра, попавших в текущий кадр при сдвиге вперёд на ds.

    Это и есть проверка широкопольного потока на согласие с моделью собственного
    движения — то, чем заняты тангенциальные клетки лобулярной пластинки.
    """
    d = prev_pts[:, 0] - ds
    u = prev_pts[:, 1]
    h = prev_pts[:, 2]
    m = d > 2.0
    if m.sum() < 50:
        return 0.0
    d, u, h = d[m], u[m], h[m]
    rho = np.hypot(d, u)
    r = np.sqrt(rho * rho + h * h)
    az = np.degrees(np.arctan2(u, d))
    el = np.degrees(np.arctan2(h, rho))
    ring, col, ok = indexer(az, el)
    rc = r_cur[ring, col]
    good = ok & valid_cur[ring, col]
    if good.sum() < 50:
        return 0.0
    err = np.abs(rc[good] - r[good])
    tol = np.maximum(0.25, 0.015 * r[good])
    return float(np.mean(err < tol))


class EgoMotionEstimator:
    """LPTC-аналог: одна широкопольная оценка собственного движения на кадр."""

    def __init__(self, dt_nominal: float = 0.1):
        self.prev_profile: np.ndarray | None = None
        self.prev_az: np.ndarray | None = None
        self.prev_stamp: float | None = None
        self.prev_ds: float | None = None
        self.prev_pts: np.ndarray | None = None
        self.indexer: RayIndexer | None = None
        self.n_sample = 6000
        self.v_filt: float | None = None
        self.n_updates = 0
        self.dt_nominal = dt_nominal

    def update(self, tf, stamp: float) -> EgoMotion:
        dt = self.dt_nominal
        if self.prev_stamp is not None:
            got = stamp - self.prev_stamp
            if 0.01 < got < 1.0:
                dt = got

        prof = longitudinal_profile(tf.d, tf.valid)
        az = np.log1p(tf.valid.sum(axis=0)).astype(np.float32)

        ds, conf, yaw = 0.0, 0.0, 0.0
        if self.prev_profile is not None and np.any(prof):
            max_bins = int(MAX_SPEED * dt / PROFILE_BIN) + 2
            prior = None if self.prev_ds is None else self.prev_ds / PROFILE_BIN
            shift, conf = match_shift(self.prev_profile, prof, max_bins, prior_bins=prior)
            ds = shift * PROFILE_BIN

        # независимая грубая оценка: дальние фронтальные поверхности приближаются
        # ровно на пройденный путь
        direct = self._direct_advance(tf)
        seeds = [s for s in (ds, direct, self.prev_ds, 0.0) if s is not None]

        # уточнение сопоставлением кадров: перебор сдвига с проверкой согласия
        refined = self._refine(tf, seeds, dt)
        if refined is not None:
            ds, conf = refined
        elif direct is not None and conf < 0.45:
            ds, conf = direct, max(conf, 0.3)

        if self.prev_az is not None and self.prev_az.size == az.size:
            yaw_bins, _ = _centred_shift(self.prev_az, az, max_shift=40)
            yaw = yaw_bins * float(tf.layout.az_step_deg)

        ds = self._filter_speed(ds, conf, dt)

        self.prev_profile = prof
        self.prev_az = az
        self.prev_stamp = stamp
        self.prev_r = np.where(tf.valid, tf.r, np.nan).astype(np.float32)
        self.prev_pts = _sample_points(tf, self.n_sample)
        self.prev_ds = ds if self.prev_ds is None else 0.6 * ds + 0.4 * self.prev_ds

        return EgoMotion(ds=ds, speed=ds / dt, yaw_deg=yaw, conf=float(conf), dt=dt)

    def _filter_speed(self, ds: float, conf: float, dt: float) -> float:
        """Сгладить оценку скорости с учётом физики состава.

        Сопоставление кадров иногда «срывается» на резкой смене обстановки —
        например на переходе круглого тоннеля в двухпутный — и выдаёт то ноль,
        то предел диапазона поиска. Поезд так не умеет: за 0.1 с скорость не
        меняется больше чем на a·dt. Измерение принимается с весом, равным
        согласию перепроекции, и ограничивается физическим пределом ускорения,
        поэтому редкий срыв сглаживается, а настоящее торможение отслеживается
        за десятые доли секунды.
        """
        self.n_updates += 1
        v_meas = ds / max(dt, 1e-3)
        if self.v_filt is None or self.n_updates <= WARMUP_FRAMES:
            self.v_filt = v_meas
            return ds
        gain = float(np.clip(conf, 0.05, 0.6))
        v = self.v_filt + gain * (v_meas - self.v_filt)
        limit = MAX_ACCEL * dt
        v = float(np.clip(v, self.v_filt - limit, self.v_filt + limit))
        self.v_filt = max(v, 0.0)
        return self.v_filt * dt

    def _refine(self, tf, seeds: list[float], dt: float):
        """Двухэтапный перебор сдвига вокруг стартовых гипотез."""
        pts = getattr(self, "prev_pts", None)
        if pts is None or pts.shape[0] < 500:
            return None
        if self.indexer is None or self.indexer.n_az != tf.layout.n_az:
            self.indexer = RayIndexer(tf.layout)

        hi = MAX_SPEED * dt
        # гипотезы от дешёвых оценок плюс редкая сетка на весь диапазон;
        # округление до 10 см убирает дубликаты и держит число проб низким
        grid = {round(float(np.clip(s, 0.0, hi)), 1) for s in seeds}
        for s in list(grid):
            grid.update(round(float(np.clip(s + o, 0.0, hi)), 1) for o in (-0.5, 0.5))
        grid.update(round(float(x), 1) for x in np.linspace(0.0, hi, 8))
        cand = np.array(sorted(grid))

        sc = np.array([advance_score(pts, tf.r, tf.valid, self.indexer, s) for s in cand])
        best = float(cand[int(np.argmax(sc))])

        fine = np.clip(best + np.linspace(-0.2, 0.2, 5), 0.0, hi)
        sf = np.array([advance_score(pts, tf.r, tf.valid, self.indexer, s) for s in fine])
        i = int(np.argmax(sf))
        if sf[i] < 0.05:
            return None
        # уточнение параболой по трём точкам вокруг лучшей
        ds = float(fine[i])
        if 0 < i < fine.size - 1:
            y0, y1, y2 = sf[i - 1], sf[i], sf[i + 1]
            den = y0 - 2 * y1 + y2
            if abs(den) > 1e-9:
                ds += 0.5 * (y0 - y2) / den * (fine[1] - fine[0])
        return float(np.clip(ds, 0.0, hi)), float(sf[i])

    def _direct_advance(self, tf, d_lo: float = 35.0, d_hi: float = 200.0,
                        grad_max: float = 0.6) -> float | None:
        """Медианное приближение дальних поверхностей, обращённых к сенсору.

        Стены тоннеля идут почти вдоль движения, и их дальность при езде почти
        не меняется, поэтому они отбрасываются по градиенту дальности вдоль
        строки: остаются только фронтальные поверхности, для которых убывание
        дальности равно пройденному пути.
        """
        cur = np.where(tf.valid, tf.r, np.nan).astype(np.float32)
        prev = getattr(self, "prev_r", None)
        if prev is None or prev.shape != cur.shape:
            return None
        with np.errstate(invalid="ignore"):
            grad = np.abs(np.gradient(cur, axis=1))
            m = (np.isfinite(cur) & np.isfinite(prev) & (cur > d_lo) & (cur < d_hi)
                 & (grad < grad_max))
            if m.sum() < 150:
                return None
            diff = prev[m] - cur[m]
        diff = diff[np.abs(diff) < MAX_SPEED * 0.12]
        if diff.size < 100:
            return None
        return float(np.clip(np.median(diff), 0.0, MAX_SPEED * 0.12))


def _sample_points(tf, n: int) -> np.ndarray:
    """Равномерная выборка точек кадра: (N, 3) = (d, u, z) в системе сенсора.

    Берётся именно z сенсора, а не высота над рельсом: перепроекция идёт в
    решётку лучей, а она задана относительно сенсора.
    """
    m = tf.valid & (tf.d > 8.0) & (tf.d < 200.0)
    idx = np.flatnonzero(m.ravel())
    if idx.size == 0:
        return np.zeros((0, 3), np.float32)
    if idx.size > n:
        idx = idx[:: max(1, idx.size // n)][:n]
    return np.stack([tf.d.ravel()[idx], tf.u.ravel()[idx], tf.z.ravel()[idx]],
                    axis=1).astype(np.float32)


def _centred_shift(prev: np.ndarray, cur: np.ndarray, max_shift: int) -> tuple[float, float]:
    """Сдвиг в обе стороны — для рыскания."""
    a, b = _norm(prev), _norm(cur)
    n = a.size
    ks = np.arange(-max_shift, max_shift + 1)
    sc = np.empty(ks.size, np.float32)
    for i, k in enumerate(ks):
        if k >= 0:
            m = n - k
            sc[i] = float(np.dot(b[:m], a[k:k + m]))
        else:
            m = n + k
            sc[i] = float(np.dot(b[-k:-k + m], a[:m]))
    i = int(np.argmax(sc))
    return float(ks[i]), float(sc[i])


# --------------------------------------------------------------------------- T4/T5

class EmdBank:
    """Коррелятор Хассенштайна–Райхардта на четыре направления.

    Работает на прорежённой решётке: широкопольным клеткам мухи тоже не нужна
    полная разрешающая способность фасеток, им важна статистика по полю.
    """

    DIRECTIONS = ((0, 1), (0, -1), (1, 0), (-1, 0))   # (Δкольцо, Δстолбец)

    def __init__(self, decimate: tuple[int, int] = (2, 4), tau_frames: float = 1.5):
        self.dec = decimate
        self.alpha = float(np.exp(-1.0 / max(tau_frames, 1e-3)))
        self.delayed: np.ndarray | None = None

    def _down(self, a: np.ndarray) -> np.ndarray:
        dh, dw = self.dec
        h = a.shape[0] // dh * dh
        w = a.shape[1] // dw * dw
        return a[:h, :w].reshape(h // dh, dh, w // dw, dw).mean(axis=(1, 3))

    def update(self, signal: np.ndarray) -> np.ndarray:
        """Вернуть (4, h, w) откликов на движение в четырёх направлениях."""
        s = self._down(signal).astype(np.float32)
        if self.delayed is None or self.delayed.shape != s.shape:
            self.delayed = s.copy()
            return np.zeros((4,) + s.shape, np.float32)

        d = self.delayed
        out = np.zeros((4,) + s.shape, np.float32)
        for k, (di, dj) in enumerate(self.DIRECTIONS):
            a = np.roll(s, (di, dj), axis=(0, 1))
            ad = np.roll(d, (di, dj), axis=(0, 1))
            out[k] = d * a - s * ad          # задержанный × соседний, антисимметрично
        self.delayed = self.alpha * d + (1.0 - self.alpha) * s
        return out


def looming(emd: np.ndarray) -> np.ndarray:
    """LPLC2: отклик на поток, расходящийся из центра рецептивного поля.

    Дендриты LPLC2 разложены по четырём слоям так, что каждый слой принимает
    T4/T5 «своего» направления с той стороны поля, куда поток должен уходить при
    надвигании. Сумма четырёх слоёв и есть дивергенция потока.
    """
    right, left, down, up = emd
    div = np.zeros_like(right)
    div[:, 1:-1] += right[:, 2:] - left[:, :-2]
    div[1:-1, :] += down[2:, :] - up[:-2, :]
    return np.maximum(div, 0.0)