#!/usr/bin/env python3
"""
target_math.py — 纯 numpy 的"像素 -> 三维目标"定位核心（不依赖 ROS，方便离线测试）。

输入：一帧点云（雷达系）+ 查询像素
输出：目标在相机系/雷达系下的坐标、距离、置信度、是否有多个深度层等
"""
import numpy as np

import calib


def _fit_plane_ray(P, ray, max_residual=0.05):
    """PCA 拟合平面，再让射线与平面求交。返回 (点, 残差)。"""
    if P.shape[0] < 6:
        return None, None
    c = P.mean(axis=0)
    _, _, vt = np.linalg.svd(P - c)
    n = vt[2]
    denom = float(n @ ray)
    if abs(denom) < 1e-6:
        return None, None
    lam = float(n @ c) / denom
    if lam <= 0:
        return None, None
    p = lam * ray
    residual = float(np.sqrt(np.mean(((P - c) @ n) ** 2)))
    if residual > max_residual:
        return None, residual
    return p, residual


def locate(xyz, query_u, query_v, conf=None, radius=10.0, depth_band=0.6,
           fit_plane=True, max_residual=0.05, min_conf=None, min_z=0.05):
    """在点云里找查询像素方向上最近的可见点。

    参数
        xyz        : (N,3) 雷达系点云
        query_u/v  : 查询像素
        conf       : (N,) 置信度，可选
        radius     : 候选像素半径
    返回 dict（未找到时 found=False）
    """
    xyz = np.asarray(xyz, dtype=np.float64)
    good = np.isfinite(xyz).all(axis=1) & (np.abs(xyz).sum(axis=1) > 1e-6)
    if conf is not None and min_conf is not None:
        good &= np.asarray(conf) >= min_conf
    idx = np.where(good)[0]
    pts = xyz[good]

    p_cam = calib.lidar_to_cam(pts)
    u, v, ok = calib.project_cam_to_pixel(p_cam)

    dpx = np.hypot(u - query_u, v - query_v)
    cand = ok & (dpx <= radius) & (p_cam[:, 2] > min_z)

    rng = np.linalg.norm(p_cam, axis=1)
    res = {
        'found': False, 'n_valid': int(pts.shape[0]), 'n_cand': int(cand.sum()),
        'radius': float(radius), 'query_u': float(query_u), 'query_v': float(query_v),
    }
    if not cand.any():
        return res

    cand_idx = np.where(cand)[0]
    order = cand_idx[np.argsort(rng[cand_idx])]
    sel = int(order[0])

    ray = calib.pixel_to_ray(query_u, query_v)
    p_near = p_cam[sel]

    # 深度分层：候选距离跨度大 => 可能擦过边缘
    r_sorted = np.sort(rng[cand_idx])
    multi = bool(len(r_sorted) > 3 and (r_sorted[-1] - r_sorted[0]) > 2.0)

    p_final = p_near
    plane_used = False
    residual = None
    if fit_plane:
        band = p_cam[cand_idx][rng[cand_idx] <= rng[sel] + depth_band]
        p_fit, residual = _fit_plane_ray(band, ray, max_residual)
        if p_fit is not None:
            p_final = p_fit
            plane_used = True

    res.update({
        'found': True,
        'p_cam': p_final,
        'range': float(np.linalg.norm(p_final)),
        'z': float(p_final[2]),
        'p_lidar': calib.cam_to_lidar(p_final)[0],
        'hit_u': float(u[sel]),
        'hit_v': float(v[sel]),
        'px_err': float(dpx[sel]),
        'lidar_index': int(idx[sel]),
        'plane_used': plane_used,
        'plane_residual': residual,
        'multi_layer': multi,
        'cand_ranges': rng[cand_idx],
        'confidence': float(conf[idx[sel]]) if conf is not None else None,
    })
    return res


class TargetIndex:
    """一帧点云的「像素格子 -> 点索引」倒排表。

    为什么需要：
        locate() 每次点击都会把 N 个点全量投影一遍。N≈5 万时一次几毫秒到
        几十毫秒，点一次无所谓；但同一帧反复点、或在板子上跑时，这部分是
        重复劳动。这里在【每帧】只做一次全量投影 + 分桶，之后每次点击只
        检查查询像素周围几个格子里的点，平面拟合也只在局部候选里做。

    用法：
        idx = TargetIndex(xyz, conf=conf)      # 每帧建一次
        res = idx.locate(u, v, radius=10)      # 每次点击，返回结构和 locate 一样

    注意：
        - 桶按点的像素坐标分，格子越大，查询时扫到的点越多；默认 32 px。
        - 只索引落在图像内的点；图像外的点本来也不会被点到。
    """

    def __init__(self, xyz, conf=None, cell=32, min_z=0.05):
        xyz = np.asarray(xyz, dtype=np.float64)
        good = np.isfinite(xyz).all(axis=1) & (np.abs(xyz).sum(axis=1) > 1e-6)
        self.idx = np.where(good)[0]
        self.xyz = xyz[good]
        self.conf = None if conf is None else np.asarray(conf)[good]

        self.p_cam = calib.lidar_to_cam(self.xyz)
        self.u, self.v, ok = calib.project_cam_to_pixel(self.p_cam)
        self.rng = np.linalg.norm(self.p_cam, axis=1)
        self.ok = ok & (self.p_cam[:, 2] > float(min_z))

        self.cell = max(4, int(cell))
        self.nx = int(np.ceil(calib.IMAGE_WIDTH / float(self.cell)))
        self.ny = int(np.ceil(calib.IMAGE_HEIGHT / float(self.cell)))

        sel = np.where(self.ok)[0]
        if sel.size:
            cx = np.clip((self.u[sel] // self.cell).astype(np.int64),
                         0, self.nx - 1)
            cy = np.clip((self.v[sel] // self.cell).astype(np.int64),
                         0, self.ny - 1)
            keys = cy * self.nx + cx
            order = np.argsort(keys, kind='stable')
            keys = keys[order]
            sel = sel[order]
            uniq, starts = np.unique(keys, return_index=True)
            ends = np.append(starts[1:], keys.size)
            self.buckets = {int(k): sel[s:e]
                            for k, s, e in zip(uniq, starts, ends)}
        else:
            self.buckets = {}

    def _gather(self, query_u, query_v, radius):
        c = self.cell
        x0 = max(0, int(np.floor((query_u - radius) / c)))
        x1 = min(self.nx - 1, int(np.floor((query_u + radius) / c)))
        y0 = max(0, int(np.floor((query_v - radius) / c)))
        y1 = min(self.ny - 1, int(np.floor((query_v + radius) / c)))

        parts = []
        for cy in range(y0, y1 + 1):
            base = cy * self.nx
            for cx in range(x0, x1 + 1):
                arr = self.buckets.get(base + cx)
                if arr is not None and arr.size:
                    parts.append(arr)
        if not parts:
            return np.empty(0, dtype=np.int64)
        return np.concatenate(parts)

    def locate(self, query_u, query_v, radius=10.0, depth_band=0.6,
               fit_plane=True, max_residual=0.05, min_conf=None):
        """局部查找，返回与 locate() 相同的字段。"""
        query_u = float(query_u)
        query_v = float(query_v)

        res = {
            'found': False, 'n_valid': int(self.idx.size),
            'n_cand': 0, 'radius': float(radius),
            'query_u': query_u, 'query_v': query_v,
            'indexed': True,
        }

        cand = self._gather(query_u, query_v, radius)
        if cand.size == 0:
            return res

        dpx = np.hypot(self.u[cand] - query_u, self.v[cand] - query_v)
        keep = dpx <= radius
        if min_conf is not None and self.conf is not None:
            keep &= self.conf[cand] >= min_conf
        cand = cand[keep]
        dpx = dpx[keep]
        res['n_cand'] = int(cand.size)
        if cand.size == 0:
            return res

        # 距离最近的点
        order = np.argsort(self.rng[cand])
        cand = cand[order]
        dpx = dpx[order]
        sel = int(cand[0])

        ray = calib.pixel_to_ray(query_u, query_v)
        r_sorted = np.sort(self.rng[cand])
        multi = bool(r_sorted.size > 3 and (r_sorted[-1] - r_sorted[0]) > 2.0)

        p_final = self.p_cam[sel]
        plane_used = False
        residual = None
        if fit_plane:
            band = self.p_cam[cand[self.rng[cand] <= self.rng[sel] + depth_band]]
            p_fit, residual = _fit_plane_ray(band, ray, max_residual)
            if p_fit is not None:
                p_final = p_fit
                plane_used = True

        res.update({
            'found': True,
            'p_cam': p_final,
            'p_lidar': calib.cam_to_lidar(p_final)[0],
            'hit_u': float(self.u[sel]),
            'hit_v': float(self.v[sel]),
            'px_err': float(dpx[0]),
            'lidar_index': int(self.idx[sel]),
            'plane_used': plane_used,
            'plane_residual': residual,
            'multi_layer': multi,
            'confidence': None if self.conf is None else float(self.conf[sel]),
        })
        return res