"""Per-epoch augmentation + subsampling for cached scene samples.

These transforms apply to the dict returned by `HoHoCachedDataset` *after* it
has been loaded from a semantic cache shard. They mutate `scene_xyz`, the per-point
parallel arrays, and `verts_gt[:, :3]` in place-equivalent ways (returning
new tensors) so that:

  - Random 3D rotation: rotate scene + GT vertices by a random 3D rotation matrix.
  - Random horizontal flip: optional X→−X mirror, applied to scene + GT.
  - Random XYZ jitter: small Gaussian noise on `scene_xyz` only.
  - Optional random subsample to `n_pts`: only active when the cache has more
    points than the configured model input. At eval time we instead take a
    deterministic prefix slice.

All operations are torch-native to keep work on whichever device the
DataLoader returns; in practice they run on CPU before the batch ships to
the GPU.
"""
from __future__ import annotations

from typing import Optional

import torch


# Per-point arrays that must be co-permuted when subsampling/dropping points.
POINT_KEYS = (
    "scene_xyz",
    "scene_type_ids",
    "scene_gestalt_ids",
    "scene_gestalt_id2",
    "scene_gestalt_w1",
    "scene_ade_ids",
    "scene_geom_conf",
    "scene_sem_conf",
    "scene_rgb",
)


def _yaw_rotation_matrix(theta: torch.Tensor) -> torch.Tensor:
    """3x3 rotation around Y (Y-up). theta: scalar tensor in radians."""
    c = torch.cos(theta)
    s = torch.sin(theta)
    R = torch.eye(3, dtype=torch.float32, device=theta.device)
    R[0, 0] = c
    R[0, 2] = s
    R[2, 0] = -s
    R[2, 2] = c
    return R


def random_yaw_rotation(sample: dict, generator: Optional[torch.Generator] = None) -> dict:
    """Rotate scene_xyz and verts_gt[:, :3] by the same random Y-axis angle."""
    theta = torch.rand((), generator=generator) * (2.0 * torch.pi)
    R = _yaw_rotation_matrix(theta)

    sample["scene_xyz"] = sample["scene_xyz"] @ R.T
    if "verts_gt" in sample:
        v = sample["verts_gt"]
        # v[:, :3] only — keep the validity flag at column 3 untouched.
        v_xyz = v[:, :3] @ R.T
        sample["verts_gt"] = torch.cat([v_xyz, v[:, 3:4]], dim=-1)
    return sample


def random_horizontal_flip(sample: dict, p: float = 0.5,
                            generator: Optional[torch.Generator] = None) -> dict:
    """With probability `p`, mirror the scene along X (X → −X)."""
    if torch.rand((), generator=generator).item() >= p:
        return sample
    sample["scene_xyz"] = sample["scene_xyz"] * torch.tensor([-1.0, 1.0, 1.0])
    if "verts_gt" in sample:
        v = sample["verts_gt"]
        v_xyz = v[:, :3] * torch.tensor([-1.0, 1.0, 1.0])
        sample["verts_gt"] = torch.cat([v_xyz, v[:, 3:4]], dim=-1)
    return sample


def random_xyz_jitter(sample: dict, sigma: float = 0.01,
                       generator: Optional[torch.Generator] = None) -> dict:
    """Add small Gaussian noise to scene_xyz only (verts_gt stays clean)."""
    if sigma <= 0.0:
        return sample
    noise = torch.randn(sample["scene_xyz"].shape, generator=generator) * sigma
    sample["scene_xyz"] = sample["scene_xyz"] + noise
    return sample


def random_point_subsample(sample: dict, n_pts: int,
                            generator: Optional[torch.Generator] = None) -> dict:
    """Randomly keep `n_pts` of the cached points (training augmentation)."""
    n_cache = sample["scene_xyz"].shape[0]
    if n_cache == n_pts:
        return sample
    if n_cache < n_pts:
        # Pad with random repeats — should be rare since cache is built oversized.
        extra = torch.randint(0, n_cache, (n_pts - n_cache,), generator=generator)
        idx = torch.cat([torch.arange(n_cache), extra])
    else:
        idx = torch.randperm(n_cache, generator=generator)[:n_pts]
    for k in POINT_KEYS:
        if k in sample:
            sample[k] = sample[k][idx]
    return sample


def deterministic_point_slice(sample: dict, n_pts: int) -> dict:
    """Take the first `n_pts` cached points (eval/inference use)."""
    n_cache = sample["scene_xyz"].shape[0]
    if n_cache == n_pts:
        return sample
    if n_cache < n_pts:
        # Cache shorter than expected — pad by repeating the last point.
        last = n_cache - 1
        pad = torch.full((n_pts - n_cache,), last, dtype=torch.long)
        idx = torch.cat([torch.arange(n_cache), pad])
    else:
        idx = torch.arange(n_pts)
    for k in POINT_KEYS:
        if k in sample:
            sample[k] = sample[k][idx]
    return sample


def apply_train_transforms(
    sample: dict,
    n_pts: int,
    yaw: bool = True,
    flip: bool = True,
    jitter_sigma: float = 0.01,
    generator: Optional[torch.Generator] = None,
) -> dict:
    """Full training-time augmentation pipeline. Order: subsample → flip → yaw → jitter."""
    sample = random_point_subsample(sample, n_pts, generator=generator)
    if flip:
        sample = random_horizontal_flip(sample, generator=generator)
    if yaw:
        sample = random_yaw_rotation(sample, generator=generator)
    if jitter_sigma > 0.0:
        sample = random_xyz_jitter(sample, sigma=jitter_sigma, generator=generator)
    return sample


def apply_eval_transforms(sample: dict, n_pts: int) -> dict:
    """Eval/inference-time pipeline: deterministic slice, no rotation, no jitter."""
    return deterministic_point_slice(sample, n_pts)