"""Robustness / generalization studies (§4.6).

* pose noise : inject up to N degrees of rotation (+ optional translation) on the
  non-reference context cameras and measure the reconstruction drop. MAGT's
  map-anchored scaffold should degrade more gracefully than pixel-aligned methods.
* map noise  : reuse :class:`mapgs.data.MapNoise` to build an imperfect map
  (height noise / lane offsets / dropped segments) and re-evaluate — validates
  the *weak-supervision* claim (no hard map dependence).
* cross-dataset: train on one dataset, evaluate (zero-shot / finetuned) on another
  via the dataset adapters; map structure should transfer.
"""

from __future__ import annotations

import math
from dataclasses import replace
from typing import Dict, List, Optional

import torch

from mapgs.eval.metrics import psnr, ssim
from mapgs.train.render import render_scene_views, extract_scene_dynamic


def _random_small_rotation(max_deg: float, device) -> torch.Tensor:
    axis = torch.randn(3, device=device)
    axis = axis / axis.norm().clamp_min(1e-8)
    ang = math.radians(max_deg) * float(torch.empty(1).uniform_(-1, 1))
    x, y, z = axis
    c, s, C = math.cos(ang), math.sin(ang), 1 - math.cos(ang)
    return torch.tensor([
        [c + x * x * C, x * y * C - z * s, x * z * C + y * s],
        [y * x * C + z * s, c + y * y * C, y * z * C - x * s],
        [z * x * C - y * s, z * y * C + x * s, c + z * z * C],
    ], device=device)


def perturb_context_poses(sample, rot_deg: float, trans: float = 0.0, device="cuda"):
    """Return a copy of ``sample`` with noisy non-reference (idx>0) context poses."""
    c2w = sample.ctx_c2w.clone().to(device)
    for v in range(1, c2w.shape[0]):
        R = _random_small_rotation(rot_deg, device)
        c2w[v, :3, :3] = R @ c2w[v, :3, :3]
        if trans > 0:
            c2w[v, :3, 3] = c2w[v, :3, 3] + trans * torch.randn(3, device=device)
    return replace(sample, ctx_c2w=c2w.cpu())


@torch.no_grad()
def pose_noise_sweep(evaluator, dataset, rot_degs: Optional[List[float]] = None,
                     trans: float = 0.0, max_scenes: int = 20) -> Dict[float, Dict[str, float]]:
    rot_degs = rot_degs or [0.0, 2.0, 5.0, 10.0]
    cfg = evaluator.cfg
    H, W = cfg.data.height, cfg.data.width
    out: Dict[float, Dict[str, float]] = {}
    n = min(len(dataset), max_scenes)
    for rd in rot_degs:
        ps, ss = [], []
        for i in range(n):
            s = dataset[i]
            s = s if rd == 0.0 else perturb_context_poses(s, rd, trans, evaluator.device)
            g, dyn = evaluator._decode(s)
            rend = render_scene_views(evaluator.model, evaluator.ras, g, dyn,
                                      s.sup_K.to(evaluator.device), s.sup_c2w.to(evaluator.device),
                                      s.sup_frame.to(evaluator.device), H, W, evaluator.model.uses_features)
            gt = s.sup_images.to(evaluator.device)
            ps.append(float(psnr(rend["rgb"], gt)))
            ss.append(float(ssim(rend["rgb"], gt)))
        out[rd] = {"PSNR": sum(ps) / len(ps), "SSIM": sum(ss) / len(ss), "n": n}
    return out