# Modified for SplatAnything: inference-only release.
from typing import Callable, Optional, Tuple, Union
import torch
from torch import Tensor
import torch.nn as nn
import torch.nn.functional as F
from itertools import repeat
import collections.abc

def make_2tuple(x):
    if isinstance(x, tuple):
        assert len(x) == 2
        return x
    assert isinstance(x, int)
    return (x, x)

class PatchEmbed(nn.Module):
    """
    2D image to patch embedding: (B,C,H,W) -> (B,N,D)

    Args:
        img_size: Image size.
        patch_size: Patch token size.
        in_chans: Number of input image channels.
        embed_dim: Number of linear projection output channels.
        norm_layer: Normalization layer.
    """

    def __init__(self, img_size: Union[int, Tuple[int, int]]=224, patch_size: Union[int, Tuple[int, int]]=16, in_chans: int=3, embed_dim: int=768, norm_layer: Optional[Callable]=None, flatten_embedding: bool=True) -> None:
        super().__init__()
        image_HW = make_2tuple(img_size)
        patch_HW = make_2tuple(patch_size)
        patch_grid_size = (image_HW[0] // patch_HW[0], image_HW[1] // patch_HW[1])
        self.img_size = image_HW
        self.patch_size = patch_HW
        self.patches_resolution = patch_grid_size
        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
        self.in_chans = in_chans
        self.embed_dim = embed_dim
        self.flatten_embedding = flatten_embedding
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x: Tensor) -> Tensor:
        (_, _, H, W) = x.shape
        (patch_H, patch_W) = self.patch_size
        assert H % patch_H == 0, f'Input image height {H} is not a multiple of patch height {patch_H}'
        assert W % patch_W == 0, f'Input image width {W} is not a multiple of patch width: {patch_W}'
        x = self.proj(x)
        (H, W) = (x.size(2), x.size(3))
        x = x.flatten(2).transpose(1, 2)
        x = self.norm(x)
        if not self.flatten_embedding:
            x = x.reshape(-1, H, W, self.embed_dim)
        return x

class PatchEmbed_Mlp(PatchEmbed):

    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten_embedding=True):
        super().__init__(img_size, patch_size, in_chans, embed_dim, norm_layer, flatten_embedding)
        self.proj = nn.Sequential(PixelUnshuffle(patch_size), Permute((0, 2, 3, 1)), Mlp(in_chans * patch_size ** 2, 4 * embed_dim, embed_dim), Permute((0, 3, 1, 2)))

class PixelUnshuffle(nn.Module):

    def __init__(self, downscale_factor):
        super().__init__()
        self.downscale_factor = downscale_factor

    def forward(self, input):
        if input.numel() == 0:
            (C, H, W) = input.shape[-3:]
            assert H and W and (H % self.downscale_factor == W % self.downscale_factor == 0)
            return input.view(*input.shape[:-3], C * self.downscale_factor ** 2, H // self.downscale_factor, W // self.downscale_factor)
        else:
            return F.pixel_unshuffle(input, self.downscale_factor)

class Permute(torch.nn.Module):
    dims: tuple[int, ...]

    def __init__(self, dims: tuple[int, ...]) -> None:
        super().__init__()
        self.dims = tuple(dims)

    def __repr__(self):
        return f'Permute{self.dims}'

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return input.permute(*self.dims)

def _ntuple(n):

    def parse(x):
        if isinstance(x, collections.abc.Iterable) and (not isinstance(x, str)):
            return x
        return tuple(repeat(x, n))
    return parse
to_2tuple = _ntuple(2)

class Mlp(nn.Module):
    """ MLP as used in Vision Transformer, MLP-Mixer and related networks"""

    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, drop=0.0):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        bias = to_2tuple(bias)
        drop_probs = to_2tuple(drop)
        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias[0])
        self.act = act_layer()
        self.drop1 = nn.Dropout(drop_probs[0])
        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias[1])
        self.drop2 = nn.Dropout(drop_probs[1])

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop1(x)
        x = self.fc2(x)
        x = self.drop2(x)
        return x