import torch
import numpy as np
 
from .native import native_bilinear
from typing import Optional
from einops import rearrange
def linear_interp(v0, v1, d, L):
    return v0 * (1 - d) / L + v1 * d / L


def interp2d(q00, q10, q01, q11, dy, dx):
    f0 = linear_interp(q00, q01, dx, 1)
    f1 = linear_interp(q10, q11, dx, 1)
    return linear_interp(f0, f1, dy, 1)


def bilinear(
    img: torch.Tensor, grid: torch.Tensor, out: torch.Tensor
) -> torch.Tensor:
    b, _, h, w = img.shape

    min_grid = torch.floor(grid).type(torch.int64)
    max_grid = min_grid + 1
    d_grid = grid - min_grid

    min_grid[:, 0, :, :] %= h
    min_grid[:, 1, :, :] %= w
    max_grid[:, 0, :, :] %= h
    max_grid[:, 1, :, :] %= w

    # FIXME: anyway to do efficient batch?
    for i in range(b):
        dy = d_grid[i, 0, ...]
        dx = d_grid[i, 1, ...]
        min_ys = min_grid[i, 0, ...]
        min_xs = min_grid[i, 1, ...]
        max_ys = max_grid[i, 0, ...]
        max_xs = max_grid[i, 1, ...]

        min_ys %= h
        min_xs %= w

        p00 = img[i][:, min_ys, min_xs]
        p10 = img[i][:, max_ys, min_xs]
        p01 = img[i][:, min_ys, max_xs]
        p11 = img[i][:, min_ys, max_xs]

        out[i, ...] = interp2d(p00, p10, p01, p11, dy, dx)

    return out


def create_xyz_grid(
    w_face: int,
    batch: Optional[int] = None,
    dtype: torch.dtype = torch.float32,
    device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
    """xyz coordinates of the faces of the cube"""

    ratio = (w_face - 1) / w_face

    out = torch.zeros((w_face, w_face * 6, 3), dtype=dtype, device=device)
    rng = torch.linspace(
        -0.5 * ratio, 0.5 * ratio, w_face, dtype=dtype, device=device
    )

    # NOTE: https://github.com/pytorch/pytorch/issues/15301
    # Torch meshgrid behaves differently than numpy

    # Front face (x = 0.5)
    out[:, 0 * w_face : 1 * w_face, [2, 1]] = torch.stack(
        torch.meshgrid([-rng, rng], indexing="ij"), -1
    )
    out[:, 0 * w_face : 1 * w_face, 0] = 0.5

    # Right face (y = -0.5)
    out[:, 1 * w_face : 2 * w_face, [2, 0]] = torch.stack(
        torch.meshgrid([-rng, -rng], indexing="ij"), -1
    )
    out[:, 1 * w_face : 2 * w_face, 1] = 0.5

    # Back face (x = -0.5)
    out[:, 2 * w_face : 3 * w_face, [2, 1]] = torch.stack(
        torch.meshgrid([-rng, -rng], indexing="ij"), -1
    )
    out[:, 2 * w_face : 3 * w_face, 0] = -0.5

    # Left face (y = 0.5)
    out[:, 3 * w_face : 4 * w_face, [2, 0]] = torch.stack(
        torch.meshgrid([-rng, rng], indexing="ij"), -1
    )
    out[:, 3 * w_face : 4 * w_face, 1] = -0.5

    # Up face (z = 0.5)
    out[:, 4 * w_face : 5 * w_face, [0, 1]] = torch.stack(
        torch.meshgrid([rng, rng], indexing="ij"), -1
    )
    out[:, 4 * w_face : 5 * w_face, 2] = 0.5

    # Down face (z = -0.5)
    out[:, 5 * w_face : 6 * w_face, [0, 1]] = torch.stack(
        torch.meshgrid([-rng, rng], indexing="ij"), -1
    )
    out[:, 5 * w_face : 6 * w_face, 2] = -0.5

    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?
        out = torch.cat([out.unsqueeze(0)] * batch)
        # grid shape is (b, h, w, 3)

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

    if method == "robust":
        ui = (theta / (2 * np.pi) - 0.5) * w_equi - 0.5
        uj = (0.5 - phi / np.pi) * h_equi - 0.5
        ui %= w_equi
        uj %= h_equi
    elif method == "faster":
        ui = (theta / (2 * np.pi) - 0.5) * w_equi - 0.5
        uj = (0.5 - phi / np.pi) * h_equi - 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 torch_grid_sample(
    img: torch.Tensor,
    grid: torch.Tensor,
    out: Optional[torch.Tensor] = None,
    mode: str = "bilinear",
    backend: str = "native",
) -> torch.Tensor:
    """Torch grid sampling algorithm

    params:
    - img (torch.Tensor)
    - grid (torch.Tensor)
    - out (Optional[torch.Tensor]): defaults to None
    - mode (str): ('bilinear', 'bicubic', 'nearest')
    - backend (str): ('native', 'pure')

    return:
    - img (torch.Tensor)

    NOTE: for `backend`, `pure` is relatively efficient since grid doesn't need
    to be in the same device as the `img`. However, `native` is faster.

    NOTE: for `pure` backends, we need to pass reference to `out`.

    NOTE: for `native` backends, we should pass anything for `out`

    """

    if backend == "native":
        if out is not None:
            # NOTE: out is created
            warnings.warn(
                "don't need to pass preallocated `out` to `grid_sample`"
            )
        assert img.device == grid.device, (
            f"ERR: when using {backend}, the devices of `img` and `grid` need"
            "to be on the same device"
        )
        if mode == "nearest":
            out = native_nearest(img, grid)
        elif mode == "bilinear":
            out = native_bilinear(img, grid)
        elif mode == "bicubic":
            out = native_bicubic(img, grid)
        else:
            raise ValueError(f"ERR: {mode} is not supported")
    elif backend == "pure":
        # NOTE: img and grid can be on different devices, but grid should be on the cpu
        # FIXME: since bilinear implementation depends on `grid` being on device, I'm removing
        # this warning and will put `grid` onto the same device until a fix is found
        # if grid.device.type == "cuda":
        #     warnings.warn("input `grid` should be on the cpu, but got a cuda tensor")
        assert (
            out is not None
        ), "ERR: need to pass reference to `out`, but got None"
        assert img.device == grid.device, (
            f"ERR: when using {backend}, the devices of `img` and `grid` need"
            "to be on the same device"
        )
        if mode == "nearest":
            out = nearest(img, grid, out)
        elif mode == "bilinear":
            out = bilinear(img, grid, out)
        elif mode == "bicubic":
            out = bicubic(img, grid, out)
        else:
            raise ValueError(f"ERR: {mode} is not supported")
    else:
        raise ValueError(f"ERR: {backend} is not supported")

    return out



def get_cubemap(equi_tensor):
    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

    res = w//4
    out = torch.empty((batch_size, c, res, res * 6)).to(device=equi_tensor.device, dtype=equi_tensor.dtype)
    xyz = create_xyz_grid(res,batch_size,device=equi_tensor.device, dtype=equi_tensor.dtype)
    # TODO: add rotation
    grid = convert_grid(xyz=xyz, h_equi=h, w_equi=w, method="robust").to(device=equi_tensor.device)
    out = torch_grid_sample(img=equi_tensor, grid=grid, out=out, backend="pure")
    
    out = rearrange(out, 'b c h (n w) -> b n c h w',n=6)
    if resize_back:
        out = rearrange(out,'(b f) n c h w -> b n c f h w',f=f)
    return out