from typing import List

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


def conv_layer(
    in_channels,
    out_channels,
    kernel_size,
    stride=1,
    padding=0,
    bias=True,
    activation="silu",
    batch_norm=False,
):
    layers = [
        nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            bias=bias,
        )
    ]
    if batch_norm:
        layers.append(nn.BatchNorm2d(out_channels))
    if activation is not None:
        if activation == "silu":
            layers.append(nn.SiLU())
        elif activation == "leakyrelu":
            layers.append(nn.LeakyReLU(negative_slope=0.2))
        else:
            layers.append(nn.ReLU())
    return nn.Sequential(*layers)


def fc_layer(in_channels, out_channels, bias=True, activation=nn.ReLU, batch_norm=False):
    layers = [nn.Linear(int(in_channels), int(out_channels), bias=bias)]
    if batch_norm:
        layers.append(nn.BatchNorm1d(out_channels))
    if activation is not None:
        layers.append(activation())
    return nn.Sequential(*layers)


def rgb_to_grayscale(x):
    weights = x.new_tensor([0.2989, 0.5870, 0.1140]).view(1, 3, 1, 1)
    return (x * weights).sum(dim=1, keepdim=True)


def slicing(grid, guide):
    n, _, h, w = guide.shape
    device = grid.device
    hh, ww = torch.meshgrid(
        torch.arange(h, device=device),
        torch.arange(w, device=device),
        indexing="ij",
    )
    hh = hh / (h - 1) * 2 - 1
    ww = ww / (w - 1) * 2 - 1
    guide = guide * 2 - 1
    hh = hh[None, :, :, None].repeat(n, 1, 1, 1)
    ww = ww[None, :, :, None].repeat(n, 1, 1, 1)
    guide = guide.permute(0, 2, 3, 1)
    guide_coords = torch.cat([ww, hh, guide], dim=3).unsqueeze(1)
    sliced = F.grid_sample(grid, guide_coords, align_corners=False, padding_mode="border")
    return sliced.squeeze(2)


def apply(sliced, fullres):
    rr = torch.sum(fullres * sliced[:, 0:3, :, :], dim=1) + sliced[:, 3, :, :]
    gg = torch.sum(fullres * sliced[:, 4:7, :, :], dim=1) + sliced[:, 7, :, :]
    bb = torch.sum(fullres * sliced[:, 8:11, :, :], dim=1) + sliced[:, 11, :, :]
    return torch.stack([rr, gg, bb], dim=1)


class Guide(nn.Module):
    def __init__(self, c_in=3):
        super().__init__()
        self.nrelus = 16
        self.c_in = c_in
        self.M = nn.Parameter(
            torch.eye(c_in, dtype=torch.float32) + torch.randn(1, dtype=torch.float32) * 1e-4
        )
        self.M_bias = nn.Parameter(torch.zeros(c_in, dtype=torch.float32))
        thresholds = np.linspace(0, 1, self.nrelus, endpoint=False, dtype=np.float32)
        thresholds = torch.tensor(thresholds)[None, None, None, :].repeat(1, 1, c_in, 1)
        self.thresholds = nn.Parameter(thresholds)
        slopes = torch.zeros(1, 1, 1, c_in, self.nrelus, dtype=torch.float32)
        slopes[:, :, :, :, 0] = 1.0
        self.slopes = nn.Parameter(slopes)
        self.relu = nn.ReLU()
        self.bias = nn.Parameter(torch.tensor(0, dtype=torch.float32))

    def forward(self, x):
        x = x.permute(0, 2, 3, 1)
        old_shape = x.shape
        x = torch.matmul(x.reshape(-1, self.c_in), self.M) + self.M_bias
        x = x.reshape(old_shape).unsqueeze(4)
        x = torch.sum(self.slopes * self.relu(x - self.thresholds), dim=4)
        x = x.permute(0, 3, 1, 2)
        x = torch.sum(x, dim=1, keepdim=True) / self.c_in
        return torch.clamp(x + self.bias, 0, 1)


class Biliteral_Grid_Joint(nn.Module):
    def __init__(
        self,
        in_channels: int = 3,
        channels: List[int] = None,
        fix_guide: bool = False,
        grid_res: int = 16,
        grid_bins: int = 8,
    ):
        super().__init__()
        del in_channels, channels
        bn = False
        activation = "relu"
        self.grid_res = grid_res
        self.grid_bins = grid_bins

        if grid_res == 16:
            self.down1 = conv_layer(320, 256, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.down2 = conv_layer(256, 128, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.local1 = conv_layer(128, 128, kernel_size=3, padding=1, bias=False, activation=None)
            self.global1 = conv_layer(128, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global2 = conv_layer(64, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global3 = fc_layer(4 * 4 * 64, 256, batch_norm=bn)
            self.global4 = fc_layer(256, 128, activation=None)
        elif grid_res == 32:
            self.down1 = conv_layer(320, 256, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.down2 = conv_layer(256, 128, kernel_size=3, padding=1, batch_norm=bn, activation=activation)
            self.local1 = conv_layer(128, 128, kernel_size=3, padding=1, bias=False, activation=None)
            self.global1 = conv_layer(128, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global2 = conv_layer(64, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global2_1 = conv_layer(64, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global3 = fc_layer(4 * 4 * 64, 256, batch_norm=bn)
            self.global4 = fc_layer(256, 128, activation=None)
        elif grid_res == 8:
            self.down1 = conv_layer(320, 256, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.down2 = conv_layer(256, 128, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.down3 = conv_layer(128, 128, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.local1 = conv_layer(128, 128, kernel_size=3, padding=1, bias=False, activation=None)
            self.global1 = conv_layer(128, 64, kernel_size=3, padding=1, batch_norm=bn, activation=activation)
            self.global2 = conv_layer(64, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global3 = fc_layer(4 * 4 * 64, 256, batch_norm=bn)
            self.global4 = fc_layer(256, 128, activation=None)
        elif grid_res == 64:
            self.down1 = conv_layer(320, 256, kernel_size=3, padding=1, batch_norm=bn, activation=activation)
            self.down2 = conv_layer(256, 128, kernel_size=3, padding=1, batch_norm=bn, activation=activation)
            self.local1 = conv_layer(128, 128, kernel_size=3, padding=1, bias=False, activation=None)
            self.global1 = conv_layer(128, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global2 = conv_layer(64, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global2_1 = conv_layer(64, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global3 = fc_layer(8 * 8 * 64, 256, batch_norm=bn)
            self.global4 = fc_layer(256, 128, activation=None)
        else:
            raise NotImplementedError(f"unsupported grid_res={grid_res}")

        self.pred = conv_layer(128, 12 * self.grid_bins, kernel_size=1, activation=None)
        self.relu = nn.ReLU()
        self.fix_guide = fix_guide
        self.guide = Guide()

    def forward(self, feat, x_full, only_gen_grid=False, return_guide=False):
        x = feat[0]
        n = x.shape[0]
        x = self.down1(x)
        x = self.down2(x)
        if self.grid_res == 8:
            x = self.down3(x)
        downed = x

        local_out = self.local1(x)
        x = self.global1(downed)
        x = self.global2(x)
        if self.grid_res in (32, 64):
            x = self.global2_1(x)
        x = x.reshape(n, -1)
        x = self.global3(x)
        global_out = self.global4(x)
        fusion = self.relu(local_out + global_out[:, :, None, None])

        x = self.pred(fusion)
        x = x.view(n, 12, self.grid_bins, self.grid_res, self.grid_res)
        coeffs = x.reshape(x.shape[0], 12, -1, x.shape[-2], x.shape[-1])
        if only_gen_grid:
            return coeffs

        guide = rgb_to_grayscale(x_full) if self.fix_guide else self.guide(x_full)
        out = apply(slicing(coeffs, guide), x_full)
        if return_guide:
            return out, coeffs, guide
        return out, coeffs


class Bilateral_Grid_Joint_Flux(nn.Module):
    def __init__(
        self,
        hidden_dim: int = 3072,
        proj_channels: int = 256,
        img_token_start: int = 0,
        fix_guide: bool = False,
        grid_res: int = 16,
        grid_bins: int = 8,
    ):
        super().__init__()
        bn = False
        activation = "relu"
        self.grid_res = grid_res
        self.grid_bins = grid_bins
        self.img_token_start = img_token_start
        self.channel_proj = nn.Sequential(
            nn.Linear(hidden_dim, proj_channels),
            nn.GELU(),
            nn.Linear(proj_channels, proj_channels),
        )

        if grid_res == 16:
            self.down1 = conv_layer(proj_channels, 128, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.local1 = conv_layer(128, 128, kernel_size=3, padding=1, bias=False, activation=None)
            self.global1 = conv_layer(128, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global2 = conv_layer(64, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global3 = fc_layer(4 * 4 * 64, 256, batch_norm=bn)
            self.global4 = fc_layer(256, 128, activation=None)
        elif grid_res == 32:
            self.down1 = conv_layer(proj_channels, 128, kernel_size=3, padding=1, batch_norm=bn, activation=activation)
            self.local1 = conv_layer(128, 128, kernel_size=3, padding=1, bias=False, activation=None)
            self.global1 = conv_layer(128, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global2 = conv_layer(64, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global2_1 = conv_layer(64, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global3 = fc_layer(4 * 4 * 64, 256, batch_norm=bn)
            self.global4 = fc_layer(256, 128, activation=None)
        elif grid_res == 8:
            self.down1 = conv_layer(proj_channels, 128, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.down2 = conv_layer(128, 128, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.local1 = conv_layer(128, 128, kernel_size=3, padding=1, bias=False, activation=None)
            self.global1 = conv_layer(128, 64, kernel_size=3, padding=1, batch_norm=bn, activation=activation)
            self.global2 = conv_layer(64, 64, kernel_size=3, stride=2, padding=1, batch_norm=bn, activation=activation)
            self.global3 = fc_layer(4 * 4 * 64, 256, batch_norm=bn)
            self.global4 = fc_layer(256, 128, activation=None)
        else:
            raise NotImplementedError(f"unsupported grid_res={grid_res}")

        self.pred = conv_layer(128, 12 * self.grid_bins, kernel_size=1, activation=None)
        self.relu = nn.ReLU()
        self.fix_guide = fix_guide
        self.guide = Guide()

    def _extract_image_tokens(self, feat, latent_h, latent_w):
        bsz = feat.shape[0]
        num_img_tokens = latent_h * latent_w
        img_tokens = feat[:, self.img_token_start : self.img_token_start + num_img_tokens, :]
        img_tokens = self.channel_proj(img_tokens)
        return img_tokens.permute(0, 2, 1).reshape(bsz, -1, latent_h, latent_w)

    def forward(
        self,
        feat,
        x_full,
        only_gen_grid=False,
        return_guide=False,
        latent_h: int = 32,
        latent_w: int = 32,
    ):
        if isinstance(feat, (list, tuple)):
            feat = feat[0]
        x = self._extract_image_tokens(feat, latent_h, latent_w)
        n = x.shape[0]

        x = self.down1(x)
        if self.grid_res == 8:
            x = self.down2(x)
        downed = x
        local_out = self.local1(x)

        x = self.global1(downed)
        x = self.global2(x)
        if self.grid_res == 32:
            x = self.global2_1(x)
        x = x.reshape(n, -1)
        x = self.global3(x)
        global_out = self.global4(x)
        fusion = self.relu(local_out + global_out[:, :, None, None])

        x = self.pred(fusion)
        x = x.view(n, 12, self.grid_bins, self.grid_res, self.grid_res)
        coeffs = x.reshape(x.shape[0], 12, -1, x.shape[-2], x.shape[-1])
        if only_gen_grid:
            return coeffs

        guide = rgb_to_grayscale(x_full) if self.fix_guide else self.guide(x_full)
        out = apply(slicing(coeffs, guide), x_full)
        if return_guide:
            return out, coeffs, guide
        return out, coeffs