"""Downstream perception protocols (§4.4, §4.5).

These are pluggable hooks: you supply a *frozen* perception model as a callable
(we do not ship CondLaneNet / BEVFormer weights), and these functions run it on
MapGS renders and compute the protocol metric. Use on real data with the real
detectors for paper numbers.

* lane_consistency_with_detector (§4.4): run a lane detector on extrapolated
  renders, report lane mIoU vs the projected map lanes.
* detection_drift_with_detector (§4.5): run a 3D detector on rendered surround
  video, report translation / scale / rotation drift of detected boxes vs GT.
"""

from __future__ import annotations

from typing import Callable, Dict, List

import torch

from mapgs.eval.metrics import lane_iou
from mapgs.hdmap.rasterize_map import project_polylines, render_lane_mask
from mapgs.losses import perturb_pose
from mapgs.model.dynamic import place_dynamic_gaussians


@torch.no_grad()
def lane_consistency_with_detector(
    evaluator,
    dataset,
    lane_detector: Callable[[torch.Tensor], torch.Tensor],  # rgb [3,H,W] in [0,1] -> lane mask [H,W]
    max_scenes: int = 30,
    shift: float = 3.0,
    frame: int = None,
) -> Dict[str, float]:
    """§4.4 with a real lane detector (e.g. frozen CondLaneNet)."""
    cfg = evaluator.cfg
    H, W = cfg.data.height, cfg.data.width
    frame = frame if frame is not None else cfg.data.num_frames // 2
    ious: List[float] = []
    n = min(len(dataset), max_scenes)
    for i in range(n):
        s = dataset[i]
        if len(s.lanes) == 0:
            continue
        g, dyn = evaluator._decode(s)
        g = g if dyn is None else place_dynamic_gaussians(
            g, dyn["box_centers"], dyn["box_rots"], dyn["canon_idx"], frame)
        base = (dataset.get_scene(i).cam2world[frame, 1] if hasattr(dataset, "get_scene")
                else s.ctx_c2w[1]).to(evaluator.device)
        K = (dataset.get_scene(i).K[1] if hasattr(dataset, "get_scene") else s.ctx_K[1]).to(evaluator.device)
        dev = perturb_pose(base, lateral=shift)
        out = evaluator.ras.render(g, K[None], dev[None], H, W)
        rgb = (evaluator.model.feature_to_rgb(out.color) if evaluator.model.uses_features
               else out.color[:, :3].clamp(0, 1))[0]
        pred_mask = lane_detector(rgb).to(evaluator.device)
        map_uv = project_polylines([l.to(evaluator.device) for l in s.lanes], K[None], dev[None], H, W)[0]
        ious.append(float(lane_iou(pred_mask, render_lane_mask(map_uv, H, W))))
    return {"lane_mIoU_detector": sum(ious) / max(len(ious), 1), "n_scenes": len(ious)}


@torch.no_grad()
def detection_drift_with_detector(
    evaluator,
    dataset,
    detector: Callable[[torch.Tensor], List[dict]],  # rgb [3,H,W] -> [{center:[3],size:[3],yaw:float}]
    max_scenes: int = 30,
    frame: int = None,
    match_radius: float = 4.0,
) -> Dict[str, float]:
    """§4.5: translation / scale / rotation drift of detections on MapGS renders
    vs GT boxes. Detections are matched to GT by nearest projected center."""
    cfg = evaluator.cfg
    H, W = cfg.data.height, cfg.data.width
    frame = frame if frame is not None else cfg.data.num_frames // 2
    t_err, s_err, r_err, matched = [], [], [], 0
    n = min(len(dataset), max_scenes)
    for i in range(n):
        s = dataset[i]
        if s.box_centers.shape[0] == 0:
            continue
        g, dyn = evaluator._decode(s)
        g = g if dyn is None else place_dynamic_gaussians(
            g, dyn["box_centers"], dyn["box_rots"], dyn["canon_idx"], frame)
        scene = dataset.get_scene(i) if hasattr(dataset, "get_scene") else None
        K = (scene.K[1] if scene else s.ctx_K[1]).to(evaluator.device)
        c2w = (scene.cam2world[frame, 1] if scene else s.ctx_c2w[1]).to(evaluator.device)
        out = evaluator.ras.render(g, K[None], c2w[None], H, W)
        rgb = (evaluator.model.feature_to_rgb(out.color) if evaluator.model.uses_features
               else out.color[:, :3].clamp(0, 1))[0]
        dets = detector(rgb)
        gt_centers = s.box_centers[:, frame].to(evaluator.device)  # [I,3]
        gt_size = s.box_size.to(evaluator.device)
        for d in dets:
            c = torch.as_tensor(d["center"], device=evaluator.device).float()
            j = (gt_centers - c).norm(dim=-1).argmin()
            if (gt_centers[j] - c).norm() > match_radius:
                continue
            matched += 1
            t_err.append(float((gt_centers[j] - c).norm()))
            s_err.append(float((gt_size[j] - torch.as_tensor(d["size"], device=evaluator.device)).abs().mean()))
            r_err.append(abs(float(d.get("yaw", 0.0))))
    return {
        "trans_drift_m": sum(t_err) / max(len(t_err), 1),
        "scale_drift_m": sum(s_err) / max(len(s_err), 1),
        "rot_drift_rad": sum(r_err) / max(len(r_err), 1),
        "matched": matched,
    }