import torch
import numpy as np
import random
from .native import native_bilinear
from typing import Optional, List, Dict
from einops import rearrange
from .get_cubemap import torch_grid_sample

def get_random_rotate():
    roll_ang = 0
    pitch_ang = random.randint(-10, 10)
    yaw_ang = random.randint(0, 360)
    rot = [{ # TODO: define random rotate
        "roll": np.deg2rad(roll_ang),  #
        "pitch": np.deg2rad(pitch_ang),  # vertical
        "yaw": np.deg2rad(yaw_ang),  # horizontal
    }]
    return rot


def matmul(m: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
    M = torch.matmul(R[:, None, None, ...], m)
    M = M.squeeze(-1)

    return M

def create_normalized_grid(
    height: int,
    width: int,
    batch: Optional[int] = None,
    dtype: torch.dtype = torch.float32,
    device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
    """Create coordinate grid with height and width

    NOTE: primarly used for equi2equi

    params:
    - height (int)
    - width (int)
    - batch (Optional[int])
    - dtype (torch.dtype)

    return:
    - grid (torch.Tensor)

    """

    # NOTE: RuntimeError: "linspace_cpu" not implemented for Half
    if device.type == "cpu":
        assert dtype in (torch.float32, torch.float64), (
            f"ERR: {dtype} is not supported by {device.type}\n"
            "If device is `cpu`, use float32 or float64"
        )

    xs = torch.linspace(0, width - 1, width, dtype=dtype, device=device)
    ys = torch.linspace(0, height - 1, height, dtype=dtype, device=device)
    theta = xs * 2 * np.pi / width - np.pi
    phi = ys * np.pi / height - np.pi / 2
    phi, theta = torch.meshgrid([phi, theta], indexing="ij")
    a = torch.stack((theta, phi), dim=-1)
    norm_A = 1
    x = norm_A * torch.cos(a[..., 1]) * torch.cos(a[..., 0])
    y = norm_A * torch.cos(a[..., 1]) * torch.sin(a[..., 0])
    z = norm_A * torch.sin(a[..., 1])
    grid = torch.stack((x, y, z), dim=-1)

    # batched (stacked copies)
    if batch is not None:
        assert isinstance(
            batch, int
        ), f"ERR: batch needs to be integer: batch={batch}"
        assert (
            batch > 0
        ), f"ERR: batch size needs to be larger than 0: batch={batch}"
        # FIXME: faster way of copying?
        grid = torch.cat([grid.unsqueeze(0)] * batch)
        # grid shape is (b, h, w, 3)

    return grid

def convert_grid(
    M: torch.Tensor, h_equi: int, w_equi: int, method: str = "robust"
) -> torch.Tensor:
    # convert to rotation
    phi = torch.asin(M[..., 2] / torch.norm(M, dim=-1))
    theta = torch.atan2(M[..., 1], M[..., 0])

    if method == "robust":
        ui = (theta - np.pi) * w_equi / (2 * np.pi)
        uj = (phi - np.pi / 2) * h_equi / np.pi
        ui += 0.5
        uj += 0.5
        ui %= w_equi
        uj %= h_equi
    elif method == "faster":
        ui = (theta - np.pi) * w_equi / (2 * np.pi)
        uj = (phi - np.pi / 2) * h_equi / np.pi
        ui += 0.5
        uj += 0.5
        ui = torch.where(ui < 0, ui + w_equi, ui)
        ui = torch.where(ui >= w_equi, ui - w_equi, ui)
        uj = torch.where(uj < 0, uj + h_equi, uj)
        uj = torch.where(uj >= h_equi, uj - h_equi, uj)
    else:
        raise ValueError(f"ERR: {method} is not supported")

    # stack the pixel maps into a grid
    grid = torch.stack((uj, ui), dim=-3)

    return grid

def create_rotation_matrix(
    roll: float,
    pitch: float,
    yaw: float,
    z_down: bool = True,
    dtype: torch.dtype = torch.float32,
    device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
    """Create Rotation Matrix

    params:
    - roll, pitch, yaw (float): in radians
    - z_down (bool): flips pitch and yaw directions
    - dtype (torch.dtype): data types

    returns:
    - R (torch.Tensor): 3x3 rotation matrix
    """

    # calculate rotation about the x-axis
    R_x = torch.tensor(
        [
            [1.0, 0.0, 0.0],
            [0.0, np.cos(roll), -np.sin(roll)],
            [0.0, np.sin(roll), np.cos(roll)],
        ],
        dtype=dtype,
    )
    # calculate rotation about the y-axis
    if not z_down:
        pitch = -pitch
    R_y = torch.tensor(
        [
            [np.cos(pitch), 0.0, np.sin(pitch)],
            [0.0, 1.0, 0.0],
            [-np.sin(pitch), 0.0, np.cos(pitch)],
        ],
        dtype=dtype,
    )
    # calculate rotation about the z-axis
    if not z_down:
        yaw = -yaw
    R_z = torch.tensor(
        [
            [np.cos(yaw), -np.sin(yaw), 0.0],
            [np.sin(yaw), np.cos(yaw), 0.0],
            [0.0, 0.0, 1.0],
        ],
        dtype=dtype,
    )
    R = R_z @ R_y @ R_x
    return R.to(device)

def create_rotation_matrices(
    rots: List[Dict[str, float]],
    z_down: bool = True,
    dtype: torch.dtype = torch.float32,
    device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
    """Create rotation matrices from batch of rotations

    This methods creates a bx3x3 np.ndarray where `b` referes to the number
    of rotations (rots) given in the input
    """

    R = torch.empty((len(rots), 3, 3), dtype=dtype, device=device)
    for i, rot in enumerate(rots):
        # FIXME: maybe default to `create_rotation_matrix_at_once`?
        # NOTE: at_once is faster with cpu, while slower on GPU
        R[i, ...] = create_rotation_matrix(
            **rot, z_down=z_down, dtype=dtype, device=device
        )

    return R

def rotate_equi(equi_tensor, rots, z_down=True):
    if equi_tensor.ndim == 4: # image
        b,c,h,w = equi_tensor.shape
        batch_size = b
        resize_back = False
    else:
        b,c,f,h,w = equi_tensor.shape
        equi_tensor = rearrange(equi_tensor, 'b c f h w -> (b f) c h w')
        batch_size = b*f
        resize_back = True  
    
    out = torch.empty((batch_size, c, h, w)).to(device=equi_tensor.device, dtype=equi_tensor.dtype)

    m = create_normalized_grid(
        height=h, width=w, batch=batch_size, device=equi_tensor.device, dtype=equi_tensor.dtype
    )
    m = m.unsqueeze(-1)

    # create batched rotation matrices
    R = create_rotation_matrices(
        rots=rots, z_down=z_down, device=equi_tensor.device, dtype=equi_tensor.dtype
    )

    # rotate the grid
    M = matmul(m, R)

    grid = convert_grid(M=M, h_equi=h, w_equi=w, method="robust")
    # grid sample
    out = torch_grid_sample(
        img=equi_tensor,
        grid=grid,
        out=out,  # FIXME: is this necessary?
        backend="pure",
    )

    if resize_back:
        out = rearrange(out,'(b f) c h w -> b c f h w',f=f)
    return out