from typing import Optional
from jaxtyping import Shaped, Bool, Float
from torch import Tensor

from itertools import product

import torch

try:
    import triton
    import triton.language as tl
    from .casting import scalar_cast as cast

    TRITON_ENABLED = True
except ImportError:
    from torchlinops.utils import fake_triton as triton, fake_tl as tl

    TRITON_ENABLED = False

from .nblocks import get_nblocks

__all__ = ["unfold"]

# Maximum size of the cuda grid per dimension
# 3 dimensions maximum
MAX_GRID_PER_DIM = 1024

MAX_TENSOR_POW_OF_2 = 20


def unfold(
    x: Shaped[Tensor, "..."],
    block_size: tuple,
    stride: Optional[tuple] = None,
    mask: Optional[Bool[Tensor, "..."]] = None,
    output: Optional[Tensor] = None,
) -> Tensor:
    """Wrapper that dispatches complex and real tensors
    Also precomputes some shapes

    Parameters
    ----------
    x : Tensor
        Shape [B..., *im_size]

    Returns
    -------
    Tensor: Shape [B..., *blocks, *block_size]
        If mask is not None, block_size will be an int equal to the number of True elements in the mask
        Otherwise it will be the full block shape.


    """
    x_flat, shapes, is_complex = _prep_unfold(x, block_size, stride, mask)
    if is_complex:
        x_flat = torch.view_as_real(x_flat)
        x_flat = torch.flatten(x_flat, -2, -1)  # Flatten real/imag into last dim
    y_flat = _unfold(x_flat, output=output, **shapes)
    y = y_flat.reshape(
        *shapes["batch_shape"],
        *shapes["nblocks"],
        *shapes["block_size"],
    )
    if is_complex:
        y = y.reshape(*y.shape[:-1], y.shape[-1] // 2, 2)
        y = torch.view_as_complex(y)
    if mask is not None:
        y = y[..., mask]
    return y


def _unfold(
    x: Shaped[Tensor, "B ..."],
    block_size: tuple[int, ...],
    stride: tuple[int, ...],
    ndim: int,
    im_size: tuple[int, ...],
    nblocks: tuple[int, ...],
    nbatch: int,
    output: Optional[Tensor] = None,
    **kwargs,
) -> Shaped[Tensor, "B ..."]:
    """Implementation of unfold"""
    # Check dtype if output buffer is provided
    if output is not None:
        if not output.dtype == x.dtype:
            raise ValueError(
                f"Output and input dtypes must match but got output {output.dtype} != input {x.dtype}"
            )
    if tuple(x.shape[-ndim:]) != tuple(im_size):
        raise RuntimeError(
            f"Unfold expected input with full size {im_size} but got {x.shape}"
        )
    if x.is_cuda and ndim in UNFOLD.keys():
        x = x.contiguous()  # Ensure contiguity
        with torch.cuda.device(x.device):
            if output is None:
                # Allocate output
                y = torch.zeros(
                    nbatch,
                    *nblocks,
                    *block_size,
                    device=x.device,
                    dtype=x.dtype,
                )
            else:
                # Use existing buffer
                y = output.reshape(nbatch, *nblocks, *block_size).zero_()
            grid = _get_grid(ndim, nbatch, nblocks)
            BLOCK_SIZE = tuple(
                min(
                    triton.next_power_of_2(blk_size),
                    2 ** (MAX_TENSOR_POW_OF_2 // ndim),
                )
                for blk_size in block_size
            )
            UNFOLD[ndim][grid](
                x,
                y,
                nbatch,
                *nblocks,
                *block_size,
                *im_size,
                *stride,
                *BLOCK_SIZE,
            )
    else:
        y = _unfold_torch(
            x, block_size, stride, ndim, im_size, nblocks, nbatch, out=output
        )
    return y


def _get_grid(ndim: int, nbatch, nblocks: tuple[int, ...]):
    if ndim == 1:
        grid = lambda meta: (  # noqa: E731
            nbatch * triton.cdiv(nblocks[0], meta["x_blocks_per_grid"]),
        )
    elif ndim == 2:
        grid = lambda meta: (  # noqa: E731
            nbatch * triton.cdiv(nblocks[0], meta["x_blocks_per_grid"]),
            triton.cdiv(nblocks[1], meta["y_blocks_per_grid"]),
        )
    elif ndim == 3:
        grid = lambda meta: (  # noqa: E731
            nbatch * triton.cdiv(nblocks[0], meta["x_blocks_per_grid"]),
            triton.cdiv(nblocks[1], meta["y_blocks_per_grid"]),
            triton.cdiv(nblocks[2], meta["z_blocks_per_grid"]),
        )
    else:
        raise ValueError(f"Invalid ndim = {ndim}")
    return grid


@triton.heuristics(
    values={
        "x_blocks_per_grid": lambda args: max(
            1, triton.cdiv(args["x_nblocks"], MAX_GRID_PER_DIM)
        ),
        "x_BLOCKS_per_block": lambda args: max(
            1, triton.cdiv(args["x_block_dim"], 2**MAX_TENSOR_POW_OF_2)
        ),
    },
)
@triton.jit
def _unfold1d(
    in_ptr,
    out_ptr,
    # Number of batches
    nbatch: int,
    # Number of blocks
    x_nblocks: int,
    # Size of each block
    x_block_dim: int,
    # Size of the input data
    x_size: int,
    # Stride of the blocks
    x_stride: int,
    # Size of the triton block (power of 2)
    X_BLOCK_SIZE: tl.constexpr,
    # Number of blocks per grid pid
    x_blocks_per_grid: int,
    # Number of triton blocks per block
    x_BLOCKS_per_block: int,
):
    """
    Note: Cannot use make_block_ptr for out_ptr because the output block
    might require masking.
    """
    pid_0 = tl.program_id(0)
    # Batch index, Block index
    NBx = pid_0 * x_blocks_per_grid
    N, Bx = NBx // x_nblocks, NBx % x_nblocks

    # Convert types
    x_nblocks = cast(x_nblocks, tl.uint64)
    x_size = cast(x_size, tl.uint64)
    x_block_dim = cast(x_block_dim, tl.uint64)

    in_size = x_size
    nblocks = x_nblocks
    block_dim = x_block_dim

    in_blk_ptr = tl.make_block_ptr(
        in_ptr,
        shape=(nbatch, x_size),
        strides=(in_size, 1),
        offsets=(N, Bx * x_stride),
        block_shape=(1, X_BLOCK_SIZE),
        order=(0, 1),
    )
    x_base_range = tl.arange(0, X_BLOCK_SIZE)
    for i in range(x_blocks_per_grid):
        if Bx + i < x_nblocks:
            for u in range(x_BLOCKS_per_block):
                x_range = x_base_range + u * X_BLOCK_SIZE
                x_mask = x_range < x_block_dim
                blk_range = x_range
                blk_mask = x_mask

                out_range = blk_range[None]
                out_mask = blk_mask[None]
                blk = load_subblock1d(in_blk_ptr, u, X_BLOCK_SIZE)
                # Save block to output
                out_offset = N * nblocks * block_dim + (Bx + i) * block_dim
                tl.store(out_ptr + out_offset + out_range, blk, out_mask)
        in_blk_ptr = tl.advance(in_blk_ptr, (0, x_stride))


@triton.jit
def load_subblock1d(in_blk_ptr, x_idx: int, X_BLOCK_SIZE: tl.constexpr):
    return tl.load(
        in_blk_ptr.advance((0, x_idx * X_BLOCK_SIZE)),
        boundary_check=(1,),
        padding_option="zero",
    )


@triton.heuristics(
    values={
        "x_blocks_per_grid": lambda args: max(
            1, triton.cdiv(args["x_nblocks"], MAX_GRID_PER_DIM)
        ),
        "y_blocks_per_grid": lambda args: max(
            1, triton.cdiv(args["y_nblocks"], MAX_GRID_PER_DIM)
        ),
        "x_BLOCKS_per_block": lambda args: max(
            1, triton.cdiv(args["x_block_dim"], 2 ** (MAX_TENSOR_POW_OF_2 // 2))
        ),
        "y_BLOCKS_per_block": lambda args: max(
            1, triton.cdiv(args["y_block_dim"], 2 ** (MAX_TENSOR_POW_OF_2 // 2))
        ),
    },
)
@triton.jit
def _unfold2d(
    in_ptr,
    out_ptr,
    # Number of batches
    nbatch: int,
    # Number of blocks
    x_nblocks: int,
    y_nblocks: int,
    # Size of each block
    x_block_dim: int,
    y_block_dim: int,
    # Size of the input data
    x_size: int,
    y_size: int,
    # Stride of the blocks
    x_stride: int,
    y_stride: int,
    # Size of the triton block (power of 2)
    X_BLOCK_SIZE: tl.constexpr,
    Y_BLOCK_SIZE: tl.constexpr,
    # Number of blocks per grid pid
    x_blocks_per_grid: int,
    y_blocks_per_grid: int,
    # Number of triton blocks per block
    x_BLOCKS_per_block: int,
    y_BLOCKS_per_block: int,
):
    pid_0 = tl.program_id(0)
    pid_1 = tl.program_id(1)
    # Batch index, Block index
    NBx = pid_0 * x_blocks_per_grid
    N, Bx = NBx // x_nblocks, NBx % x_nblocks
    By = pid_1 * y_blocks_per_grid

    # Convert types
    x_nblocks = cast(x_nblocks, tl.uint64)
    y_nblocks = cast(y_nblocks, tl.uint64)
    x_size = cast(x_size, tl.uint64)
    y_size = cast(y_size, tl.uint64)
    x_block_dim = cast(x_block_dim, tl.uint64)
    y_block_dim = cast(y_block_dim, tl.uint64)

    # global sizes
    in_size = x_size * y_size
    nblocks = x_nblocks * y_nblocks
    block_dim = x_block_dim * y_block_dim

    in_blk_ptr = tl.make_block_ptr(
        in_ptr,
        shape=(nbatch, x_size, y_size),
        strides=(in_size, y_size, 1),
        offsets=(N, Bx * x_stride, By * y_stride),
        block_shape=(1, X_BLOCK_SIZE, Y_BLOCK_SIZE),
        order=(0, 1, 2),
    )
    x_base_range = tl.arange(0, X_BLOCK_SIZE)
    y_base_range = tl.arange(0, Y_BLOCK_SIZE)

    for i in range(x_blocks_per_grid):
        if Bx + i < x_nblocks:
            x_blk_offset = (Bx + i) * y_nblocks * block_dim
            in_blk_ptr_x = in_blk_ptr
            for j in range(y_blocks_per_grid):
                if By + j < y_nblocks:
                    y_blk_offset = (By + j) * block_dim
                    # Loop over blocks within block
                    for u in range(x_BLOCKS_per_block):
                        for v in range(y_BLOCKS_per_block):
                            x_range = x_base_range + u * X_BLOCK_SIZE
                            y_range = y_base_range + v * Y_BLOCK_SIZE
                            x_mask = x_range < x_block_dim
                            y_mask = y_range < y_block_dim
                            blk_range = (
                                x_range[:, None] * y_block_dim + y_range[None, :]
                            )
                            blk_mask = x_mask[:, None] & y_mask[None, :]
                            # out_offset = N * x_nblocks * x_block_dim + Bx * x_block_dim
                            # add batch dim
                            out_range = blk_range[None, :, :]
                            out_mask = blk_mask[None, :, :]
                            # blk = tl.load(in_blk_ptr)
                            blk = load_subblock2d(
                                in_blk_ptr, u, v, X_BLOCK_SIZE, Y_BLOCK_SIZE
                            )
                            out_offset = (
                                N * nblocks * block_dim + x_blk_offset + y_blk_offset
                            )
                            tl.store(out_ptr + out_offset + out_range, blk, out_mask)
                in_blk_ptr = tl.advance(in_blk_ptr, (0, 0, y_stride))
            in_blk_ptr = in_blk_ptr_x
        in_blk_ptr = tl.advance(in_blk_ptr, (0, x_stride, 0))


@triton.jit
def load_subblock2d(in_blk_ptr, x_idx, y_idx, X_BLOCK_SIZE, Y_BLOCK_SIZE):
    return tl.load(
        in_blk_ptr.advance((0, x_idx * X_BLOCK_SIZE, y_idx * Y_BLOCK_SIZE)),
        boundary_check=(1, 2),
        padding_option="zero",
    )


@triton.heuristics(
    values={
        "x_blocks_per_grid": lambda args: max(
            1, triton.cdiv(args["x_nblocks"], MAX_GRID_PER_DIM)
        ),
        "y_blocks_per_grid": lambda args: max(
            1, triton.cdiv(args["y_nblocks"], MAX_GRID_PER_DIM)
        ),
        "z_blocks_per_grid": lambda args: max(
            1, triton.cdiv(args["z_nblocks"], MAX_GRID_PER_DIM)
        ),
        "x_BLOCKS_per_block": lambda args: max(
            1, triton.cdiv(args["x_block_dim"], 2 ** (MAX_TENSOR_POW_OF_2 // 3))
        ),
        "y_BLOCKS_per_block": lambda args: max(
            1, triton.cdiv(args["y_block_dim"], 2 ** (MAX_TENSOR_POW_OF_2 // 3))
        ),
        "z_BLOCKS_per_block": lambda args: max(
            1, triton.cdiv(args["z_block_dim"], 2 ** (MAX_TENSOR_POW_OF_2 // 3))
        ),
    },
)
@triton.jit
def _unfold3d(
    in_ptr,
    out_ptr,
    # Number of batches
    nbatch: int,
    # Number of blocks
    x_nblocks: int,
    y_nblocks: int,
    z_nblocks: int,
    # Size of each block
    x_block_dim: int,
    y_block_dim: int,
    z_block_dim: int,
    # Size of the input data
    x_size: int,
    y_size: int,
    z_size: int,
    # Stride of the blocks
    x_stride: int,
    y_stride: int,
    z_stride: int,
    # Size of the triton block (power of 2)
    X_BLOCK_SIZE: tl.constexpr,
    Y_BLOCK_SIZE: tl.constexpr,
    Z_BLOCK_SIZE: tl.constexpr,
    # Number of blocks per grid pid
    x_blocks_per_grid: int,
    y_blocks_per_grid: int,
    z_blocks_per_grid: int,
    # Number of triton blocks per block
    x_BLOCKS_per_block: int,
    y_BLOCKS_per_block: int,
    z_BLOCKS_per_block: int,
):
    """"""
    pid_0 = tl.program_id(0)
    pid_1 = tl.program_id(1)
    pid_2 = tl.program_id(2)
    # Batch index, Block index
    NBx = pid_0 * x_blocks_per_grid
    N, Bx = NBx // x_nblocks, NBx % x_nblocks
    By = pid_1 * y_blocks_per_grid
    Bz = pid_2 * z_blocks_per_grid

    # Convert types
    x_nblocks = cast(x_nblocks, tl.uint64)
    y_nblocks = cast(y_nblocks, tl.uint64)
    z_nblocks = cast(z_nblocks, tl.uint64)
    x_size = cast(x_size, tl.uint64)
    y_size = cast(y_size, tl.uint64)
    z_size = cast(z_size, tl.uint64)
    x_block_dim = cast(x_block_dim, tl.uint64)
    y_block_dim = cast(y_block_dim, tl.uint64)
    z_block_dim = cast(z_block_dim, tl.uint64)

    # global sizes
    in_size = x_size * y_size * z_size
    nblocks = x_nblocks * y_nblocks * z_nblocks
    block_dim = x_block_dim * y_block_dim * z_block_dim

    in_blk_ptr = tl.make_block_ptr(
        in_ptr,
        shape=(nbatch, x_size, y_size, z_size),
        strides=(in_size, y_size * z_size, z_size, 1),
        offsets=(N, Bx * x_stride, By * y_stride, Bz * z_stride),
        block_shape=(1, X_BLOCK_SIZE, Y_BLOCK_SIZE, Z_BLOCK_SIZE),
        order=(0, 1, 2, 3),
    )
    x_base_range = tl.arange(0, X_BLOCK_SIZE)
    y_base_range = tl.arange(0, Y_BLOCK_SIZE)
    z_base_range = tl.arange(0, Z_BLOCK_SIZE)

    for i in range(x_blocks_per_grid):
        if Bx + i < x_nblocks:
            x_blk_offset = (Bx + i) * y_nblocks * z_nblocks * block_dim
            in_blk_ptr_x = in_blk_ptr
            for j in range(y_blocks_per_grid):
                if By + j < y_nblocks:
                    y_blk_offset = (By + j) * z_nblocks * block_dim
                    in_blk_ptr_y = in_blk_ptr
                    for k in range(z_blocks_per_grid):
                        if Bz + k < z_nblocks:
                            z_blk_offset = (Bz + k) * block_dim
                            # Loop over subblocks
                            for u in range(x_BLOCKS_per_block):
                                for v in range(y_BLOCKS_per_block):
                                    for w in range(z_BLOCKS_per_block):
                                        x_range = x_base_range + u * X_BLOCK_SIZE
                                        y_range = y_base_range + v * Y_BLOCK_SIZE
                                        z_range = z_base_range + w * Z_BLOCK_SIZE
                                        x_mask = x_range < x_block_dim
                                        y_mask = y_range < y_block_dim
                                        z_mask = z_range < z_block_dim
                                        blk_range = (
                                            x_range[:, None, None] * y_block_dim
                                            + y_range[None, :, None]
                                        ) * z_block_dim + z_range[None, None, :]
                                        blk_mask = x_mask[:, None, None] & (
                                            y_mask[None, :, None]
                                            & z_mask[None, None, :]
                                        )

                                        out_range = blk_range[None]
                                        out_mask = blk_mask[None]
                                        blk = load_subblock3d(
                                            in_blk_ptr,
                                            u,
                                            v,
                                            w,
                                            X_BLOCK_SIZE,
                                            Y_BLOCK_SIZE,
                                            Z_BLOCK_SIZE,
                                        )
                                        out_offset = (
                                            N * nblocks * block_dim
                                            + x_blk_offset
                                            + y_blk_offset
                                            + z_blk_offset
                                        )
                                        tl.store(
                                            out_ptr + out_offset + out_range,
                                            blk,
                                            out_mask,
                                        )
                        in_blk_ptr = tl.advance(in_blk_ptr, (0, 0, 0, z_stride))
                    in_blk_ptr = in_blk_ptr_y
                in_blk_ptr = tl.advance(in_blk_ptr, (0, 0, y_stride, 0))
            in_blk_ptr = in_blk_ptr_x
        in_blk_ptr = tl.advance(in_blk_ptr, (0, x_stride, 0, 0))


@triton.jit
def load_subblock3d(
    in_blk_ptr, x_idx, y_idx, z_idx, X_BLOCK_SIZE, Y_BLOCK_SIZE, Z_BLOCK_SIZE
):
    return tl.load(
        in_blk_ptr.advance(
            (0, x_idx * X_BLOCK_SIZE, y_idx * Y_BLOCK_SIZE, z_idx * Z_BLOCK_SIZE)
        ),
        boundary_check=(1, 2, 3),
        padding_option="zero",
    )


if TRITON_ENABLED:
    UNFOLD = {1: _unfold1d, 2: _unfold2d, 3: _unfold3d}
else:
    UNFOLD = {}


def _unfold_torch(
    x: Float[Tensor, "B ..."],
    block_size: tuple[int, ...],
    stride: tuple[int, ...],
    ndim: int,
    im_size: tuple[int, ...],
    nblocks: tuple[int, ...],
    nbatch: int,
    out: Optional[Tensor] = None,
) -> Float[Tensor, "B I ..."]:
    """Fallback option

    Note: Compile takes forever
    """
    if out is None:
        out = torch.zeros(
            (nbatch, *nblocks, *block_size), device=x.device, dtype=x.dtype
        )
    else:
        out = out.reshape(nbatch, *nblocks, *block_size).zero_()
    # Python implementation
    for batch in range(nbatch):
        for blk in product(*(range(nblk) for nblk in nblocks)):
            blk_slc = tuple(
                slice(iblk * st, iblk * st + blk_sz)
                for iblk, st, blk_sz in zip(blk, stride, block_size)
            )
            out_idx = (batch, *blk)
            in_idx = (batch, *blk_slc)
            out[out_idx] = x[in_idx]
    return out


def _prep_unfold(
    x,
    block_size: tuple,
    stride: Optional[tuple] = None,
    mask: Optional[Bool[Tensor, "..."]] = None,
):
    is_complex = torch.is_complex(x)
    # Infer some shapes
    ndim = len(block_size)
    im_size = x.shape[-ndim:]
    stride = stride if stride is not None else (1,) * ndim
    nblocks = get_nblocks(im_size, block_size, stride)
    if any(b == 0 for b in nblocks):
        raise ValueError(
            f"Found 0 in nblocks: {nblocks} with im_size {im_size}, block_size {block_size} and stride {stride} - make sure there is at least one block in each direction."
        )

    if is_complex:
        im_size = list(im_size)
        im_size[-1] *= 2
        block_size = list(block_size)
        block_size[-1] *= 2
        stride = list(stride)
        stride[-1] *= 2

    # Add or infer batch dim
    batch_shape = x.shape[:-ndim]
    if ndim < len(x.shape):
        x_flat = x.flatten(0, len(x.shape) - ndim - 1)
    else:
        x_flat = x[None]
    nbatch = x_flat.shape[0]

    # Handle mask
    if mask is not None:
        mask = mask.to(x.device)

    return (
        x_flat,
        {
            "ndim": ndim,
            "im_size": im_size,
            "stride": stride,
            "nblocks": nblocks,
            "nbatch": nbatch,
            "batch_shape": batch_shape,
            "mask": mask,
            "block_size": block_size,
        },
        is_complex,
    )
