"""LAMINA — локальный контраст, разделение ON/OFF.

Первый нейропиль за фоторецепторами. Клетки L1 и L2 получают один и тот же вход
от R1–R6 и расходятся на два канала: L1 → ON (стало ярче), L2 → OFF (стало
темнее). Оба канала предварительно проходят **латеральное торможение** от
амакриновых клеток и Dm9 — классическое «центр минус окружение», которое
подавляет ровный фон и оставляет только локальные отклонения.

Перенос на лидар:

* «Яркость» — это **диспаритет** δ = 1/R, а не сама дальность. Так правильно
  по двум причинам: угловой размер предмета пропорционален 1/R, и шум лидара
  в диспаритете почти однороден, тогда как в дальности растёт квадратично.
* **ON** = объект ближе своего окружения — выступ, то есть препятствие.
* **OFF** = дальше окружения или эха нет вовсе — провал, то есть окклюзионная
  тень **за** препятствием. Тень часто во много раз крупнее самого предмета,
  и именно она даёт шанс увидеть мелкий объект на большой дальности.
* Размер предмета в лучах меняется с дальностью на два порядка, поэтому
  окружение берётся **на нескольких масштабах** сразу — как у колонковых
  нейронов лобулы с разными размерами рецептивных полей, сходящихся на один
  нисходящий нейрон.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.ndimage import uniform_filter

# (радиус центра, радиус окружения) в лучах: кольца × столбцы
SCALES: tuple[tuple[int, int], ...] = ((1, 6), (3, 14), (7, 30))


@dataclass
class LaminaOutput:
    """Каналы ламины для одного кадра."""

    disp: np.ndarray        # (H, W) диспаритет 1/R, 0 там, где эха нет
    on: np.ndarray          # (H, W) ON-контраст, максимум по масштабам, 1/м
    off: np.ndarray         # (H, W) OFF-контраст, 1/м
    on_scale: np.ndarray    # (H, W) int8 — на каком масштабе отклик максимален
    surround: np.ndarray    # (H, W) диспаритет окружения на среднем масштабе
    hole: np.ndarray        # (H, W) доля «нет эха» в окрестности


def _masked_mean(v: np.ndarray, m: np.ndarray, size: tuple[int, int]) -> np.ndarray:
    """Среднее по прямоугольному окну только по валидным отсчётам."""
    num = uniform_filter(v, size=size, mode="nearest")
    den = uniform_filter(m, size=size, mode="nearest")
    return num, den


def _annulus_mean(v: np.ndarray, m: np.ndarray, r_in: int, r_out: int):
    """Среднее по кольцу: большое окно минус вырезанный центр.

    Реализовано через два равномерных фильтра (каждый разделим и работает за
    O(N)), поэтому стоимость не зависит от размера окна.
    """
    s_in = (2 * r_in + 1, 2 * r_in + 1)
    s_out = (2 * r_out + 1, 4 * r_out + 1)     # шире по азимуту: решётка анизотропна
    n_in = s_in[0] * s_in[1]
    n_out = s_out[0] * s_out[1]

    num_i, den_i = _masked_mean(v, m, s_in)
    num_o, den_o = _masked_mean(v, m, s_out)

    num = num_o * n_out - num_i * n_in
    den = den_o * n_out - den_i * n_in
    out = np.divide(num, den, out=np.zeros_like(num), where=den > 0.5)
    return out, den


def process(r: np.ndarray, valid: np.ndarray, *, r_max: float = 300.0) -> LaminaOutput:
    """Посчитать ON/OFF-каналы ламины по дальностному образу."""
    v = valid.astype(np.float32)
    disp = np.zeros_like(r, dtype=np.float32)
    np.divide(1.0, r, out=disp, where=valid & (r > 0.05))
    disp *= v

    on = np.zeros_like(disp)
    off = np.zeros_like(disp)
    on_scale = np.zeros(disp.shape, np.int8)
    surround_mid = None

    for k, (r_in, r_out) in enumerate(SCALES):
        sur, cnt = _annulus_mean(disp, v, r_in, r_out)
        enough = cnt > 8.0
        c = np.where(enough, disp - sur, 0.0)
        pos = np.maximum(c, 0.0) * v            # ближе окружения
        # провал считается и там, где эха нет: 1/∞ = 0 — это тоже сигнал
        neg = np.maximum(-(disp - sur), 0.0) * enough
        better = pos > on
        on = np.where(better, pos, on)
        on_scale = np.where(better, np.int8(k), on_scale)
        off = np.maximum(off, neg)
        if k == 1:
            surround_mid = sur

    # доля лучей без эха в окрестности — мера «дыры» в поверхности
    hole = 1.0 - uniform_filter(v, size=(5, 15), mode="nearest")

    # диспаритет физически ограничен снизу дальностью прибора
    np.clip(on, 0.0, 1.0 / max(r_max, 1.0) * 1e4, out=on)
    return LaminaOutput(disp=disp, on=on, off=off, on_scale=on_scale,
                        surround=surround_mid if surround_mid is not None else np.zeros_like(disp),
                        hole=hole.astype(np.float32))


def contrast_to_depth_gap(on: np.ndarray, r: np.ndarray) -> np.ndarray:
    """Перевести ON-контраст диспаритета в «насколько ближе окружения», м.

    δ − δ_sur = 1/R − 1/R_sur ⇒ R_sur − R = on · R · R_sur. Для оценки берётся
    R_sur = R/(1 − on·R), что даёт разрыв по глубине в метрах.
    """
    x = np.clip(on * r, 0.0, 0.999)
    with np.errstate(divide="ignore", invalid="ignore"):
        gap = r * x / (1.0 - x)
    return np.nan_to_num(gap, nan=0.0, posinf=1e4).astype(np.float32)