"""
GPU-accelerated data augmentation for training.

All augmentations operate on GPU tensors and are applied in the training
loop after the batch has been moved to the GPU.

data.augmentation keys:
- rotation_augmentation: rotation about the vertical (z) axis. One angle per
  batch, a multiple of angle_step_size, with probability prob_use_rotation;
  with per_sample: true every sample gets its own uniform angle in [0, 360)
  (probability prob_use_rotation per sample).
- reflection_augmentation: per sample with probability prob, the scene is
  mirrored in the vertical x = 0 plane (x -> -x). Mirrored rigid bodies under
  gravity, contact and isotropic friction follow the mirrored trajectories,
  so the mirrored sequence is again a valid one.
- rigid_noise: the input frames (all but the last frame of a sample) get a
  rigid random-walk perturbation per object (rotation about the object's
  centroid plus a translation), written to batch['obj_positions_input'];
  batch['obj_positions'] stays the ground truth the losses compare against.
  pos_std (m) and rot_std_deg are the standard deviations at the last input
  frame for step size 1, scaled by step_size ** step_size_exponent.
"""

import torch
import math
from typing import Dict
import logging

logger = logging.getLogger(__name__)


def rotvec_to_matrix(rotvec: torch.Tensor) -> torch.Tensor:
    """Rotation matrices (..., 3, 3) of rotation vectors (..., 3) (Rodrigues' formula)."""
    theta = rotvec.norm(dim=-1, keepdim=True)  # (..., 1)
    small = theta < 1e-8
    axis = rotvec / torch.where(small, torch.ones_like(theta), theta)
    x, y, z = axis.unbind(-1)
    zero = torch.zeros_like(x)
    K = torch.stack([zero, -z, y, z, zero, -x, -y, x, zero], dim=-1).view(*rotvec.shape[:-1], 3, 3)
    s = torch.sin(theta).unsqueeze(-1)
    c = torch.cos(theta).unsqueeze(-1)
    eye = torch.eye(3, dtype=rotvec.dtype, device=rotvec.device).expand(K.shape)
    R = eye + s * K + (1.0 - c) * (K @ K)
    return torch.where(small.unsqueeze(-1), eye, R)


class GPUAugmentation:
    """
    GPU-accelerated augmentation pipeline.

    Features:
    - Rotation augmentation (Z-axis only, no ground mesh), per batch or per sample
    - Reflection augmentation (x -> -x, per sample)
    - Rigid random-walk noise on the input frames
    - All operations on GPU tensors
    """

    def __init__(self, config: Dict, device: torch.device):
        """
        Initialize GPU augmentation.

        Args:
            config: Augmentation configuration
            device: Target GPU device
        """
        self.device = device
        self.config = config
        aug_cfg = config.get('data', {}).get('augmentation', {})

        # Rotation augmentation settings
        rotation_cfg = aug_cfg.get('rotation_augmentation', {})
        self.rotation_enabled = rotation_cfg.get('enable', False)
        self.rotation_angle_step = rotation_cfg.get('angle_step_size', 30)
        self.rotation_prob = rotation_cfg.get('prob_use_rotation', 0.3)
        self.rotation_per_sample = bool(rotation_cfg.get('per_sample', False))

        # Pre-compute rotation angles (only if rotation enabled and valid step size)
        if self.rotation_enabled and self.rotation_angle_step > 0:
            self.rotation_angles = torch.arange(
                self.rotation_angle_step, 360, self.rotation_angle_step,
                dtype=torch.float32, device=device
            )
        else:
            # Rotation disabled or invalid step size - create empty tensor
            self.rotation_angles = torch.tensor([], dtype=torch.float32, device=device)

        reflection_cfg = aug_cfg.get('reflection_augmentation', {}) or {}
        self.reflection_enabled = bool(reflection_cfg.get('enable', False))
        self.reflection_prob = float(reflection_cfg.get('prob', 0.5))

        noise_cfg = aug_cfg.get('rigid_noise', {}) or {}
        self.noise_enabled = bool(noise_cfg.get('enable', False))
        self.noise_pos_std = float(noise_cfg.get('pos_std', 0.0))
        self.noise_rot_std = math.radians(float(noise_cfg.get('rot_std_deg', 0.0)))
        self.noise_step_exponent = float(noise_cfg.get('step_size_exponent', 1.0))
        if self.noise_enabled and config.get('data', {}).get('center_scene_by_t0', False):
            # centring shifts the network inputs in place; with separate input frames the targets would not follow
            raise ValueError("data.augmentation.rigid_noise cannot be combined with data.center_scene_by_t0")

        # Counters for periodic logging.
        self.rotation_total_calls = 0
        self.rotation_applied_count = 0

        logger.info(f"\n{'='*80}")
        logger.info(f"[GPU Augmentation] Initialization")
        logger.info(f"{'='*80}")
        logger.info(f"Rotation Augmentation:")
        logger.info(f"  - Enabled: {self.rotation_enabled}")
        if self.rotation_enabled:
            if self.rotation_per_sample:
                logger.info(f"  - Per sample, uniform angle in [0, 360)")
            else:
                logger.info(f"  - Angle step size: {self.rotation_angle_step}°")
                logger.info(f"  - Rotation angles: {self.rotation_angles.cpu().tolist()}")
            logger.info(f"  - Probability: {self.rotation_prob:.2%}")
        logger.info(f"Reflection Augmentation (x -> -x): enabled={self.reflection_enabled}"
                    + (f", probability {self.reflection_prob:.2%} per sample" if self.reflection_enabled else ""))
        logger.info(f"Rigid input noise: enabled={self.noise_enabled}"
                    + (f", pos_std={self.noise_pos_std} m, rot_std={math.degrees(self.noise_rot_std):.3f} deg at the last "
                       f"input frame (step size 1), x step_size^{self.noise_step_exponent:g}" if self.noise_enabled else ""))
        logger.info(f"{'='*80}\n")

    @property
    def enabled(self) -> bool:
        return self.rotation_enabled or self.reflection_enabled or self.noise_enabled

    def apply_rotation_augmentation(
        self,
        batch: Dict[str, torch.Tensor],
        split: str = 'train'
    ) -> Dict[str, torch.Tensor]:
        """
        Apply rotation augmentation around Z-axis (GPU version, no ground mesh).

        Args:
            batch: Batch dict with GPU tensors
            split: 'train', 'val', or 'test' (only train gets augmented)

        Returns:
            Augmented batch dict
        """
        # Only apply to training split
        if split != 'train':
            batch['rotation_angle_deg'] = 0.0
            return batch

        self.rotation_total_calls += 1

        if not self.rotation_enabled:
            batch['rotation_angle_deg'] = 0.0
            return batch

        if self.rotation_per_sample:
            return self._rotate_per_sample(batch)

        if len(self.rotation_angles) == 0:
            raise ValueError(
                f"Rotation enabled but no valid angles available. "
                f"angle_step_size={self.rotation_angle_step} must be > 0"
            )

        # Decide whether to rotate (probabilistic)
        rand_val = torch.rand(1, device=self.device).item()
        if rand_val > self.rotation_prob:
            batch['rotation_angle_deg'] = 0.0
            return batch

        # Randomly select rotation angle
        angle_idx = torch.randint(0, len(self.rotation_angles), (1,), device=self.device)
        angle_deg = self.rotation_angles[angle_idx].item()

        if angle_deg == 0:
            batch['rotation_angle_deg'] = 0.0
            return batch

        self.rotation_applied_count += 1
        if self.rotation_applied_count <= 5 or self.rotation_total_calls % 100 == 0:
            apply_rate = self.rotation_applied_count / self.rotation_total_calls
            logger.info(
                f"[GPU Rotation Aug] Applying {angle_deg:.0f}° | "
                f"Batch {self.rotation_total_calls} | "
                f"Applied: {self.rotation_applied_count}/{self.rotation_total_calls} ({apply_rate:.2%})"
            )

        # Create rotation matrix (Z-axis)
        angle_rad = math.radians(angle_deg)
        cos_a = math.cos(angle_rad)
        sin_a = math.sin(angle_rad)

        R = torch.tensor([
            [cos_a, -sin_a, 0],
            [sin_a,  cos_a, 0],
            [0,      0,     1]
        ], dtype=torch.float32, device=self.device)

        # Rotate obj_positions (B, T, Nobj, Nv, 3); all objects rotate since
        # there is no ground mesh.
        obj_pos = batch['obj_positions']
        obj_pos_flat = obj_pos.reshape(-1, 3)
        obj_pos_rotated = torch.matmul(obj_pos_flat, R.T)
        batch['obj_positions'] = obj_pos_rotated.view(obj_pos.shape)

        # Rotate ref_obj_positions: (B, Nobj, Nv, 3)
        ref_pos = batch['ref_obj_positions']
        ref_pos_flat = ref_pos.reshape(-1, 3)
        ref_pos_rotated = torch.matmul(ref_pos_flat, R.T)
        batch['ref_obj_positions'] = ref_pos_rotated.view(ref_pos.shape)


        # Record the applied angle so ground-truth overlays can be rotated to match.
        batch['rotation_angle_deg'] = angle_deg

        return batch

    def _per_sample_linear(self, batch: Dict[str, torch.Tensor], M: torch.Tensor) -> None:
        """Apply the per-sample linear maps M (B, 3, 3) to obj_positions (B,T,Nobj,Nv,3) and ref_obj_positions."""
        # einsum may return non-contiguous strides; the model reshapes these tensors with .view()
        batch['obj_positions'] = torch.einsum('bij,btonj->btoni', M, batch['obj_positions']).contiguous()
        ref = batch['ref_obj_positions']
        if ref.dim() == 4:
            batch['ref_obj_positions'] = torch.einsum('bij,bonj->boni', M, ref).contiguous()
        else:
            raise ValueError(f"per-sample augmentation needs per-sample ref_obj_positions (B,Nobj,Nv,3); got {tuple(ref.shape)}")

    def _rotate_per_sample(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        B = batch['obj_positions'].shape[0]
        angle = torch.rand(B, device=self.device) * (2 * math.pi)
        angle = torch.where(torch.rand(B, device=self.device) < self.rotation_prob, angle, torch.zeros_like(angle))
        c, s = torch.cos(angle), torch.sin(angle)
        zero, one = torch.zeros_like(c), torch.ones_like(c)
        R = torch.stack([c, -s, zero, s, c, zero, zero, zero, one], dim=-1).view(B, 3, 3)
        self._per_sample_linear(batch, R)
        batch['rotation_angle_deg'] = 0.0  # no single angle: ground-truth overlays cannot be rotated to match
        batch['rotation_angles_deg_per_sample'] = torch.rad2deg(angle)
        return batch

    def apply_reflection(self, batch: Dict[str, torch.Tensor], split: str = 'train') -> Dict[str, torch.Tensor]:
        if split != 'train' or not self.reflection_enabled:
            return batch
        B = batch['obj_positions'].shape[0]
        flip = torch.rand(B, device=self.device) < self.reflection_prob
        sx = torch.where(flip, -torch.ones(B, device=self.device), torch.ones(B, device=self.device))
        M = torch.diag_embed(torch.stack([sx, torch.ones_like(sx), torch.ones_like(sx)], dim=-1))
        self._per_sample_linear(batch, M)
        batch['reflected'] = flip
        return batch

    def apply_rigid_noise(self, batch: Dict[str, torch.Tensor], split: str = 'train') -> Dict[str, torch.Tensor]:
        if split != 'train' or not self.noise_enabled:
            return batch
        pos = batch['obj_positions']  # (B,T,Nobj,Nv,3)
        B, T, Nobj, Nv, _ = pos.shape
        Tin = T - 1  # the network inputs are frames 0..T-2; frame T-1 is only a target
        if Tin < 1:
            return batch
        scale = torch.ones(B, device=pos.device)
        if 'skip_step_size' in batch:
            scale = batch['skip_step_size'].to(pos.device).float() ** self.noise_step_exponent
        scale = scale.view(B, 1, 1, 1)
        # random walk over the input frames: the last input frame has the configured standard deviation
        rotvec = (torch.randn(B, Tin, Nobj, 3, device=pos.device) * (self.noise_rot_std / math.sqrt(Tin))).cumsum(1) * scale
        shift = (torch.randn(B, Tin, Nobj, 3, device=pos.device) * (self.noise_pos_std / math.sqrt(Tin))).cumsum(1) * scale
        R = rotvec_to_matrix(rotvec)  # (B,Tin,Nobj,3,3)
        x = pos[:, :Tin]
        vmask = batch['valid_vtx_mask'][:, :Tin].to(torch.bool)  # (B,Tin,Nobj,Nv)
        w = vmask.unsqueeze(-1).to(x.dtype)
        centroid = (x * w).sum(3, keepdim=True) / w.sum(3, keepdim=True).clamp_min(1.0)  # (B,Tin,Nobj,1,3)
        xn = torch.einsum('btoij,btonj->btoni', R, x - centroid) + centroid + shift.unsqueeze(3)
        xn = torch.where(vmask.unsqueeze(-1), xn, x)  # padding stays as it is
        batch['obj_positions_input'] = torch.cat([xn, pos[:, Tin:]], dim=1).contiguous()
        return batch

    def __call__(
        self,
        batch: Dict[str, torch.Tensor],
        split: str = 'train'
    ) -> Dict[str, torch.Tensor]:
        """
        Apply all enabled augmentations.

        Args:
            batch: Batch dict with GPU tensors
            split: 'train', 'val', or 'test'

        Returns:
            Augmented batch dict
        """
        if self.rotation_enabled:
            batch = self.apply_rotation_augmentation(batch, split)
        else:
            # Rotation disabled: record a zero angle for consistency.
            batch['rotation_angle_deg'] = 0.0
        batch = self.apply_reflection(batch, split)
        # last: the noise perturbs the (rotated, mirrored) input frames
        batch = self.apply_rigid_noise(batch, split)

        return batch