"""
Pose and action utility functions adapted from WorldMem/framear.
"""
import numpy as np
import torch
from einops import rearrange


ACTION_KEYS = [
    "inventory",    # 0
    "ESC",          # 1
    "hotbar.1",     # 2
    "hotbar.2",     # 3
    "hotbar.3",     # 4
    "hotbar.4",     # 5
    "hotbar.5",     # 6
    "hotbar.6",     # 7
    "hotbar.7",     # 8
    "hotbar.8",     # 9
    "hotbar.9",     # 10
    "forward",      # 11
    "back",         # 12
    "left",         # 13
    "right",        # 14
    "cameraY",      # 15
    "cameraX",      # 16
    "jump",         # 17
    "sneak",        # 18
    "sprint",       # 19
    "swapHands",    # 20
    "attack",       # 21
    "use",          # 22
    "pickItem",     # 23
    "drop",         # 24
]


def convert_action_space(actions: np.ndarray) -> torch.Tensor:
    """Convert raw WorldMem action array to 25-dim one-hot vector.

    Args:
        actions: np.ndarray of shape (T, 8+), raw action indices from NPZ.

    Returns:
        Tensor of shape (T, 25) with float values.
    """
    vec_25 = torch.zeros(len(actions), len(ACTION_KEYS))
    vec_25[actions[:, 0] == 1, 11] = 1   # forward
    vec_25[actions[:, 0] == 2, 12] = 1   # backward
    vec_25[actions[:, 4] == 11, 16] = -1  # camera X left
    vec_25[actions[:, 4] == 13, 16] = 1   # camera X right
    vec_25[actions[:, 3] == 11, 15] = -1  # camera Y up
    vec_25[actions[:, 3] == 13, 15] = 1   # camera Y down
    vec_25[actions[:, 5] == 6, 24] = 1   # place box
    vec_25[actions[:, 5] == 1, 24] = 1   # place box
    vec_25[actions[:, 1] == 1, 13] = 1   # left
    vec_25[actions[:, 1] == 2, 14] = 1   # right
    vec_25[actions[:, 7] == 1, 2] = 1    # hotbar 1
    return vec_25


def euler_to_rotation_matrix(pitch: torch.Tensor, yaw: torch.Tensor) -> torch.Tensor:
    """Convert pitch and yaw angles (radians) to 3x3 rotation matrix.

    Args:
        pitch: Tensor of pitch angles, shape (...,).
        yaw:   Tensor of yaw angles, shape (...,).

    Returns:
        Rotation matrix of shape (..., 3, 3).
    """
    cos_pitch, sin_pitch = torch.cos(pitch), torch.sin(pitch)
    cos_yaw, sin_yaw = torch.cos(yaw), torch.sin(yaw)

    R_pitch = torch.stack([
        torch.ones_like(pitch), torch.zeros_like(pitch), torch.zeros_like(pitch),
        torch.zeros_like(pitch), cos_pitch, -sin_pitch,
        torch.zeros_like(pitch), sin_pitch, cos_pitch
    ], dim=-1).reshape(-1, 3, 3)

    R_yaw = torch.stack([
        cos_yaw, torch.zeros_like(yaw), sin_yaw,
        torch.zeros_like(yaw), torch.ones_like(yaw), torch.zeros_like(yaw),
        -sin_yaw, torch.zeros_like(yaw), cos_yaw
    ], dim=-1).reshape(-1, 3, 3)

    return torch.matmul(R_yaw, R_pitch)


def euler_to_camera_to_world_matrix(pose: torch.Tensor) -> torch.Tensor:
    """Convert (x, y, z, pitch, yaw) pose vector to 4x4 camera-to-world matrix.

    Args:
        pose: Tensor of shape (5,), (f, 5), or (f, b, 5).

    Returns:
        Tensor of shape (4,4), (f, 4, 4), or (f, b, 4, 4).
    """
    origin_dim = pose.ndim
    if origin_dim == 1:
        pose = pose.unsqueeze(0).unsqueeze(0)
    elif origin_dim == 2:
        pose = pose.unsqueeze(0)

    x, y, z, pitch, yaw = (
        pose[..., 0], pose[..., 1], pose[..., 2],
        pose[..., 3], pose[..., 4]
    )
    pitch = torch.deg2rad(pitch)
    yaw = torch.deg2rad(yaw)

    R = euler_to_rotation_matrix(pitch, yaw)

    eye = torch.eye(4, dtype=pose.dtype, device=pose.device)
    c2w = eye.repeat(R.shape[0], 1, 1)
    c2w[:, :3, :3] = R
    c2w[:, :3, 3] = torch.stack([x.reshape(-1), y.reshape(-1), z.reshape(-1)], dim=-1)

    if origin_dim == 3:
        return c2w.view(pose.shape[0], pose.shape[1], 4, 4)
    elif origin_dim == 2:
        return c2w.view(pose.shape[0], 4, 4)
    else:
        return c2w.squeeze(0).squeeze(0)


def c2w_to_w2c(c2w: torch.Tensor) -> torch.Tensor:
    """Invert camera-to-world to world-to-camera matrix.

    Args:
        c2w: Tensor of shape (..., 4, 4).

    Returns:
        w2c: Tensor of shape (..., 4, 4).
    """
    R = c2w[..., :3, :3]
    t = c2w[..., :3, 3]
    Rt = R.transpose(-1, -2)
    t_new = -Rt @ t.unsqueeze(-1)
    w2c = torch.zeros_like(c2w)
    w2c[..., :3, :3] = Rt
    w2c[..., :3, 3] = t_new.squeeze(-1)
    w2c[..., 3, 3] = 1.0
    return w2c


def extract_pose_action_from_npz(
    data_path: str,
    start_frame: int,
    num_frames: int,
    dtype=torch.float32
) -> tuple:
    """Load action and pose from WorldMem NPZ file.

    Args:
        data_path:   Path to .npz file containing 'actions' and 'poses'.
        start_frame: Starting frame index.
        num_frames:  Number of frames to load.
        dtype:       Output tensor dtype.

    Returns:
        actions: Tensor (num_frames, 25) - encoded action vectors.
        poses:   Tensor (num_frames, 5)  - raw pose vectors (x,y,z,pitch,yaw).
    """
    data = np.load(data_path)
    actions_raw = data["actions"]          # (N, 8+)
    poses_pool = data["poses"].copy()      # (N, 5)

    # Fix potentially corrupted height in first frame
    poses_pool[0, 1] = poses_pool[1, 1]

    # Pad poses if shorter than actions
    if len(poses_pool) < len(actions_raw):
        poses_pool = np.pad(poses_pool, ((1, 0), (0, 0)))

    actions_raw = actions_raw[start_frame: start_frame + num_frames]
    poses = poses_pool[start_frame: start_frame + num_frames]

    # Normalize pose: translate relative to first frame, negate yaw
    poses[:, :3] -= poses[:1, :3]
    poses[:, -1] = -poses[:, -1]
    poses[:, 3:] %= 360

    actions = convert_action_space(actions_raw)  # (T, 25)
    return actions.to(dtype), torch.tensor(poses, dtype=dtype)


def poses_to_viewmats(poses: torch.Tensor) -> torch.Tensor:
    """Convert (B, T, 5) pose vectors to (B, T, 4, 4) world-to-camera matrices.

    Args:
        poses: Tensor of shape (B, T, 5).

    Returns:
        viewmats: Tensor of shape (B, T, 4, 4).
    """
    B, T, _ = poses.shape
    # (B, T, 5) -> (T, B, 5) for euler_to_camera_to_world_matrix
    pose_tbd = rearrange(poses, "b t d -> t b d")
    c2w = euler_to_camera_to_world_matrix(pose_tbd)  # (T, B, 4, 4)
    w2c = c2w_to_w2c(c2w)                            # (T, B, 4, 4)
    viewmats = rearrange(w2c, "t b p q -> b t p q")  # (B, T, 4, 4)
    return viewmats