"""equirect_projector.py — Equirectangular projection for 4× IMX219 fisheye cameras.

Equirectangular (lat-lon) projection:
  Output pixel (u, v) corresponds to:
    longitude lambda = (u / W - 0.5) * 2pi     [-pi, pi]
    latitude  phi    = (0.5 - v / H) * pi       [-pi/2, pi/2]
  3D unit vector: x=cos(phi)cos(lambda), y=cos(phi)sin(lambda), z=sin(phi)

For each camera (with known yaw, and mount position relative to base_link):
  - Rotate the 3D unit vector from world frame to camera frame (apply inverse yaw)
  - Project using fisheye equidistant model: r = f * theta, where theta = angle from
    optical axis
  - theta = atan2(sqrt(x^2+y^2), z)  [from optical axis = +z direction]
  - psi   = atan2(y, x)              [azimuth in camera image plane]
  - r     = f * theta  (equidistant fisheye)
  - u_cam = cx + r * cos(psi)
  - v_cam = cy + r * sin(psi)
  - Valid if theta < FoV_half (160deg/2 = 80deg)

Precomputed remap maps (computed once, applied with cv2.remap):
  For each camera: map_x, map_y of shape (H_out, W_out)
  Each (v, u) -> (u_cam, v_cam) lookup into source camera image
"""

import math
import numpy as np
import cv2
from typing import NamedTuple


class CameraConfig(NamedTuple):
    name: str
    yaw_deg: float    # camera facing direction (0=front, -90=right, 180=rear, 90=left)
    pitch_deg: float  # camera tilt (0=horizontal, positive=down)
    K: np.ndarray     # 3x3 camera matrix
    D: np.ndarray     # distortion coefficients (4x1 equidistant)
    image_size: tuple  # (width, height)
    fov_deg: float = 160.0  # total field of view


def _world_to_camera_frame(xyz_world: np.ndarray, yaw_deg: float,
                            pitch_deg: float) -> np.ndarray:
    """Rotate 3D world unit vectors to camera frame.

    World frame: x=forward, y=left, z=up (base_link ROS convention)
    Camera frame: z=optical axis (out of lens), x=right in image, y=down in image

    Args:
        xyz_world: [N, 3] unit vectors in world frame
        yaw_deg: camera facing direction (0=forward, CCW positive)
        pitch_deg: camera tilt down (positive = tilted down)

    Returns:
        [N, 3] unit vectors in camera frame (z=along optical axis)
    """
    yaw = math.radians(yaw_deg)
    pitch = math.radians(pitch_deg)

    # Rotation from world to camera:
    # 1. Rotate by -yaw around z (so camera forward = world x direction becomes z=0)
    # 2. Rotate by pitch around y (tilt)
    # 3. Apply axis swap: world(x=fwd,y=left,z=up) -> camera(z=fwd,x=right,y=down)
    #    i.e., camera_x = -world_y, camera_y = -world_z, camera_z = world_x

    cos_y, sin_y = math.cos(-yaw), math.sin(-yaw)
    cos_p, sin_p = math.cos(pitch), math.sin(pitch)

    # R_yaw (around z):
    Rz = np.array([[cos_y, -sin_y, 0],
                   [sin_y,  cos_y, 0],
                   [0,      0,     1]], dtype=np.float64)
    # R_pitch (around y, positive = tilt down = rotate -pitch around world y):
    Ry = np.array([[ cos_p, 0, sin_p],
                   [ 0,     1, 0    ],
                   [-sin_p, 0, cos_p]], dtype=np.float64)
    # Axis swap: world(x,y,z) -> camera(z_cam=x_world_rotated, x_cam=-y_world_rotated,
    #                                   y_cam=-z_world_rotated)
    R_swap = np.array([[0, -1,  0],
                       [0,  0, -1],
                       [1,  0,  0]], dtype=np.float64)

    R = R_swap @ Ry @ Rz
    return (R @ xyz_world.T).T  # [N, 3]


def build_equirect_remap(
    out_w: int,
    out_h: int,
    cam_config: CameraConfig,
) -> tuple:
    """Build cv2.remap maps for one camera's contribution to equirectangular output.

    Returns:
        map_x: float32 [out_h, out_w] -- source column lookup
        map_y: float32 [out_h, out_w] -- source row lookup
        mask:  uint8   [out_h, out_w] -- 255 where this camera contributes, 0 elsewhere
    """
    K = cam_config.K
    fx, fy = K[0, 0], K[1, 1]
    cx, cy = K[0, 2], K[1, 2]
    f_mean = (fx + fy) / 2.0   # equidistant: use mean focal length

    fov_half_rad = math.radians(cam_config.fov_deg / 2.0)

    # Build longitude/latitude grid
    u_idx = np.arange(out_w, dtype=np.float64)
    v_idx = np.arange(out_h, dtype=np.float64)
    lam = (u_idx / out_w - 0.5) * 2.0 * math.pi  # [-pi, pi]
    phi = (0.5 - v_idx / out_h) * math.pi          # [pi/2, -pi/2]

    lam_grid, phi_grid = np.meshgrid(lam, phi)     # both [H, W]
    lam_flat = lam_grid.ravel()
    phi_flat = phi_grid.ravel()

    # 3D unit vectors in world frame (x=forward, y=left, z=up)
    x_w = np.cos(phi_flat) * np.cos(lam_flat)
    y_w = np.cos(phi_flat) * np.sin(lam_flat)
    z_w = np.sin(phi_flat)
    xyz_world = np.stack([x_w, y_w, z_w], axis=1)  # [N, 3]

    # Rotate to camera frame
    xyz_cam = _world_to_camera_frame(xyz_world, cam_config.yaw_deg, cam_config.pitch_deg)

    z_cam = xyz_cam[:, 2]
    x_cam = xyz_cam[:, 0]
    y_cam = xyz_cam[:, 1]

    # Angle from optical axis (z direction)
    r_xy = np.sqrt(x_cam**2 + y_cam**2)
    theta = np.arctan2(r_xy, z_cam)   # [0, pi]
    psi   = np.arctan2(y_cam, x_cam)  # azimuth in camera plane

    # Equidistant projection: r_img = f * theta
    r_img = f_mean * theta

    # Image coordinates
    u_cam = cx + r_img * np.cos(psi)
    v_cam = cy + r_img * np.sin(psi)

    # Validity mask: behind camera OR outside FoV OR outside image bounds
    img_w, img_h = cam_config.image_size
    valid = (
        (z_cam > 0) &
        (theta < fov_half_rad) &
        (u_cam >= 0) & (u_cam < img_w - 1) &
        (v_cam >= 0) & (v_cam < img_h - 1)
    )

    map_x = u_cam.reshape(out_h, out_w).astype(np.float32)
    map_y = v_cam.reshape(out_h, out_w).astype(np.float32)
    mask  = (valid.reshape(out_h, out_w) * 255).astype(np.uint8)

    # Invalidate map coords outside valid region (set to -1 so remap produces black)
    map_x[mask == 0] = -1.0
    map_y[mask == 0] = -1.0

    return map_x, map_y, mask


def build_alpha_blend_weights(
    masks: list,
    blend_margin_px: int = 30,
) -> list:
    """Compute per-pixel alpha weights for smooth blending in overlap zones.

    Each mask is distance-transformed to get a weight proportional to distance
    from the mask boundary. Weights are normalized so they sum to 1 per pixel.

    Args:
        masks: list of uint8 [H, W] masks (255=valid, 0=invalid)
        blend_margin_px: blending zone width in pixels

    Returns:
        list of float32 [H, W] weight maps, summing to <=1 per pixel
    """
    weights = []
    for mask in masks:
        # Distance transform within the valid region
        dist = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
        # Clamp to blend_margin and normalize to [0, 1]
        w = np.clip(dist / blend_margin_px, 0.0, 1.0).astype(np.float32)
        weights.append(w)

    # Normalize: sum of weights per pixel
    weight_sum = sum(weights)
    weight_sum = np.maximum(weight_sum, 1e-6)
    return [w / weight_sum for w in weights]


class EquirectProjector:
    """Manages precomputed remap maps for 4-camera equirectangular stitching."""

    def __init__(self, out_w: int, out_h: int, cam_configs: list,
                 blend_margin_px: int = 30):
        self._out_w = out_w
        self._out_h = out_h
        self._configs = cam_configs

        print('Building equirect remap maps...')
        self._maps = []
        masks = []
        for cfg in cam_configs:
            mx, my, mask = build_equirect_remap(out_w, out_h, cfg)
            self._maps.append((mx, my))
            masks.append(mask)
            print(f'  {cfg.name}: coverage {mask.mean()/255*100:.1f}%')

        self._weights = build_alpha_blend_weights(masks, blend_margin_px)
        print('Remap maps ready.')

    def stitch(self, images: dict) -> np.ndarray:
        """Stitch 4 camera images into equirectangular panorama.

        Args:
            images: dict {cam_name: bgr_image}

        Returns:
            bgr panorama of shape (out_h, out_w, 3)
        """
        out = np.zeros((self._out_h, self._out_w, 3), dtype=np.float32)

        for i, cfg in enumerate(self._configs):
            img = images.get(cfg.name)
            if img is None:
                continue
            mx, my = self._maps[i]
            w = self._weights[i][:, :, np.newaxis]  # [H, W, 1]

            warped = cv2.remap(img.astype(np.float32), mx, my,
                               interpolation=cv2.INTER_LINEAR,
                               borderMode=cv2.BORDER_CONSTANT, borderValue=0)
            out += warped * w

        return np.clip(out, 0, 255).astype(np.uint8)
