import torch
import torch.nn as nn
from models_trans import BasicTransformer, LinearTransformer, SparseTransformer
from timm.models.layers import trunc_normal_


def create_block(d_model=192, n_heads=8, d_head=24, dropout=0.1, map_name="elu+1", 
                 block_name="basic", return_attn=False):
    if block_name == "basic":
        return BasicTransformer(d_model, n_heads, d_head, dropout=dropout, 
                                return_attn=return_attn)
    elif block_name == "basic-gated":
        return BasicTransformer(d_model, n_heads, d_head, dropout=dropout, 
                                is_gated=True, return_attn=return_attn)
    elif block_name == "flash":
        return BasicTransformer(d_model, n_heads, d_head, dropout=dropout, 
                                use_flash_attention=True, return_attn=return_attn)
    elif block_name == "flash-gated":
        return BasicTransformer(d_model, n_heads, d_head, dropout=dropout, 
                                use_flash_attention=True, is_gated=True, return_attn=return_attn)
    elif block_name == "linear":
        return LinearTransformer(d_model, dropout=dropout, map_name=map_name)
    elif block_name == "sparse":
        return SparseTransformer(d_model, n_heads, dropout=dropout)
    else:
        raise NotImplementedError(f"Block {block_name} not implemented")


from timm.models.layers import trunc_normal_, lecun_normal_
import math
import time
from functools import partial


# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
def _init_weights(
    module,
    n_layer,
    initializer_range=0.02,  # Now only used for embedding layer.
    rescale_prenorm_residual=True,
    n_residuals_per_layer=1,  # Change to 2 if we have MLP
):
    if isinstance(module, nn.Linear):
        if module.bias is not None:
            if not getattr(module.bias, "_no_reinit", False):
                nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        nn.init.normal_(module.weight, std=initializer_range)

    if rescale_prenorm_residual:
        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
        #
        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
        for name, p in module.named_parameters():
            if name in ["out_proj.weight", "fc2.weight"]:
                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
                # We need to reinit p since this code could be called multiple times
                # Having just p *= scale would repeatedly scale it down
                nn.init.kaiming_uniform_(p, a=math.sqrt(5))
                with torch.no_grad():
                    p /= math.sqrt(n_residuals_per_layer * n_layer)


def segm_init_weights(m):
    if isinstance(m, nn.Linear):
        trunc_normal_(m.weight, std=0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.constant_(m.bias, 0)
    elif isinstance(m, (nn.Conv2d, nn.Conv1d)):
        # NOTE conv was left to pytorch default in my original init
        lecun_normal_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
        nn.init.zeros_(m.bias)
        nn.init.ones_(m.weight)


class StrideEmbed(nn.Module):
    def __init__(self, arr_length=1600, stride_size=4, in_chans=1, embed_dim=192):
        super().__init__()
        assert arr_length % stride_size == 0
        self.num_patches = arr_length // stride_size
        self.proj = nn.Conv1d(in_chans, embed_dim, kernel_size=stride_size, stride=stride_size)
        
    def forward(self, x):
        """
        x: [B, N]
        """
        return self.proj(x).transpose(1, 2) # [B, N, D]


class NetTransformer(nn.Module):
    def __init__(self, 
                 arr_length=1600,
                 stride_size=4,
                 in_chans=1,
                 embed_dim=192, depth=4, 
                 decoder_embed_dim=128, decoder_depth=2,
                 num_classes=1000,
                 n_heads=8, block_name="basic",
                 norm_pix_loss=False,
                 drop_rate=0.,
                 is_pretrain=False,
                 if_cls_token=True,
                 device=None, dtype=None,
                 return_attn=False,
                 **kwargs):
        super().__init__()

        factory_kwargs = {"device": device, "dtype": dtype}
        # add factory_kwargs into kwargs
        kwargs.update(factory_kwargs) 
        self.num_classes = num_classes
        self.d_model = self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.is_pretrain = is_pretrain
        self.return_attn = return_attn
        self.stride_size = stride_size

        # --------------------------------------------------------------------------
        # NetMamba encoder specifics
        self.patch_embed = StrideEmbed(arr_length=arr_length, stride_size=stride_size, embed_dim=embed_dim)
        self.num_patches = self.patch_embed.num_patches
        self.if_cls_token = if_cls_token
        if if_cls_token:
            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
            self.num_cls_token = 1
        else:
            self.num_cls_token = 0
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + self.num_cls_token, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_rate)
        # Mamba blocks
        self.blocks = nn.ModuleList([
            create_block(d_model=embed_dim, n_heads=n_heads, d_head=embed_dim // n_heads, dropout=0.1, 
                         block_name=block_name, return_attn=return_attn)
            for _ in range(depth)])
        # --------------------------------------------------------------------------

        if is_pretrain:
            # --------------------------------------------------------------------------
            # NetMamba decoder specifics
            self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
            self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
            self.decoder_pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + self.num_cls_token, decoder_embed_dim))
            self.decoder_blocks = nn.ModuleList([
                create_block(d_model=decoder_embed_dim, n_heads=n_heads, d_head=decoder_embed_dim // n_heads, dropout=0.1, 
                             block_name=block_name, return_attn=return_attn)
                for _ in range(decoder_depth)])
            self.decoder_pred = nn.Linear(decoder_embed_dim, stride_size * in_chans, bias=True)  # decoder to stride
            # --------------------------------------------------------------------------
        else:
            # --------------------------------------------------------------------------
            # NetMamba classifier specifics
            self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
            # --------------------------------------------------------------------------

        self.norm_pix_loss = norm_pix_loss
        self.initialize_weights(depth)

    def initialize_weights(self, depth):
        self.patch_embed.apply(segm_init_weights)
        if not self.is_pretrain:
            self.head.apply(segm_init_weights)
        trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)
        if self.is_pretrain:
            trunc_normal_(self.decoder_pos_embed, std=.02)
            trunc_normal_(self.mask_token, std=.02)
        # initialize nn.Linear and nn.LayerNorm
        self.apply(partial(_init_weights, n_layer=depth,))
    
    @torch.jit.ignore
    def no_weight_decay(self):
        return {"pos_embed", "cls_token", "dist_token", "cls_token_head", "cls_token_tail"}
    
    def stride_patchify(self, imgs):
        """
        imgs: (N, 1, H, W)
        x: (N, L, patch_size**2 *1)
        """
        B, C, H, W = imgs.shape
        assert C == 1, "Input images should be grayscale"
        stride_size = self.stride_size
        x = imgs.reshape(B, H*W // stride_size, stride_size)
        return x

    def random_masking(self, x, mask_ratio):
        """
        Perform per-sample random masking by per-sample shuffling.
        Per-sample shuffling is done by argsort random noise.
        x: [B N D], sequence
        """
        B, N, D = x.shape  # batch, length, dim
        len_keep = int(N * (1 - mask_ratio))

        noise = torch.rand(B, N, device=x.device)  # noise in [0, 1]

        # sort noise for each sample
        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
        ids_restore = torch.argsort(ids_shuffle, dim=1) # ids_restore[i] = i-th noise element's rank

        # keep the first subset
        ids_keep = ids_shuffle[:, :len_keep]
        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D)) # x_masked are acctually non-masked elements

        # generate the binary mask: 0 is keep, 1 is remove
        mask = torch.ones([B, N], device=x.device)
        mask[:, :len_keep] = 0
        # unshuffle to get the binary mask
        mask = torch.gather(mask, dim=1, index=ids_restore)

        return x_masked, mask, ids_restore

    def forward_encoder(self, x, mask_ratio, if_mask=True):
        """
        x: [B, 1, H, W]
        """
        # embed patches
        B, C, H, W = x.shape
        x = self.patch_embed(x.reshape(B, C, -1))

        # add pos embed w/o cls token
        if self.if_cls_token:
            x = x + self.pos_embed[:, :-1, :]
        else:
            x = x + self.pos_embed

        # masking: length -> length * mask_ratio
        if if_mask:
            x, mask, ids_restore = self.random_masking(x, mask_ratio)

        # append cls token
        if self.if_cls_token:
            cls_token = self.cls_token + self.pos_embed[:, -1, :]
            cls_tokens = cls_token.expand(x.shape[0], -1, -1)
            x = torch.cat((x, cls_tokens), dim=1)
        x = self.pos_drop(x)

        # apply Mamba blocks
        attn_list = []
        for blk in self.blocks:
            if self.return_attn:
                x, attn = blk(x)
                attn_list.append(attn)
            else:
                x = blk(x)
        if if_mask:
            return x, mask, ids_restore
        else:
            # return x
            if self.return_attn:
                return x, attn_list
            else:
                return x

    def forward_decoder(self, x, ids_restore):
        # embed tokens
        x = self.decoder_embed(x)

        # append mask tokens to sequence
        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + self.num_cls_token - x.shape[1], 1)
        if self.if_cls_token:
            visible_tokens = x[:, :-1, :]
        else:
            visible_tokens = x
        x_ = torch.cat([visible_tokens, mask_tokens], dim=1)  # no cls token
        x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle
        if self.if_cls_token:
            x = torch.cat([x_, x[:, -1:, :]], dim=1)  # append cls token
        else:
            x = x_

        # add pos embed
        x = x + self.decoder_pos_embed

        # apply Mamba blocks
        for blk in self.decoder_blocks:
            x = blk(x)

        # predictor projection
        x = self.decoder_pred(x)

        # remove cls token
        if self.if_cls_token:
            x = x[:, :-1, :]
        return x

    def forward_rec_loss(self, imgs, pred, mask):
        """
        imgs: [N, 1, H, W]
        pred: [N, L, p*p*1]
        mask: [N, L], 0 is keep, 1 is remove,
        """
        target = self.stride_patchify(imgs)
        if self.norm_pix_loss:
            mean = target.mean(dim=-1, keepdim=True)
            var = target.var(dim=-1, keepdim=True)
            target = (target - mean) / (var + 1.e-6) ** .5

        loss = (pred - target) ** 2
        loss = loss.mean(dim=-1)  # [N, L], mean loss per patch

        loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
        return loss

    def forward(self, imgs, mask_ratio=0.9, **kwargs):
        # imgs: [B, 1, H, W]
        B, C, H, W = imgs.shape
        assert C == 1, "Input images should be grayscale"
        if self.is_pretrain:
            latent, mask, ids_restore = self.forward_encoder(imgs, 
                    mask_ratio=mask_ratio,)
            pred = self.forward_decoder(latent, ids_restore)
            loss = self.forward_rec_loss(imgs, pred, mask)
            return loss, pred, mask
        else:
            if self.return_attn:
                x, attn_list = self.forward_encoder(imgs, mask_ratio=mask_ratio, if_mask=False)
                if self.if_cls_token:
                    return self.head(x[:, -1, :]), attn_list
                else:
                    return self.head(torch.mean(x, dim=1)), attn_list
            else:
                x = self.forward_encoder(imgs, mask_ratio=mask_ratio, if_mask=False)
                if self.if_cls_token:
                    return self.head(x[:, -1, :])
                else:
                    return self.head(torch.mean(x, dim=1))


def net_bt_base_pretrain(**kwargs):
    model = NetTransformer(
        is_pretrain=True, embed_dim=192, depth=4,
        decoder_embed_dim=128, decoder_depth=2, **kwargs)
    return model

def net_bt_base_classifier(**kwargs):
    model = NetTransformer(
        is_pretrain=False, embed_dim=192, depth=4,
        **kwargs)
    return model

def net_bt_medium_pretrain(**kwargs):
    model = NetTransformer(
        is_pretrain=True, embed_dim=256, depth=4,
        decoder_embed_dim=128, decoder_depth=2, **kwargs)
    return model

def net_bt_meidum_classifier(**kwargs):
    model = NetTransformer(
        is_pretrain=False, embed_dim=256, depth=4,
        **kwargs)
    return model

def net_bgt_base_pretrain(**kwargs):
    model = NetTransformer(
        is_pretrain=True, embed_dim=192, depth=4,
        decoder_embed_dim=128, decoder_depth=2, block_name="basic-gated", **kwargs)
    return model

def net_bgt_base_classifier(**kwargs):
    model = NetTransformer(
        is_pretrain=False, embed_dim=192, depth=4, block_name="basic-gated",
        **kwargs)
    return model

def net_bgt_medium_pretrain(**kwargs):
    model = NetTransformer(
        is_pretrain=True, embed_dim=256, depth=4,
        decoder_embed_dim=128, decoder_depth=2, block_name="basic-gated", **kwargs)
    return model

def net_bgt_medium_classifier(**kwargs):
    model = NetTransformer(
        is_pretrain=False, embed_dim=256, depth=4, block_name="basic-gated",
        **kwargs)
    return model

def net_ft_base_pretrain(**kwargs):
    model = NetTransformer(
        is_pretrain=True, embed_dim=192, depth=4,
        decoder_embed_dim=128, decoder_depth=2, block_name="flash", **kwargs)
    return model

def net_ft_base_classifier(**kwargs):
    model = NetTransformer(
        is_pretrain=False, embed_dim=192, depth=4,
        block_name="flash", **kwargs)
    return model

def net_fgt_base_pretrain(**kwargs):
    model = NetTransformer(
        is_pretrain=True, embed_dim=192, depth=4,
        decoder_embed_dim=128, decoder_depth=2, block_name="flash-gated", **kwargs)
    return model

def net_fgt_base_classifier(**kwargs):
    model = NetTransformer(
        is_pretrain=False, embed_dim=192, depth=4,
        block_name="flash-gated", **kwargs)
    return model

def net_fgt_medium_pretrain(**kwargs):
    model = NetTransformer(
        is_pretrain=True, embed_dim=256, depth=4,
        decoder_embed_dim=128, decoder_depth=2, block_name="flash-gated", **kwargs)
    return model

def net_fgt_medium_classifier(**kwargs):
    model = NetTransformer(
        is_pretrain=False, embed_dim=256, depth=4,
        block_name="flash-gated", **kwargs)
    return model

def net_lt_base_pretrain(**kwargs):
    model = NetTransformer(
        is_pretrain=True, embed_dim=192, depth=4,
        decoder_embed_dim=128, decoder_depth=2, block_name="linear", **kwargs)
    return model

def net_lt_base_classifier(**kwargs):
    model = NetTransformer(
        is_pretrain=False, embed_dim=192, depth=4,
        block_name="linear", **kwargs)
    return model

def net_st_base_pretrain(**kwargs):
    model = NetTransformer(
        is_pretrain=True, embed_dim=192, depth=4,
        decoder_embed_dim=128, decoder_depth=2, block_name="sparse", 
        if_cls_token=False, **kwargs)
    return model

def net_st_base_classifier(**kwargs):
    model = NetTransformer(
        is_pretrain=False, embed_dim=192, depth=4,
        block_name="sparse", if_cls_token=False, **kwargs)
    return model