"""3D rotation and oblique plane sampling for volumetric data visualization.

This module provides utilities for rotating the viewing plane in 3D space
using camera-centric (view-relative) rotations and sampling oblique planes from 
volumetric zarr arrays using trilinear interpolation.

Camera-centric rotations:
- The viewer looks down the -Z axis (toward the YX plane)
- rotation_roll: Spin around camera normal (Z-axis) - rotates image in-place
- rotation_tilt: Tilt around camera right (X-axis) - tilts to see different Z layers
"""

import numpy as np
from typing import Tuple
import logging

logger = logging.getLogger(__name__)


def rotation_matrix_z(angle_degrees: float) -> np.ndarray:
    """Generate rotation matrix around Z-axis (camera normal - in-plane spin).
    
    Args:
        angle_degrees: Rotation angle in degrees
        
    Returns:
        3x3 rotation matrix
    """
    angle_rad = np.radians(angle_degrees)
    cos_a = np.cos(angle_rad)
    sin_a = np.sin(angle_rad)
    
    return np.array([
        [ cos_a, -sin_a,  0],
        [ sin_a,  cos_a,  0],
        [     0,      0,  1]
    ])


def rotation_matrix_x(angle_degrees: float) -> np.ndarray:
    """Generate rotation matrix around X-axis (camera right - tilt for depth).
    
    Args:
        angle_degrees: Rotation angle in degrees
        
    Returns:
        3x3 rotation matrix
    """
    angle_rad = np.radians(angle_degrees)
    cos_a = np.cos(angle_rad)
    sin_a = np.sin(angle_rad)
    
    return np.array([
        [1,      0,       0],
        [0, cos_a, -sin_a],
        [0, sin_a,  cos_a]
    ])


def combined_rotation_matrix(rotation_roll: float, rotation_tilt: float) -> np.ndarray:
    """Combine camera-centric rotations into single matrix using intrinsic composition.
    
    Rotations are applied in the camera's rotating frame (intrinsic):
    - rotation_tilt: Tilt around camera right (X-axis) in degrees, applied first
    - rotation_roll: Spin around camera normal (Z-axis) in degrees, applied second
    
    Intrinsic composition: R_final = R_roll(angle_roll) @ R_tilt(angle_tilt)
    This means tilt happens first in body frame, then roll.
    
    Args:
        rotation_roll: Spin around camera normal (Z-axis) in degrees (0-180)
        rotation_tilt: Tilt around camera right (X-axis) in degrees (0-180)
        
    Returns:
        3x3 combined rotation matrix for intrinsic rotations
    """
    R_tilt = rotation_matrix_x(rotation_tilt)
    R_roll = rotation_matrix_z(rotation_roll)
    # Intrinsic composition: apply tilt first, then roll in the tilted frame
    return R_roll @ R_tilt


def pixel_to_world_coordinates(
    pixel_i: np.ndarray,
    pixel_j: np.ndarray,
    layer_height: int,
    layer_width: int,
    rotation_roll: float,
    rotation_tilt: float,
    center_z: int,
    z_dim_size: int = 1,
    pan_offset_x: int = 0,
    pan_offset_y: int = 0,
    depth_offset: int = 0,
    voxel_scale_xyz: Tuple[float, float, float] = (1.0, 1.0, 1.0)
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Transform 2D pixel coordinates to 3D world coordinates with camera-centric rotation.
    
    The viewer's camera looks down the -Z axis at the YX plane (centered at center_z).
    Rotations are applied in the camera's reference frame:
    - rotation_roll: Spin around camera normal (Z-axis) - rotates image in place
    - rotation_tilt: Tilt around camera right (X-axis) - tilts to reveal depth layers
    
    Voxel scaling is applied so that rotation happens in physical space, not voxel space.
    For example, if Z voxels are 5x larger than X/Y, the rotation will properly account for this.
    
    Args:
        pixel_i: 2D array of row indices (shape: fov_height, fov_width)
        pixel_j: 2D array of column indices (shape: fov_height, fov_width)
        layer_height: Height of the layer in voxels
        layer_width: Width of the layer in voxels
        rotation_roll: In-plane spin around camera normal (Z-axis) in degrees
        rotation_tilt: Tilt around camera right (X-axis) to see depth in degrees
        center_z: Z coordinate of plane center in voxels (rotation pivot point)
        z_dim_size: Number of Z slices in the volume
        pan_offset_x: Pan offset in camera-X direction (screen right) in pixels
        pan_offset_y: Pan offset in camera-Y direction (screen up) in pixels
        depth_offset: Additional depth movement along camera normal (voxel units)
        voxel_scale_xyz: Voxel spacing (scale_z, scale_y, scale_x) in physical units
        
    Returns:
        Tuple of (world_y, world_x, world_z) coordinate arrays in voxel indices
    """
    # Unpack voxel scaling
    scale_z, scale_y, scale_x = voxel_scale_xyz
    
    # Get FOV dimensions
    fov_height, fov_width = pixel_i.shape
    
    # Get rotation matrix - columns are camera basis vectors in layer frame
    R = combined_rotation_matrix(rotation_roll, rotation_tilt)
    # R[:, 0] = camera_x direction in layer frame
    # R[:, 1] = camera_y direction in layer frame  
    # R[:, 2] = camera_z direction (normal) in layer frame
    
    # Transform pan from camera frame to layer frame
    pan_camera = np.array([pan_offset_x, pan_offset_y, 0.0])
    pan_layer = pan_camera @ R.T
    
    # Base center point in layer coordinates (with camera-aware pan applied)
    center_y = layer_height / 2.0 + pan_layer[1]
    center_x = layer_width / 2.0 + pan_layer[0]
    center_z_world = center_z
    
    # Normalize pixel positions to [-1, 1] relative to FOV center
    # Y (rows): positive downward in images, but in 3D we want Y positive upward
    norm_y = -(pixel_i.astype(float) - fov_height / 2.0) / (fov_height / 2.0)
    norm_x = (pixel_j.astype(float) - fov_width / 2.0) / (fov_width / 2.0)
    
    # Create offsets from center in camera frame (pre-rotation)
    # Scale each dimension based on its voxel spacing AND FOV size
    # In physical space: offset_range = (fov_size / 2) * voxel_scale
    # NOTE: We do NOT include center_z here - it will be applied AFTER rotation
    # to preserve the rotation center point
    offset_x = norm_x * (fov_width / 2.0) * scale_x   # Apply X scaling
    offset_y = norm_y * (fov_height / 2.0) * scale_y  # Apply Y scaling
    offset_z = np.zeros_like(offset_x)                # No depth offset here - applied post-rotation
    
    # Stack into 3D offset vectors: shape (fov_height, fov_width, 3)
    offsets_camera = np.stack([offset_x, offset_y, offset_z], axis=-1)
    
    # Reshape for matrix multiplication: (N, 3)
    shape_orig = offsets_camera.shape
    offsets_flat = offsets_camera.reshape(-1, 3)
    
    # Apply rotation: (N, 3) @ (3, 3).T -> (N, 3)
    # This rotates the offset vectors from camera frame to layer frame
    offsets_rotated_flat = offsets_flat @ R.T
    
    # Reshape back: (fov_height, fov_width, 3)
    offsets_rotated = offsets_rotated_flat.reshape(shape_orig)
    
    # Convert back to voxel coordinates by dividing by voxel scale
    # (rotated offsets are in physical space, need to convert back to voxel indices)
    offsets_rotated[..., 0] = offsets_rotated[..., 0] / scale_x  # X
    offsets_rotated[..., 1] = offsets_rotated[..., 1] / scale_y  # Y
    offsets_rotated[..., 2] = offsets_rotated[..., 2] / scale_z  # Z
    
    # Depth movement: move camera along its rotated normal direction
    # R[:, 2] is the camera_z direction (normal) in layer frame
    # Multiply by depth and voxel scale to get physical displacement
    # Then divide by voxel scale to get back to voxel coordinates
    camera_z_direction = R[:, 2]  # The camera normal in layer frame
    depth_displacement_layer = depth_offset * camera_z_direction  # Physical displacement
    depth_displacement_voxels = depth_displacement_layer / np.array([scale_x, scale_y, scale_z])
    
    # Convert to world coordinates: center + rotated offsets + depth movement along camera normal
    # The center point stays fixed during rotation (it's the rotation pivot)
    # But depth moves the entire plane along the camera's viewing direction
    world_y = center_y + offsets_rotated[..., 1] + depth_displacement_voxels[1]  # Y component
    world_x = center_x + offsets_rotated[..., 0] + depth_displacement_voxels[0]  # X component
    world_z = center_z_world + offsets_rotated[..., 2] + depth_displacement_voxels[2]  # Z component
    
    # NOTE: We do NOT clamp Y and X here - let them go out of bounds
    # The rotation can legitimately sample beyond layer boundaries, and zarr's
    # fancy indexing will handle out-of-bounds gracefully. Clamping here causes
    # cropping artifacts as rotated corners all sample from the layer edge.
    # Only clamp Z since it's the discrete depth dimension.
    if z_dim_size > 1:
        world_z = np.clip(world_z, 0, z_dim_size - 1)
    else:
        world_z = np.clip(world_z, 0, np.inf)  # No Z clamping for 2D images
    
    return world_y, world_x, world_z


def trilinear_interpolate(
    v000: np.ndarray, v001: np.ndarray,
    v010: np.ndarray, v011: np.ndarray,
    v100: np.ndarray, v101: np.ndarray,
    v110: np.ndarray, v111: np.ndarray,
    frac_y: np.ndarray, frac_x: np.ndarray, frac_z: np.ndarray
) -> np.ndarray:
    """Perform trilinear interpolation from 8 corner values.
    
    Corner ordering:
    - First index (0/1) represents y coordinate (0=y_floor, 1=y_ceil)
    - Second index (0/1) represents x coordinate (0=x_floor, 1=x_ceil)
    - Third index (0/1) represents z coordinate (0=z_floor, 1=z_ceil)
    
    Args:
        v000-v111: 8 corner values (each shape: fov_height, fov_width)
        frac_y: Fractional part of Y coordinate (0-1)
        frac_x: Fractional part of X coordinate (0-1)
        frac_z: Fractional part of Z coordinate (0-1)
        
    Returns:
        Interpolated values (shape: fov_height, fov_width)
    """
    # Interpolate along Z axis first
    v_y0x0 = v000 * (1 - frac_z) + v001 * frac_z
    v_y0x1 = v010 * (1 - frac_z) + v011 * frac_z
    v_y1x0 = v100 * (1 - frac_z) + v101 * frac_z
    v_y1x1 = v110 * (1 - frac_z) + v111 * frac_z
    
    # Interpolate along X axis
    v_y0 = v_y0x0 * (1 - frac_x) + v_y0x1 * frac_x
    v_y1 = v_y1x0 * (1 - frac_x) + v_y1x1 * frac_x
    
    # Interpolate along Y axis
    result = v_y0 * (1 - frac_y) + v_y1 * frac_y
    
    return result
