import numpy as np
import cv2
import math

def equirect_to_perspective(img, yaw, pitch, fov=90, size=512):
    h, w = img.shape[:2]

    fov_rad = math.radians(fov)
    half = math.tan(fov_rad / 2)

    xs, ys = np.meshgrid(
        np.linspace(-half, half, size),
        np.linspace(half, -half, size)  # Reversed: sky at top
    )

    zs = np.ones_like(xs)

    dirs = np.stack([xs, -ys, zs], axis=-1)
    norm = np.linalg.norm(dirs, axis=-1, keepdims=True)
    dirs /= norm

    cy, sy = math.cos(yaw), math.sin(yaw)
    cp, sp = math.cos(pitch), math.sin(pitch)

    rot_y = np.array([[ cy, 0, sy],
                      [  0, 1,  0],
                      [-sy, 0, cy]])

    rot_x = np.array([[1,  0,   0],
                      [0, cp, -sp],
                      [0, sp,  cp]])

    dirs = dirs @ rot_y.T @ rot_x.T

    lon = np.arctan2(dirs[...,0], dirs[...,2])
    lat = np.arcsin(dirs[...,1])

    u = (lon / (2*np.pi) + 0.5) * w
    v = (0.5 - lat / np.pi) * h

    map_x = u.astype(np.float32)
    map_y = v.astype(np.float32)

    return cv2.remap(img, map_x, map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_WRAP)


def strip_to_equirect(strip_img, width, height, face_size, order=None, flip_sides=False):
    """Convert a 6-face strip back to equirectangular projection.
    
    Args:
        order: List of indices indicating which face is at each strip position.
               e.g., [0,1,2,3,4,5] means strip position 0 has front, 1 has left, etc.
               Face indices: 0=front, 1=left, 2=back, 3=right, 4=top, 5=bottom
        flip_sides: If True, flip the first 4 faces (horizontal views) vertically before processing
    """
    if order is None:
        order = [0, 1, 2, 3, 4, 5]
    
    # Create reverse mapping: for each face type, which strip position has it?
    face_to_strip = {face_idx: strip_pos for strip_pos, face_idx in enumerate(order)}
    
    faces = [
        strip_img[:, i*face_size:(i+1)*face_size]
        for i in range(len(order))
    ]
    
    # Flip horizontal faces if requested
    if flip_sides:
        for i, face_idx in enumerate(order):
            if face_idx < 4:  # front, left, back, right
                faces[i] = cv2.flip(faces[i], 0)

    # Vectorized approach for speed
    y_coords = np.arange(height)
    x_coords = np.arange(width)
    xx, yy = np.meshgrid(x_coords, y_coords)

    lat = (0.5 - yy / height) * np.pi
    lon = (xx / width - 0.5) * 2 * np.pi

    # Direction vector
    dx = np.cos(lat) * np.sin(lon)
    dy = np.sin(lat)
    dz = np.cos(lat) * np.cos(lon)

    adx, ady, adz = np.abs(dx), np.abs(dy), np.abs(dz)

    out = np.zeros((height, width, 3), dtype=np.uint8)

    # Front face (z+) - face index 0
    mask = (adz >= adx) & (adz >= ady) & (dz > 0)
    u = dx[mask] / dz[mask]
    v = dy[mask] / dz[mask]
    px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    out[mask] = faces[face_to_strip[0]][py, px]

    # Back face (z-) - face index 2
    mask = (adz >= adx) & (adz >= ady) & (dz <= 0)
    u = dx[mask] / dz[mask]
    v = -dy[mask] / dz[mask]
    px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    out[mask] = faces[face_to_strip[2]][py, px]

    # Right face (x+) - face index 3
    mask = (adx > adz) & (adx >= ady) & (dx > 0)
    u = -dz[mask] / dx[mask]
    v = dy[mask] / dx[mask]
    px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    out[mask] = faces[face_to_strip[3]][py, px]

    # Left face (x-) - face index 1
    mask = (adx > adz) & (adx >= ady) & (dx <= 0)
    u = -dz[mask] / dx[mask]
    v = -dy[mask] / dx[mask]
    px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    out[mask] = faces[face_to_strip[1]][py, px]

    # Top face (y+) - face index 4
    mask = (ady > adx) & (ady > adz) & (dy > 0)
    u = dx[mask] / dy[mask]
    v = -dz[mask] / dy[mask]  # Flipped to match equirect_to_perspective
    px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    out[mask] = faces[face_to_strip[4]][py, px]

    # Bottom face (y-) - face index 5
    mask = (ady > adx) & (ady > adz) & (dy <= 0)
    u = -dx[mask] / dy[mask]  # Flipped horizontal
    v = -dz[mask] / dy[mask]
    px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
    out[mask] = faces[face_to_strip[5]][py, px]

    return out