"""Синтетические препятствия: трассировка лучей в реальные кадры.

Разметки в датасете нет, а организаторы прямо предупредили, что приватный тест
собран добавлением синтезированных препятствий в новые проезды. Поэтому свой
полигон строится тем же способом: берётся настоящий кадр пустого тоннеля,
в него трассировкой лучей вставляется предмет заданного размера на заданной
дистанции, и получается **размеченный** пример с точно известным ответом.

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

Модель сенсора намеренно пессимистична: добавляется шум дальности, вероятность
несостоявшегося эха растёт с расстоянием и с углом падения, а интенсивность
считается по ламбертовой модели. Лучше недооценить свой детектор, чем на
защите обнаружить, что полигон был слишком добрым.
"""
from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np

from .cdr import PointCloud2
from .geometry import RailPlane
from .retina import ScanLayout

INF = np.float32(np.inf)


# --------------------------------------------------------------------------- тела

@dataclass
class Box:
    """Параллелепипед в координатах пути, стоящий на плоскости рельсов."""

    length: float      # вдоль пути, м
    width: float       # поперёк, м
    height: float      # вверх от рельса, м
    h_base: float = 0.0

    def intersect(self, od, ou, oh, dd, du, dh) -> np.ndarray:
        lo = np.array([-self.length / 2, -self.width / 2, self.h_base], np.float32)
        hi = np.array([self.length / 2, self.width / 2, self.h_base + self.height], np.float32)
        t0 = np.full(od.shape, -INF, np.float32)
        t1 = np.full(od.shape, INF, np.float32)
        for o, d, a, b in ((od, dd, lo[0], hi[0]), (ou, du, lo[1], hi[1]),
                           (oh, dh, lo[2], hi[2])):
            with np.errstate(divide="ignore", invalid="ignore"):
                ta = (a - o) / d
                tb = (b - o) / d
            lo_t = np.where(d != 0, np.minimum(ta, tb), np.where((o >= a) & (o <= b), -INF, INF))
            hi_t = np.where(d != 0, np.maximum(ta, tb), np.where((o >= a) & (o <= b), INF, -INF))
            t0 = np.maximum(t0, lo_t)
            t1 = np.minimum(t1, hi_t)
        hit = (t1 >= np.maximum(t0, 0.0)) & np.isfinite(t0)
        return np.where(hit, np.maximum(t0, 0.0), INF)


@dataclass
class Cylinder:
    """Вертикальный цилиндр — человек, столб, ведро."""

    radius: float
    height: float
    h_base: float = 0.0

    def intersect(self, od, ou, oh, dd, du, dh) -> np.ndarray:
        a = dd * dd + du * du
        b = 2.0 * (od * dd + ou * du)
        c = od * od + ou * ou - self.radius ** 2
        disc = b * b - 4 * a * c
        ok = (disc > 0) & (a > 1e-9)
        sq = np.sqrt(np.where(ok, disc, 0.0))
        with np.errstate(divide="ignore", invalid="ignore"):
            t = (-b - sq) / (2 * a)
            t2 = (-b + sq) / (2 * a)
        t = np.where(t > 0, t, t2)
        h = oh + t * dh
        ok &= (t > 0) & (h >= self.h_base) & (h <= self.h_base + self.height)
        return np.where(ok, t, INF)


@dataclass
class Sphere:
    radius: float
    h_centre: float

    def intersect(self, od, ou, oh, dd, du, dh) -> np.ndarray:
        oz = oh - self.h_centre
        a = dd * dd + du * du + dh * dh
        b = 2.0 * (od * dd + ou * du + oz * dh)
        c = od * od + ou * ou + oz * oz - self.radius ** 2
        disc = b * b - 4 * a * c
        ok = disc > 0
        sq = np.sqrt(np.where(ok, disc, 0.0))
        with np.errstate(divide="ignore", invalid="ignore"):
            t = (-b - sq) / (2 * a)
        return np.where(ok & (t > 0), t, INF)


@dataclass
class ObjectModel:
    """Предмет: набор тел плюс отражательные свойства."""

    name: str
    parts: list = field(default_factory=list)
    reflectivity: float = 0.3     # 0…1, доля отражённого света

    def intersect(self, od, ou, oh, dd, du, dh) -> np.ndarray:
        t = np.full(od.shape, INF, np.float32)
        for p in self.parts:
            t = np.minimum(t, p.intersect(od, ou, oh, dd, du, dh))
        return t

    @property
    def size(self) -> tuple[float, float]:
        """Грубые габариты (ширина, высота) для отчётов."""
        w = h = 0.0
        for p in self.parts:
            if isinstance(p, Box):
                w = max(w, p.width); h = max(h, p.h_base + p.height)
            elif isinstance(p, Cylinder):
                w = max(w, 2 * p.radius); h = max(h, p.h_base + p.height)
            elif isinstance(p, Sphere):
                w = max(w, 2 * p.radius); h = max(h, p.h_centre + p.radius)
        return w, h


def catalogue() -> dict[str, ObjectModel]:
    """Набор предметов, встречающихся в тоннеле, от крупных к мелким."""
    return {
        "человек_стоя": ObjectModel("человек_стоя", [
            Cylinder(radius=0.22, height=1.45),
            Sphere(radius=0.11, h_centre=1.60)], reflectivity=0.35),
        "человек_сидя": ObjectModel("человек_сидя", [
            Box(0.45, 0.50, 0.85)], reflectivity=0.35),
        "ящик": ObjectModel("ящик", [Box(0.60, 0.60, 0.60)], reflectivity=0.40),
        "чемодан": ObjectModel("чемодан", [Box(0.25, 0.45, 0.55)], reflectivity=0.30),
        "ведро": ObjectModel("ведро", [Cylinder(radius=0.15, height=0.35)], reflectivity=0.45),
        "камень": ObjectModel("камень", [Sphere(radius=0.11, h_centre=0.11)], reflectivity=0.20),
        "бутылка": ObjectModel("бутылка", [Cylinder(radius=0.045, height=0.30)],
                               reflectivity=0.25),
        "кабель": ObjectModel("кабель", [Box(2.20, 0.06, 0.06)], reflectivity=0.15),
        "каска": ObjectModel("каска", [Sphere(radius=0.14, h_centre=0.10)], reflectivity=0.55),
    }


# --------------------------------------------------------------------------- вставка

@dataclass
class Placement:
    """Где стоит предмет."""

    d: float                 # вперёд от сенсора, м
    u: float = 0.0           # поперёк от оси пути, м
    yaw_deg: float = 0.0     # поворот вокруг вертикали (для вытянутых тел)


def ray_dirs_track(layout: ScanLayout, plane: RailPlane):
    """Направления лучей в координатах пути: (dd, du, dh) и начало (0, 0, H)."""
    dirs = layout.dirs
    dd = -dirs[..., 1].astype(np.float32)
    du = dirs[..., 0].astype(np.float32)
    dh = (dirs[..., 2] - plane.a * dd - plane.b * du).astype(np.float32)
    return dd, du, dh, np.float32(plane.height)


# Паспортные данные Pandar128E3X (руководство v4p5, п. 1.4 и Приложение A):
# дальность 0.3…200 м при отражательной способности 10 %, вероятность
# обнаружения на паспортной дальности PoD = 70 %, точность ±2 см на 1…200 м.
SPEC_REFLECTIVITY = 0.10
SPEC_POD = 0.70
HARD_MAX_RANGE = 230.0          # дальше в датасете возвратов не встречается

# Расходимость луча в руководстве не указана; принята равной угловому шагу
# решётки (0.1° по азимуту, 0.125° по элевации в полосе высокого разрешения) —
# это верхняя оценка, дающая консервативный результат для мелких предметов.
BEAM_DIV_H = 1.75e-3            # рад
BEAM_DIV_V = 2.18e-3            # рад

_CHANNEL_RANGE: np.ndarray | None = None


def channel_max_range() -> np.ndarray:
    """Паспортная дальность каждого канала при 10 % отражения, (128,).

    Каналы сильно неравноправны: 34–65 «дальнобойные» и берут 200 м, а 98–128
    смотрят в землю и рассчитаны только на ближнее и среднее поле. Без учёта
    этого синтетический полигон завышал бы дальность обнаружения для предметов,
    попадающих в нижние каналы.
    """
    global _CHANNEL_RANGE
    if _CHANNEL_RANGE is None:
        import csv
        from pathlib import Path
        path = Path(__file__).with_name("data") / "pandar128_channels.csv"
        try:
            rows = sorted(csv.DictReader(open(path, encoding="utf-8")),
                          key=lambda r: int(r["channel"]))
            _CHANNEL_RANGE = np.array([float(r["max_range_10pct_m"]) for r in rows],
                                      np.float32)
        except (OSError, KeyError, ValueError):
            _CHANNEL_RANGE = np.full(128, 200.0, np.float32)
    return _CHANNEL_RANGE


def dropout_probability(r: np.ndarray, reflectivity: float,
                        ang_w: np.ndarray | float = 1.0,
                        ang_h: np.ndarray | float = 1.0,
                        max_range: np.ndarray | float = 200.0) -> np.ndarray:
    """Вероятность, что эхо не вернётся.

    Принятая мощность падает как ρ·A/r², где A — доля пятна луча, закрытая
    предметом. Отсюда «эффективная дальность» r·√(ρ_паспорт/(ρ·A)), которую
    остаётся сравнить с паспортной дальностью **этого канала**. Переход сделан
    логистическим и смещён так, чтобы на паспортной дальности вероятность
    обнаружения равнялась заявленным 70 %.

    Заполнение пятна важно именно для мелких предметов: на 200 м луч шириной
    1.75 мрад покрывает 35 см, и бутылка диаметром 9 см отражает лишь четверть
    его энергии — поэтому она пропадает намного раньше человека, хотя по
    геометрии в неё ещё попадают лучи.
    """
    fill = np.clip(ang_w / BEAM_DIV_H, 0.05, 1.0) * np.clip(ang_h / BEAM_DIV_V, 0.05, 1.0)
    rho = max(reflectivity, 0.02) * fill
    eff = r * np.sqrt(SPEC_REFLECTIVITY / rho)
    width = 0.12 * np.asarray(max_range, np.float32)
    r50 = np.asarray(max_range, np.float32) + width * np.log(SPEC_POD / (1 - SPEC_POD))
    p_detect = 1.0 / (1.0 + np.exp((eff - r50) / np.maximum(width, 1e-3)))
    p_detect = np.where(r > HARD_MAX_RANGE, 0.0, p_detect)
    return np.clip(1.0 - p_detect, 0.0, 1.0).astype(np.float32)


def inject(pc: PointCloud2, layout: ScanLayout, plane: RailPlane,
           obj: ObjectModel, place: Placement, *,
           rng: np.random.Generator | None = None,
           range_noise: float = 0.02, cols: slice | None = None) -> tuple[PointCloud2, dict]:
    """Вставить предмет в облако точек. Возвращает (новое облако, разметка)."""
    rng = rng or np.random.default_rng()
    n_rings, n_az, n_echo = layout.n_rings, layout.n_az, layout.n_echo

    dd, du, dh, H = ray_dirs_track(layout, plane)
    # начало луча в системе предмета: сенсор в (0,0,H), предмет в (d, u, 0)
    od = np.full(dd.shape, -np.float32(place.d), np.float32)
    ou = np.full(dd.shape, -np.float32(place.u), np.float32)
    oh = np.full(dd.shape, H, np.float32)

    if place.yaw_deg:
        c, s = np.cos(np.radians(place.yaw_deg)), np.sin(np.radians(place.yaw_deg))
        od, ou = c * od + s * ou, -s * od + c * ou
        dd, du = c * dd + s * du, -s * dd + c * du

    t = obj.intersect(od, ou, oh, dd, du, dh)
    hit = np.isfinite(t) & (t > 1.0)
    if not hit.any():
        return pc, dict(hit_rays=0, d=place.d, u=place.u, name=obj.name)

    # шум дальности и пропуски эха
    t = t + rng.normal(0.0, range_noise, t.shape).astype(np.float32)
    w_obj, h_obj = obj.size
    ang_w = w_obj / np.maximum(t, 1.0)
    ang_h = h_obj / np.maximum(t, 1.0)
    ch_range = channel_max_range()
    per_ray_range = (ch_range[:n_rings, None] if ch_range.size >= n_rings
                     else np.float32(200.0))
    p_drop = dropout_probability(t, obj.reflectivity, ang_w, ang_h, per_ray_range)
    hit &= rng.random(t.shape) > p_drop

    pts = pc.points.copy()
    ring_i, col_i = np.nonzero(hit)
    raw_col = layout.gather[ring_i, col_i]        # выпрямленный столбец → сырой
    new_r = t[ring_i, col_i]

    # плоский индекс точки: (сырой столбец · число эх + эхо) · число колец + кольцо
    base = (raw_col * n_echo) * n_rings + ring_i
    fx, fy, fz = pts["x"], pts["y"], pts["z"]     # виды на поля, запись идёт в pts

    # предмет виден, только если он ближе уже зарегистрированного эха
    r_exist = np.full(base.shape, INF, np.float32)
    for e in range(n_echo):
        idx = base + e * n_rings
        ex, ey, ez = fx[idx], fy[idx], fz[idx]
        have = (ex != 0) | (ey != 0) | (ez != 0)
        r_e = np.where(have, np.sqrt(ex * ex + ey * ey + ez * ez), INF)
        r_exist = np.minimum(r_exist, r_e)

    closer = new_r < r_exist
    ring_i, col_i, base, new_r = ring_i[closer], col_i[closer], base[closer], new_r[closer]
    n_written = int(base.size)
    if n_written:
        dir_sel = layout.dirs[ring_i, col_i]
        nx = (dir_sel[:, 0] * new_r).astype(np.float32)
        ny = (dir_sel[:, 1] * new_r).astype(np.float32)
        nz = (dir_sel[:, 2] * new_r).astype(np.float32)
        # ламбертова интенсивность: ρ·cosθ/r², приведена к шкале прибора 0…255
        cos_inc = np.clip(np.abs(dir_sel[:, 1]), 0.05, 1.0)
        inten = np.clip(2.2e4 * obj.reflectivity * cos_inc / (new_r ** 2), 1, 255)
        for e in range(n_echo):
            idx = base + e * n_rings
            fx[idx] = nx
            fy[idx] = ny
            fz[idx] = nz
            pts["intensity"][idx] = inten.astype(np.float32)

    out = PointCloud2(stamp=pc.stamp, frame_id=pc.frame_id, height=pc.height,
                      width=pc.width, point_step=pc.point_step,
                      is_dense=pc.is_dense, points=pts)
    w, hgt = obj.size
    return out, dict(hit_rays=int(n_written), d=float(place.d), u=float(place.u),
                     name=obj.name, width=w, height=hgt,
                     reflectivity=obj.reflectivity,
                     # индексы лучей, в которые предмет реально записан: по ним
                     # разметка кандидата точная, а не «по дальности примерно»
                     rays=(ring_i, col_i))