from functools import partial

import math
import torch
import torch.nn as nn

from timm.models.vision_transformer import DropPath, Mlp


def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None) -> torch.Tensor:
    L, S = query.size(-2), key.size(-2)
    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
    attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
    if is_causal:
        assert attn_mask is None
        temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(diagonal=0)
        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
        attn_bias.to(query.dtype)

    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
        else:
            attn_bias += attn_mask
    with torch.amp.autocast("cuda", enabled=False):
        attn_weight = query @ key.transpose(-2, -1) * scale_factor
    attn_weight += attn_bias
    attn_weight = torch.softmax(attn_weight, dim=-1)
    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
    return attn_weight @ value


class CausalAttention(nn.Module):
    def __init__(
            self,
            dim: int,
            num_heads: int = 8,
            qkv_bias: bool = False,
            qk_norm: bool = False,
            attn_drop: float = 0.0,
            proj_drop: float = 0.0,
            norm_layer: nn.Module = nn.LayerNorm
    ) -> None:
        super().__init__()
        assert dim % num_heads == 0, "dim should be divisible by num_heads"
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, C = x.shape
        qkv = (
            self.qkv(x)
            .reshape(B, N, 3, self.num_heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        q, k, v = qkv.unbind(0)
        q, k = self.q_norm(q), self.k_norm(k)

        x = scaled_dot_product_attention(
            q,
            k,
            v,
            dropout_p=self.attn_drop.p if self.training else 0.0,
            is_causal=True
        )

        x = x.transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class CausalBlock(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, proj_drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = CausalAttention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=proj_drop)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=proj_drop)

    def forward(self, x):
        x = x + self.drop_path(self.attn(self.norm1(x)))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


class MlmLayer(nn.Module):

    def __init__(self, vocab_size):
        super().__init__()
        self.bias = nn.Parameter(torch.zeros(1, vocab_size))

    def forward(self, x, word_embeddings):
        word_embeddings = word_embeddings.transpose(0, 1)
        logits = torch.matmul(x, word_embeddings)
        logits = logits + self.bias
        return logits


class PixelLoss(nn.Module):
    def __init__(self, c_channels, width, depth, num_heads, r_weight=1.0):
        super().__init__()

        self.pix_mean = torch.Tensor([0.485, 0.456, 0.406])
        self.pix_std = torch.Tensor([0.229, 0.224, 0.225])

        self.cond_proj = nn.Linear(c_channels, width)
        self.r_codebook = nn.Embedding(256, width)
        self.g_codebook = nn.Embedding(256, width)
        self.b_codebook = nn.Embedding(256, width)

        self.ln = nn.LayerNorm(width, eps=1e-6)
        self.blocks = nn.ModuleList([
            CausalBlock(width, num_heads=num_heads, mlp_ratio=4.0,
                        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
                        proj_drop=0, attn_drop=0)
            for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(width, eps=1e-6)

        self.r_weight = r_weight
        self.r_mlm = MlmLayer(256)
        self.g_mlm = MlmLayer(256)
        self.b_mlm = MlmLayer(256)

        self.criterion = torch.nn.CrossEntropyLoss(reduction="none")

        self.initialize_weights()

    def initialize_weights(self):
        # parameters
        torch.nn.init.normal_(self.r_codebook.weight, std=.02)
        torch.nn.init.normal_(self.g_codebook.weight, std=.02)
        torch.nn.init.normal_(self.b_codebook.weight, std=.02)

        # initialize nn.Linear and nn.LayerNorm
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            # we use xavier_uniform following official JAX ViT:
            torch.nn.init.xavier_uniform_(m.weight)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
            if m.weight is not None:
                nn.init.constant_(m.weight, 1.0)

    def predict(self, target, cond_list):
        target = target.reshape(target.size(0), target.size(1))
        # back to [0, 255]
        mean = self.pix_mean.to(target.device).unsqueeze(0)
        std = self.pix_std.to(target.device).unsqueeze(0)
        target = target * std + mean
        # add a very small noice to avoid pixel distribution inconsistency caused by banker's rounding
        target = (target * 255 + 1e-2 * torch.randn_like(target)).round().long()

        # take only the middle condition
        cond = cond_list[0]
        x = torch.cat(
            [self.cond_proj(cond).unsqueeze(1), self.r_codebook(target[:, 0:1]), self.g_codebook(target[:, 1:2]),
             self.b_codebook(target[:, 2:3])], dim=1)
        x = self.ln(x)

        for block in self.blocks:
            x = block(x)

        x = self.norm(x)
        with torch.amp.autocast("cuda", enabled=False):
            r_logits = self.r_mlm(x[:, 0], self.r_codebook.weight)
            g_logits = self.g_mlm(x[:, 1], self.g_codebook.weight)
            b_logits = self.b_mlm(x[:, 2], self.b_codebook.weight)

        logits = torch.cat([r_logits.unsqueeze(1), g_logits.unsqueeze(1), b_logits.unsqueeze(1)], dim=1)
        return logits, target

    def forward(self, target, cond_list):
        """ training """
        logits, target = self.predict(target, cond_list)
        loss_r = self.criterion(logits[:, 0], target[:, 0])
        loss_g = self.criterion(logits[:, 1], target[:, 1])
        loss_b = self.criterion(logits[:, 2], target[:, 2])

        if self.training:
            loss = (self.r_weight * loss_r + loss_g + loss_b) / (self.r_weight + 2)
        else:
            # for NLL computation
            loss = (loss_r + loss_g + loss_b) / 3

        return loss.mean()

    def sample(self, cond_list, temperature, cfg, filter_threshold=0):
        """ generation """
        if cfg == 1.0:
            bsz = cond_list[0].size(0)
        else:
            bsz = cond_list[0].size(0) // 2
        pixel_values = torch.zeros(bsz, 3, device=cond_list[0].device)

        for i in range(3):
            if cfg == 1.0:
                logits, _ = self.predict(pixel_values, cond_list)
            else:
                logits, _ = self.predict(torch.cat([pixel_values, pixel_values], dim=0), cond_list)
            logits = logits[:, i]
            logits = logits * temperature

            if not cfg == 1.0:
                cond_logits = logits[:bsz]
                uncond_logits = logits[bsz:]

                # very unlikely conditional logits will be suppressed
                cond_probs = torch.softmax(cond_logits, dim=-1)
                mask = cond_probs < filter_threshold
                uncond_logits[mask] = torch.max(
                    uncond_logits,
                    cond_logits - torch.max(cond_logits, dim=-1, keepdim=True)[0] + torch.max(uncond_logits, dim=-1, keepdim=True)[0]
                )[mask]

                logits = uncond_logits + cfg * (cond_logits - uncond_logits)

            # get token prediction
            probs = torch.softmax(logits, dim=-1)
            sampled_ids = torch.multinomial(probs, num_samples=1).reshape(-1)
            pixel_values[:, i] = (sampled_ids.float() / 255 - self.pix_mean[i]) / self.pix_std[i]

        # back to [0, 1]
        return pixel_values