# Action codec for the VAE: encode/decode pairs that convert the canonical 20-d action
# between latent and raw forms. Analogous to the VAE's pixel <-> latent codec for high-dim
# modalities, but for low-dim actions. Pairs:
#   normalize_actions <-> unnormalize_actions    (raw scale <-> normalized [-1, 1])
#   actions_to_relative <-> actions_to_absolute  (world frame <-> per-step relative frame)
# Borrowed and cleaned up from vidar/utils/action.py.

import torch
import pytorch3d.transforms

from anydata.geometry.camera_utils import invert_extrinsics
from custom.utils.io import read_pickle


# ---------------------------------------------------------------------------
# Normalize / unnormalize
# ---------------------------------------------------------------------------

def normalize_actions(action, normalizer):
    """Scale and shift action values to [-1, 1] range. Skips the first timestep (base frame)."""
    if normalizer is None:
        return action
    scale = normalizer['scale'].to(action.device).view(1, 1, -1)
    offset = normalizer['offset'].to(action.device).view(1, 1, -1)
    action = action.clone()
    action[:, 1:] = (action[:, 1:] * scale) + offset
    return action


def unnormalize_actions(action, normalizer):
    """Inverse of normalize_actions."""
    if normalizer is None:
        return action
    scale = normalizer['scale'].to(action.device).view(1, 1, -1)
    offset = normalizer['offset'].to(action.device).view(1, 1, -1)
    action = action.clone()
    action[:, 1:] = (action[:, 1:] - offset) / scale
    return action


def prepare_action_normalizer(path, relative):
    """Build scale/offset dict from normalizer stats file."""
    if relative:
        return _prepare_normalizer_relative(path)
    else:
        return _prepare_normalizer_absolute(path)


def _prepare_normalizer_absolute(path):
    """Normalizer for absolute 3D end-effector coordinates."""
    if not path:
        return None

    stats = read_pickle(path)['input_stats']
    input_min = torch.tensor(stats['min']).view(1, 1, -1)
    input_max = torch.tensor(stats['max']).view(1, 1, -1)

    range_eps, output_min, output_max = 1e-4, -1, 1
    input_range = input_max - input_min
    ignore_dim = input_range < range_eps
    input_range[ignore_dim] = output_max - output_min

    scale = (output_max - output_min) / input_range
    offset = output_min - scale * input_min
    offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim]

    return {'scale': scale, 'offset': offset}


def _prepare_normalizer_relative(path):
    """Normalizer for relative end-effector coordinates (LBM-style)."""
    if not path:
        return None

    from externals.lbm.lbm.diffusion_policy.common.normalize_util import (
        get_linear_normalizer_from_saved_stats,
    )
    normalizer = get_linear_normalizer_from_saved_stats(path)

    action_scale = normalizer['action'].params_dict['scale'].mean(0)
    action_offset = normalizer['action'].params_dict['offset'].mean(0)

    return {'scale': action_scale, 'offset': action_offset}


# ---------------------------------------------------------------------------
# Relative / absolute action conversion (bimanual robot format: 20-dim)
# ---------------------------------------------------------------------------

def _action_to_grip_all(action, arm_idx):
    """Extract gripper value for one arm from [B, N, 20] action tensor."""
    b, n = action.shape[:2]
    col = 18 if arm_idx == 0 else 19
    return action[:, :, col:col+1]


def _action_to_pose_all(action, arm_idx):
    """Convert [B, N, 20] action to [B, N, 4, 4] poses for one arm."""
    b, n, d = action.shape
    flat = action.view(b * n, d)

    if arm_idx == 0:
        xyz, rot6d = flat[:, 0:3], flat[:, 3:9]
    else:
        xyz, rot6d = flat[:, 9:12], flat[:, 12:18]

    rot = pytorch3d.transforms.rotation_6d_to_matrix(rot6d)

    pose = torch.eye(4, device=action.device, dtype=action.dtype).unsqueeze(0).expand(b * n, -1, -1).clone()
    pose[:, :3, :3] = rot
    pose[:, :3, 3] = xyz

    return invert_extrinsics(pose).view(b, n, 4, 4)


def _pose_to_action_all(pose0, pose1, grip0, grip1):
    """Convert per-arm poses [B, N, 4, 4] + grips [B, N, 1] back to [B, N, 20] action."""
    b, n = pose0.shape[:2]

    pose0 = invert_extrinsics(pose0.view(b * n, 4, 4))
    pose1 = invert_extrinsics(pose1.view(b * n, 4, 4))

    xyz0 = pose0[:, :3, 3]
    rot0 = pytorch3d.transforms.matrix_to_rotation_6d(pose0[:, :3, :3])
    xyz1 = pose1[:, :3, 3]
    rot1 = pytorch3d.transforms.matrix_to_rotation_6d(pose1[:, :3, :3])

    action = torch.zeros(b * n, 20, device=pose0.device, dtype=pose0.dtype)
    action[:, 0:3] = xyz0
    action[:, 3:9] = rot0
    action[:, 9:12] = xyz1
    action[:, 12:18] = rot1
    action[:, 18:19] = grip0.view(b * n, 1)
    action[:, 19:20] = grip1.view(b * n, 1)

    return action.view(b, n, 20)


def actions_to_relative(action):
    """Convert absolute bimanual actions [B, N, 20] to relative (w.r.t. first timestep).
    Returns (relative_action, [base0, base1]) for later recovery."""
    pose0 = _action_to_pose_all(action, 0)
    pose1 = _action_to_pose_all(action, 1)
    grip0 = _action_to_grip_all(action, 0)
    grip1 = _action_to_grip_all(action, 1)

    b, n = pose0.shape[:2]
    pose0 = pose0.view(b * n, 4, 4)
    pose1 = pose1.view(b * n, 4, 4)

    base0 = pose0[:1].clone()  # TODO seems wrong for B > 1 ?
    base1 = pose1[:1].clone()

    pose0[1:] = pose0[1:] @ invert_extrinsics(base0)
    pose1[1:] = pose1[1:] @ invert_extrinsics(base1)

    pose0 = pose0.view(b, n, 4, 4)
    pose1 = pose1.view(b, n, 4, 4)

    return _pose_to_action_all(pose0, pose1, grip0, grip1), [base0, base1]


def actions_to_absolute(action, bases):
    """Convert relative bimanual actions [B, N, 20] back to absolute using saved bases."""
    pose0 = _action_to_pose_all(action, 0)
    pose1 = _action_to_pose_all(action, 1)
    grip0 = _action_to_grip_all(action, 0)
    grip1 = _action_to_grip_all(action, 1)

    b, n = pose0.shape[:2]
    pose0 = pose0.view(b * n, 4, 4)
    pose1 = pose1.view(b * n, 4, 4)

    if bases is not None:
        pose0[1:] = pose0[1:] @ bases[0]
        pose1[1:] = pose1[1:] @ bases[1]

    pose0 = pose0.view(b, n, 4, 4)
    pose1 = pose1.view(b, n, 4, 4)

    return _pose_to_action_all(pose0, pose1, grip0, grip1)