from functools import partial

import math
import numpy as np
import scipy.stats as stats
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
from .util import visualize_patch

from timm.models.vision_transformer import DropPath, Mlp
from .pixelloss import PixelLoss


def mask_by_order(mask_len, order, bsz, seq_len):
    masking = torch.zeros(bsz, seq_len, device=order.device)
    masking = torch.scatter(masking, dim=-1, index=order[:, :mask_len.long()], src=torch.ones(bsz, seq_len, device=order.device)).bool()
    return masking


class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
        self.scale = qk_scale or head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)

        with torch.amp.autocast('cuda', enabled=False):
            attn = (q.float() @ k.float().transpose(-2, -1)) * self.scale

        attn = attn - torch.max(attn, dim=-1, keepdim=True)[0]
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

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


class Block(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, 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 = Attention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, 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 MAR(nn.Module):
    def __init__(self, seq_len, patch_size, cond_embed_dim, embed_dim, num_blocks, num_heads, attn_dropout, proj_dropout,
                 num_conds=1, guiding_pixel=False, grad_checkpointing=False
                 ):
        super().__init__()

        self.seq_len = seq_len
        self.patch_size = patch_size

        self.num_conds = num_conds
        self.guiding_pixel = guiding_pixel

        self.grad_checkpointing = grad_checkpointing

        # --------------------------------------------------------------------------
        # variant masking ratio
        self.mask_ratio_generator = stats.truncnorm(-4, 0, loc=1.0, scale=0.25)

        # --------------------------------------------------------------------------
        # network
        self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.patch_emb = nn.Linear(3 * patch_size ** 2, embed_dim, bias=True)
        self.patch_emb_ln = nn.LayerNorm(embed_dim, eps=1e-6)
        self.cond_emb = nn.Linear(cond_embed_dim, embed_dim, bias=True)
        if self.guiding_pixel:
            self.pix_proj = nn.Linear(3, embed_dim, bias=True)
        self.pos_embed_learned = nn.Parameter(torch.zeros(1, seq_len+num_conds+self.guiding_pixel, embed_dim))

        self.blocks = nn.ModuleList([
            Block(embed_dim, num_heads, mlp_ratio=4.,
                  qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
                  proj_drop=proj_dropout, attn_drop=attn_dropout)
            for _ in range(num_blocks)
        ])
        self.norm = nn.LayerNorm(embed_dim, eps=1e-6)

        self.initialize_weights()

        if self.guiding_pixel:
            self.guiding_pixel_loss = PixelLoss(
                c_channels=cond_embed_dim,
                width=128,
                depth=3,
                num_heads=4,
                r_weight=5.0
            )

    def initialize_weights(self):
        # parameters
        torch.nn.init.normal_(self.mask_token, std=.02)
        torch.nn.init.normal_(self.pos_embed_learned, 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 patchify(self, x):
        bsz, c, h, w = x.shape
        p = self.patch_size
        h_, w_ = h // p, w // p

        x = x.reshape(bsz, c, h_, p, w_, p)
        x = torch.einsum('nchpwq->nhwcpq', x)
        x = x.reshape(bsz, h_ * w_, c * p ** 2)
        return x  # [n, l, d]

    def unpatchify(self, x):
        bsz = x.shape[0]
        p = self.patch_size
        h_, w_ = int(np.sqrt(self.seq_len)), int(np.sqrt(self.seq_len))

        x = x.reshape(bsz, h_, w_, 3, p, p)
        x = torch.einsum('nhwcpq->nchpwq', x)
        x = x.reshape(bsz, 3, h_ * p, w_ * p)
        return x  # [n, 3, h, w]

    def sample_orders(self, bsz):
        orders = torch.argsort(torch.rand(bsz, self.seq_len), dim=1).long()
        return orders

    def random_masking_uniform(self, x, orders):
        bsz, seq_len, embed_dim = x.shape
        num_masked_tokens = np.random.randint(seq_len) + 1
        mask = torch.zeros(bsz, seq_len, device=x.device)
        mask = torch.scatter(mask, dim=-1, index=orders[:, :num_masked_tokens],
                             src=torch.ones(bsz, seq_len, device=x.device))
        return mask

    def random_masking(self, x, orders):
        bsz, seq_len, embed_dim = x.shape
        mask_rates = self.mask_ratio_generator.rvs(bsz)
        num_masked_tokens = torch.Tensor(np.ceil(seq_len * mask_rates)).to(x.device)
        expanded_indices = torch.arange(seq_len, device=x.device).expand(bsz, seq_len)
        sorted_orders = torch.argsort(orders, dim=-1)
        mask = (expanded_indices < num_masked_tokens[:, None]).float()
        mask = torch.scatter(torch.zeros_like(mask), dim=-1, index=sorted_orders, src=mask)

        return mask

    def predict(self, x, mask, cond_list):
        x = self.patch_emb(x)

        # prepend conditions from prev generator
        for i in range(self.num_conds):
            x = torch.cat([self.cond_emb(cond_list[i]).unsqueeze(1), x], dim=1)

        # prepend guiding pixel
        if self.guiding_pixel:
            x = torch.cat([self.pix_proj(cond_list[-1]).unsqueeze(1), x], dim=1)

        # masking
        mask_with_cond = torch.cat([torch.zeros(x.size(0), self.num_conds+self.guiding_pixel, device=x.device), mask], dim=1).bool()
        x = torch.where(mask_with_cond.unsqueeze(-1), self.mask_token.to(x.dtype), x)

        # position embedding
        x = x + self.pos_embed_learned
        x = self.patch_emb_ln(x)

        # apply Transformer blocks
        if self.grad_checkpointing and not torch.jit.is_scripting() and self.training:
            for block in self.blocks:
                x = checkpoint(block, x)
        else:
            for block in self.blocks:
                x = block(x)
        x = self.norm(x)

        # return 5 conditions: middle, top, right, bottom, left
        middle_cond = x[:, self.num_conds+self.guiding_pixel:]
        bsz, seq_len, c = middle_cond.size()
        h = int(np.sqrt(seq_len))
        w = int(np.sqrt(seq_len))
        top_cond = middle_cond.reshape(bsz, h, w, c)
        top_cond = torch.cat([torch.zeros(bsz, 1, w, c, device=top_cond.device), top_cond[:, :-1]], dim=1)
        top_cond = top_cond.reshape(bsz, seq_len, c)

        right_cond = middle_cond.reshape(bsz, h, w, c)
        right_cond = torch.cat([right_cond[:, :, 1:], torch.zeros(bsz, h, 1, c, device=right_cond.device)], dim=2)
        right_cond = right_cond.reshape(bsz, seq_len, c)

        bottom_cond = middle_cond.reshape(bsz, h, w, c)
        bottom_cond = torch.cat([bottom_cond[:, 1:], torch.zeros(bsz, 1, w, c, device=bottom_cond.device)], dim=1)
        bottom_cond = bottom_cond.reshape(bsz, seq_len, c)

        left_cond = middle_cond.reshape(bsz, h, w, c)
        left_cond = torch.cat([torch.zeros(bsz, h, 1, c, device=left_cond.device), left_cond[:, :, :-1]], dim=2)
        left_cond = left_cond.reshape(bsz, seq_len, c)

        return [middle_cond, top_cond, right_cond, bottom_cond, left_cond]

    def forward(self, imgs, cond_list):
        """ training """
        # patchify to get gt
        patches = self.patchify(imgs)

        # mask tokens
        orders = self.sample_orders(bsz=patches.size(0)).to(patches.device)
        if self.training:
            mask = self.random_masking(patches, orders)
        else:
            # uniform random masking for NLL computation
            mask = self.random_masking_uniform(patches, orders)

        # guiding pixel
        if self.guiding_pixel:
            guiding_pixels = imgs.mean(-1).mean(-1)
            guiding_pixel_loss = self.guiding_pixel_loss(guiding_pixels, cond_list)
            cond_list.append(guiding_pixels)
        else:
            guiding_pixel_loss = torch.Tensor([0]).to(imgs.device).mean()

        # get condition for next level
        cond_list_next = self.predict(patches, mask, cond_list)

        # only keep those conditions and patches on mask
        for cond_idx in range(len(cond_list_next)):
            cond_list_next[cond_idx] = cond_list_next[cond_idx].reshape(cond_list_next[cond_idx].size(0) * cond_list_next[cond_idx].size(1), -1)
            cond_list_next[cond_idx] = cond_list_next[cond_idx][mask.reshape(-1).bool()]

        patches = patches.reshape(patches.size(0) * patches.size(1), -1)
        patches = patches[mask.reshape(-1).bool()]
        patches = patches.reshape(patches.size(0), 3, self.patch_size, self.patch_size)

        return patches, cond_list_next, guiding_pixel_loss

    def sample(self, cond_list, num_iter, cfg, cfg_schedule, temperature, filter_threshold, next_level_sample_function,
               visualize=False):
        """ generation """
        if cfg == 1.0:
            bsz = cond_list[0].size(0)
        else:
            bsz = cond_list[0].size(0) // 2

        # sample the guiding pixel
        if self.guiding_pixel:
            sampled_pixels = self.guiding_pixel_loss.sample(cond_list, temperature, cfg, filter_threshold)
            if not cfg == 1.0:
                sampled_pixels = torch.cat([sampled_pixels, sampled_pixels], dim=0)
            cond_list.append(sampled_pixels)

        # init token mask
        _device = cond_list[0].device
        mask = torch.ones(bsz, self.seq_len, device=_device)
        patches = torch.zeros(bsz, self.seq_len, 3 * self.patch_size**2, device=_device)
        orders = self.sample_orders(bsz).to(_device)
        num_iter = min(self.seq_len, num_iter)

        # sample image
        for step in range(num_iter):
            cur_patches = patches.clone()

            if not cfg == 1.0:
                patches = torch.cat([patches, patches], dim=0)
                mask = torch.cat([mask, mask], dim=0)

            # get next level conditions
            cond_list_next = self.predict(patches, mask, cond_list)

            # mask ratio for the next round, following MAR.
            mask_ratio = np.cos(math.pi / 2. * (step + 1) / num_iter)
            mask_len = torch.Tensor([np.floor(self.seq_len * mask_ratio)]).to(_device)

            # masks out at least one for the next iteration
            mask_len = torch.maximum(torch.Tensor([1]).to(_device),
                                     torch.minimum(torch.sum(mask, dim=-1, keepdims=True) - 1, mask_len))

            # get masking for next iteration and locations to be predicted in this iteration
            mask_next = mask_by_order(mask_len[0], orders, bsz, self.seq_len)
            if step >= num_iter - 1:
                mask_to_pred = mask[:bsz].bool()
            else:
                mask_to_pred = torch.logical_xor(mask[:bsz].bool(), mask_next.bool())
            mask = mask_next
            if not cfg == 1.0:
                mask_to_pred = torch.cat([mask_to_pred, mask_to_pred], dim=0)

            # sample token latents for this step
            for cond_idx in range(len(cond_list_next)):
                cond_list_next[cond_idx] = cond_list_next[cond_idx][mask_to_pred.nonzero(as_tuple=True)]

            # cfg schedule
            if cfg_schedule == "linear":
                cfg_iter = 1 + (cfg - 1) * (self.seq_len - mask_len[0]) / self.seq_len
            else:
                cfg_iter = cfg
            sampled_patches = next_level_sample_function(cond_list=cond_list_next, cfg=cfg_iter,
                                                         temperature=temperature, filter_threshold=filter_threshold)
            sampled_patches = sampled_patches.reshape(sampled_patches.size(0), -1)

            if not cfg == 1.0:
                mask_to_pred, _ = mask_to_pred.chunk(2, dim=0)

            cur_patches[mask_to_pred.nonzero(as_tuple=True)] = sampled_patches.to(cur_patches.dtype)
            patches = cur_patches.clone()

            # visualize generation process for colab
            if visualize:
                visualize_patch(self.unpatchify(patches))

        patches = self.unpatchify(patches)
        return patches