# Created by BVH, Dec 2025.

import numpy as np
import torch
from rich import print


# 3D pose translation -> 2D trajectory projection matrix (2x3).
# Maps (x, y, z) camera position to (lateral/rightward, forward).
# Uncomment the one matching the current camera convention:

# OLD1
# Vidar convention: X = right, Y = up, Z = backward.
# hence (x, -z) = (right, forward).
# CAM_TO_TRAJ_2D = [[1, 0, 0], [0, 0, -1]]

# OLD2
# AnyData convention: X = forward, Y = left, Z = up.
# hence (-y, x) = (right, forward).
# CAM_TO_TRAJ_2D = [[0, -1, 0], [1, 0, 0]]

# NEW
# AnyData convention: X = right, Y = down, Z = forward.
# hence (x, z) = (right, forward).
CAM_TO_TRAJ_2D = [[1, 0, 0], [0, 0, 1]]

# GAIA ego_pose convention: X = forward, Y = left, Z = up.
# hence (-y, x) = (right, forward).
EGO_TO_TRAJ_2D = [[0, -1, 0], [1, 0, 0]]


# ============================================================
# ROBOTICS

# See also:
# https://github.com/ToyotaResearchInstitute/lbm_eval/blob/main/TRAINING_DATA_FORMAT.md#the-actions-archive-actionsnpz


def perturb_action_basile1(action, cond_frames_raw, seed=None):
    '''
    Perturb robot action vectors by modifying xyz coords for both arms.
    RIGHT arm: indices 0, 1, 2 (xyz), 18 (gripper)
    LEFT arm: indices 9, 10, 11 (xyz), 19 (gripper)
    Both arms are perturbed in similar directions.
    :param action (B, T, 20) tensor of float.
    :param cond_frames_raw: int.
    :param seed: Optional int seed for reproducible perturbations (useful for per-sample variation).
    :return action (B, T, 20) tensor of float.
    '''
    B, T, D = action.shape
    assert D == 20, f'Expected 20-dim action vector, got {D}'
    device = action.device
    perturbed = action.clone()
    
    # Use local RNG if seed is provided (for per-sample variation)
    if seed is not None:
        rng = np.random.RandomState(seed)
    else:
        rng = np.random

    n_quadratic = 6
    num_frames = T - cond_frames_raw
    
    # Pre-compute the sum of easing factors for normalization
    # This lets us scale the step size so final offset equals target
    ease_sum = sum(min(dt / n_quadratic, 1.0) for dt in range(num_frames))
    if ease_sum < 1e-6:
        return perturbed  # No frames to perturb

    for b in range(B):
        # Sample max overall offset between 0.2 and 0.8 meters for each axis
        max_offset_x = 0.1 + 0.5 * rng.random()
        max_offset_y = 0.1 + 0.5 * rng.random()
        max_offset_z = 0.0 + 0.2 * rng.random()
        
        # Random direction for each axis (same for both arms)
        sign_x = rng.choice([-1, 1])
        sign_y = rng.choice([-1, 1])
        sign_z = rng.choice([-1, 1])
        
        # Compute per-frame step size so total offset at T equals max_offset
        step_x = max_offset_x / ease_sum
        step_y = max_offset_y / ease_sum
        step_z = max_offset_z / ease_sum

        # Quadratic easing for the first few frames, then linear, accumulates from zero
        total_offset_x = 0.0
        total_offset_y = 0.0
        total_offset_z = 0.0

        # Gradually offset xyz coordinates from original position over time
        for t in range(cond_frames_raw, T):
            dt = t - cond_frames_raw
            ease = min(dt / n_quadratic, 1.0)
            total_offset_x += ease * step_x
            total_offset_y += ease * step_y
            total_offset_z += ease * step_z
            
            # RIGHT arm: indices 0, 1, 2
            perturbed[b, t, 0] += sign_x * total_offset_x
            perturbed[b, t, 1] += sign_y * total_offset_y
            perturbed[b, t, 2] += sign_z * total_offset_z
            
            # LEFT arm: indices 9, 10, 11 (same direction as right arm)
            perturbed[b, t, 9] += sign_x * total_offset_x
            perturbed[b, t, 10] += sign_y * total_offset_y
            perturbed[b, t, 11] += sign_z * total_offset_z

    return perturbed


def perturb_action_basile2(action, cond_frames_raw, seed=None):
    '''
    Perturb robot action vectors by modifying xyz coords for both arms.
    Makes a parabolic arc: offset grows to max at midpoint, then returns to 0 at end.
    Uses sin(pi * progress) for smooth start/end and peak at midpoint.
    RIGHT arm: indices 0, 1, 2 (xyz), 18 (gripper)
    LEFT arm: indices 9, 10, 11 (xyz), 19 (gripper)
    Each arm is perturbed in independently sampled directions.
    :param action (B, T, 20) tensor of float.
    :param cond_frames_raw: int.
    :param seed: Optional int seed for reproducible perturbations (useful for per-sample variation).
    :return action (B, T, 20) tensor of float.
    '''
    B, T, D = action.shape
    assert D == 20, f'Expected 20-dim action vector, got {D}'
    device = action.device
    perturbed = action.clone()
    
    # Use local RNG if seed is provided (for per-sample variation)
    if seed is not None:
        rng = np.random.RandomState(seed)
    else:
        rng = np.random

    num_frames = T - cond_frames_raw
    if num_frames <= 1:
        return perturbed  # No frames to perturb

    for b in range(B):
        # Sample max overall offset between 0.1 and 0.5 meters for each axis
        max_offset_x = 0.0 + 0.5 * rng.random()
        max_offset_y = 0.0 + 0.5 * rng.random()
        max_offset_z = 0.0 + 0.2 * rng.random()
        
        # Random direction for each axis (independent for each arm)
        sign_x_r = rng.choice([-1, 1])
        sign_y_r = rng.choice([-1, 1])
        sign_z_r = rng.choice([-1, 1])
        sign_x_l = rng.choice([-1, 1])
        sign_y_l = rng.choice([-1, 1])
        sign_z_l = rng.choice([-1, 1])

        # Apply parabolic arc offset using raised cosine (Hann window)
        # Maps progress [0,1] to [-pi, pi] with (sin+1)/2, giving derivative=0 at both ends
        # This gives: 0 at start, max at midpoint, 0 at end, with smooth entry/exit
        for t in range(cond_frames_raw, T):
            dt = t - cond_frames_raw
            progress = dt / (num_frames - 1)  # 0 to 1
            arc = (1.0 - np.cos(2.0 * np.pi * progress)) / 2.0  # raised cosine
            
            offset_x = max_offset_x * arc
            offset_y = max_offset_y * arc
            offset_z = max_offset_z * arc
            
            # RIGHT arm: indices 0, 1, 2
            perturbed[b, t, 0] += sign_x_r * offset_x
            perturbed[b, t, 1] += sign_y_r * offset_y
            perturbed[b, t, 2] += sign_z_r * offset_z
            
            # LEFT arm: indices 9, 10, 11 (independent direction)
            perturbed[b, t, 9] += sign_x_l * offset_x
            perturbed[b, t, 10] += sign_y_l * offset_y
            perturbed[b, t, 11] += sign_z_l * offset_z

        # With 50% probability, invert gripper status for a random subclip
        if rng.random() < 0.5:
            t_start = rng.randint(cond_frames_raw, T - 5)
            t_end = rng.randint(t_start + 1, T + 1)
            
            # Invert gripper: 0.1 - g (indices 18 for right, 19 for left)
            # Value range seems to be [0.0, 0.1] meters
            perturbed[b, t_start:t_end, 18] = 0.1 - perturbed[b, t_start:t_end, 18]
            perturbed[b, t_start:t_end, 19] = 0.1 - perturbed[b, t_start:t_end, 19]

    return perturbed


# ============================================================
# DRIVING


def compute_trajectory_from_extrinsics(anydata_dict, avail_times, ego_traj_proxy):
    '''
    Compute 2D ego trajectory from a weighted combination of camera origins.
        NOTE(bvh): we assume a reasonable reference camera has already been applied to cams.
    :param avail_times: list of int.
    :param ego_traj_proxy: list of (cam_idx, weight) tuples. Weights normalize internally,
        so e.g. [(front, 0.5), (back, 0.5)] uses the midpoint of front and back centers
        as a proxy for vehicle center. See EGO_TRAJ_PROXY in unified_anydrive.py.
    :return (B, T, 2) float tensor of (lateral, forward) offsets relative to the first frame,
        where typically positive = rightward and positive = forward.
    '''
    cams = anydata_dict['cams']  # dict mapping (time, cam_idx) to Camera objects.
    w_sum = sum(w for _, w in ego_traj_proxy)
    centers_per_t = [
        sum(w * cams[(t, c)].get_center() for c, w in ego_traj_proxy) / w_sum
        for t in avail_times
    ]  # list of (B, 3) tensors
    all_trans = torch.stack(centers_per_t, dim=1)  # (B, T, 3)

    # Convert absolute to relative (first frame = origin)
    rel_trans = all_trans - all_trans[:, 0:1, :]  # (B, T, 3)

    # Project 3D positions to 2D trajectory via CAM_TO_TRAJ_2D
    proj = torch.tensor(CAM_TO_TRAJ_2D, dtype=rel_trans.dtype, device=rel_trans.device)
    trajectory = rel_trans @ proj.T  # (B, T, 2): (lateral, forward)
    
    # NOTE(bvh): NaN values sometimes appear in trajectories, likely due to invalid camera
    # extrinsics/intrinsics in a subset of episodes in some datasets => can crash visuals code.
    if trajectory.isnan().any():
        trajectory = trajectory.nan_to_num(0.0)

    return trajectory


def compute_trajectory_from_ego_pose(anydata_dict, avail_times):
    '''
    Compute 2D ego trajectory from 6DoF ego vehicle pose.
        NOTE(bvh): we use the first ego pose of each sample as the frame of reference.
    :param avail_times: list of int.
    :return (B, T, 2) float tensor of (lateral, forward) offsets relative to the first frame,
        where typically positive = rightward and positive = forward.
    '''
    ego_pose = anydata_dict['ego_pose']  # dict mapping time -> (B, 4, 4) "Tcw" tensor (cam dim stripped in post_process_sample).
    ego_stack = torch.stack([ego_pose[t] for t in avail_times], dim=1)  # (B, T, 4, 4).
    
    # NOTE(bvh): do everything in high precision to minimize numerical errors
    ego_stack = ego_stack.to(torch.float64)
    ref_pose = ego_stack[:, 0:1, :, :]  # (B, 1, 4, 4).
    ref_pose_inv = torch.linalg.inv(ref_pose)  # (B, 1, 4, 4).
    rel_pose = ref_pose_inv @ ego_stack  # (B, T, 4, 4).
    rel_pose = rel_pose.to(torch.float32)  # (B, T, 4, 4).
    
    all_trans = rel_pose[:, :, 0:3, 3]  # (B, T, 3) = (x, y, z).
    rel_trans = all_trans  # first one is already (0, 0, 0) by definition of rel_pose
    # assert (rel_trans[:, 0, :].abs() < 1e-2).all(), f'rel_trans invalid: {rel_trans}'
    
    # Project 3D positions to 2D trajectory via EGO_TO_TRAJ_2D
    proj = torch.tensor(EGO_TO_TRAJ_2D, dtype=rel_trans.dtype, device=rel_trans.device)
    trajectory = rel_trans @ proj.T  # (B, T, 2) = (lateral, forward).
    
    return trajectory


def perturb_traj_basile1(trajectory, cond_frames_raw, seed=None):
    '''
    :param trajectory (B, T, 2) tensor of float.
    :param cond_frames_raw: int.
    :param seed: Optional int seed for reproducible perturbations.
    :return trajectory (B, T, 2) tensor of float.
    '''
    # Use a time for loop for simplicity as requested
    B, T, _ = trajectory.shape
    device = trajectory.device
    perturbed = trajectory.clone()

    # Use local RNG if seed is provided (for per-sample variation)
    rng = np.random.RandomState(seed) if seed is not None else np.random

    for b in range(B):
        # Sample magnitude randomly in [0.05, 2.0) meters per second (= 10 frames) in either direction
        mag_x = 0.05 + 1.95 * rng.random() if rng.random() < 0.8 else 0.0
        mag_z = 0.05 + 1.95 * rng.random() if rng.random() < 0.4 else 0.0
        sign_x = rng.choice([-1, 1])
        sign_z = rng.choice([-1, 1])
        
        # Quadratic easing for the first few frames, then linear, accumulates from zero
        n_quadratic = 8
        total_offset_x = 0.0
        total_offset_z = 0.0

        # Increasingly X coordinate from original position over time
        for t in range(cond_frames_raw, T):
            dt = t - cond_frames_raw
            offset_step_x = min(dt / n_quadratic, 1.0) * mag_x / 10.0
            offset_step_z = min(dt / n_quadratic, 1.0) * mag_z / 10.0
            total_offset_x += offset_step_x
            total_offset_z += offset_step_z
            perturbed[b, t, 0] += sign_x * total_offset_x
            perturbed[b, t, 1] += sign_z * total_offset_z

    return perturbed


def perturb_traj_basile2(trajectory, cond_frames_raw, seed=None):
    '''
    :param trajectory (B, T, 2) tensor of float.
    :param cond_frames_raw: int.
    :param seed: Optional int seed for reproducible perturbations.
    :return trajectory (B, T, 2) tensor of float.
    '''
    # Use a time for loop for simplicity as requested
    B, T, _ = trajectory.shape
    device = trajectory.device
    perturbed = trajectory.clone()

    # Use local RNG if seed is provided (for per-sample variation)
    rng = np.random.RandomState(seed) if seed is not None else np.random

    for b in range(B):
        # Sample step magnitude randomly in [0.1, 3.0) meters per second (= 10 frames)
        # going in either direction
        mag_x = 0.1 + 2.9 * rng.random() if rng.random() < 0.8 else 0.0
        mag_z = 0.1 + 2.9 * rng.random() if rng.random() < 0.4 else 0.0
        sign_x = rng.choice([-1, 1])
        sign_z = rng.choice([-1, 1])

        # Further scale following square root of distance traveled
        # such that driving at ~10 meters per second is equivalent to basile1
        forward_range = (trajectory[b, :, 1].max() - trajectory[b, :, 1].min()).item()  # meters
        duration = (T / 10.0)  # seconds
        avg_speed = forward_range / duration  # meters per second
        speed_factor = np.sqrt(avg_speed / 10.0)
        mag_x *= speed_factor
        mag_z *= speed_factor
        
        # Quadratic easing for the first few frames, then linear,
        # step size accumulates from zero
        n_quadratic = 8
        total_offset_x = 0.0
        total_offset_z = 0.0

        # Increasingly X coordinate from original position over time
        for t in range(cond_frames_raw, T):
            dt = t - cond_frames_raw
            offset_step_x = min(dt / n_quadratic, 1.0) * mag_x / 10.0
            offset_step_z = min(dt / n_quadratic, 1.0) * mag_z / 10.0
            total_offset_x += offset_step_x
            total_offset_z += offset_step_z
            perturbed[b, t, 0] += sign_x * total_offset_x
            perturbed[b, t, 1] += sign_z * total_offset_z

    return perturbed


def perturb_traj_basile3(trajectory, cond_frames_raw, seed=None):
    '''
    :param trajectory (B, T, 2) tensor of float.
    :param cond_frames_raw: int.
    :param seed: Optional int seed for reproducible perturbations.
    :return trajectory (B, T, 2) tensor of float.
    '''
    # Use a time for loop for simplicity as requested
    B, T, _ = trajectory.shape
    device = trajectory.device
    perturbed = trajectory.clone()

    # Use local RNG if seed is provided (for per-sample variation)
    rng = np.random.RandomState(seed) if seed is not None else np.random

    for b in range(B):
        # Sample step magnitude randomly in [0.1, 6.0) meters per second (= 10 frames)
        # going in either direction
        if rng.random() < 0.2:
            mag_x = 3.0 + 3.0 * rng.random()
        elif rng.random() < 0.8:
            mag_x = 0.1 + 2.9 * rng.random()
        else:
            mag_x = 0.0
        
        if rng.random() < 0.2:
            mag_z = 3.0 + 5.0 * rng.random()
        elif rng.random() < 0.4:
            mag_z = 0.1 + 2.9 * rng.random()
        else:
            mag_z = 0.0
        
        sign_x = rng.choice([-1, 1])
        sign_z = rng.choice([-1, 1])

        # Further scale following square root of distance traveled
        # such that driving at ~10 meters per second is equivalent to basile1
        forward_range = (trajectory[b, :, 1].max() - trajectory[b, :, 1].min()).item()  # meters
        duration = (T / 10.0)  # seconds
        avg_speed = forward_range / duration  # meters per second
        speed_factor = np.sqrt(avg_speed / 10.0)
        mag_x *= speed_factor
        mag_z *= speed_factor
        
        # Quadratic easing for the first few frames, then linear,
        # step size accumulates from zero
        n_quadratic = 8
        total_offset_x = 0.0
        total_offset_z = 0.0

        # Increasingly X coordinate from original position over time
        for t in range(cond_frames_raw, T):
            dt = t - cond_frames_raw
            offset_step_x = min(dt / n_quadratic, 1.0) * mag_x / 10.0
            offset_step_z = min(dt / n_quadratic, 1.0) * mag_z / 10.0
            total_offset_x += offset_step_x
            total_offset_z += offset_step_z
            perturbed[b, t, 0] += sign_x * total_offset_x
            perturbed[b, t, 1] += sign_z * total_offset_z

    return perturbed