import math
import torch
from torch import nn
import torch.nn.functional as F

from image_synthesis.utils.misc import instantiate_from_config
from image_synthesis.modeling.utils.misc import gen_attention_mask
from image_synthesis.modeling.transformers.base_transformer import BaseTransformer
try:
    from image_synthesis.modeling.modules.sparse_matmul.sparse_matmul import SparseMatmul
except:
    print("Sparse Matmul compiled Error! Using torch matmul instead!")
    SparseMatmul = None

from image_synthesis.modeling.utils.misc import logits_top_k

# NOTE: We found that our implemented sparse mat mul is slower than torch.einsum or operation @ :(
SparseMatmul = None

class SparseAttention(nn.Module):
    """
    A multi-head masked sparse self-attention layer with a projection at the end.
    It is possible to use torch.nn.MultiheadAttention here but I am including an
    explicit implementation here to show that there is nothing too scary here.
    """
    def __init__(self,
                 n_embd, # the embed dim
                 n_head, # the number of heads
                 seq_len=None, # the max length of sequence
                 attn_pdrop=0.1, # attention dropout prob
                 resid_pdrop=0.1, # residual attention dropout prob
                 causal=True,
                 attn_type='full',
                 content_spatial_size=None, # H , W
                 conv_attn_kernel_size=None, # only need for dalle_conv attention
    ):
        super().__init__()
        assert n_embd % n_head == 0
        # key, query, value projections for all heads
        self.key = nn.Linear(n_embd, n_embd)
        self.query = nn.Linear(n_embd, n_embd)
        self.value = nn.Linear(n_embd, n_embd)
        # regularization
        self.attn_drop = nn.Dropout(attn_pdrop)
        self.resid_drop = nn.Dropout(resid_pdrop)
        # output projection
        self.proj = nn.Linear(n_embd, n_embd)

        self.n_head = n_head
        self.causal = causal

        # sparse mask to ensure that attention is only applied to the left in the input sequence
        self.content_spatial_size = content_spatial_size
        self.attn_type = attn_type
        self.conv_attn_kernel_size = conv_attn_kernel_size
        content_seq_len = content_spatial_size[0] * content_spatial_size[1]
        condition_seq_len = seq_len - content_seq_len
        assert condition_seq_len >= 0, 'condition seq should be larger or equal to 0!'
        assert attn_type in ['full', 'dalle_row', 'dalle_col', 'dalle_conv']
        # if self.causal or self.attn_type != 'full':
        mask = gen_attention_mask(
            H=self.content_spatial_size[0],
            W=self.content_spatial_size[1],
            type=self.attn_type,
            causal=self.causal,
            condition_seq_len=condition_seq_len,
            kernel_size=self.conv_attn_kernel_size
        )
        self.register_buffer('mask', mask.view(1, 1, seq_len, seq_len))
        self.sparse_matmul = SparseMatmul() if SparseMatmul is not None else None

    def forward(self, x, mask=None):
        """
        x: B x T x C
        mask: None or tensor B x T, bool type. For values with False, no attention should be attened
        """
        B, T, C = x.size()
        hs = C // self.n_head
        # import pdb; pdb.set_trace()
        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        k = self.key(x).view(B, T, self.n_head, hs).transpose(1, 2) # (B, nh, T, hs)
        q = self.query(x).view(B, T, self.n_head, hs).transpose(1, 2) # (B, nh, T, hs)
        v = self.value(x).view(B, T, self.n_head, hs).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        # print(q.shape, k.shape)
        # import pdb; pdb.set_trace()
        if self.sparse_matmul is not None:
            att = self.sparse_matmul(
                q.contiguous().view(B*self.n_head, T, hs),
                k.contiguous().view(B*self.n_head, T, hs),
                mask = self.mask[0,0,:T,:T], # 2D mask is ok
            ).view(B, self.n_head, T, T) # B x nh x T x hs
            att = att * (1.0 / math.sqrt(k.size(-1)))
        else:
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) # (B, nh, T, T)
        att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))
        
        if mask is not None:
            mask = mask.view(B, 1, 1, T)
            att.masked_fill(~mask, float('-inf'))

        att = F.softmax(att, dim=-1) # (B, nh, T, T)
        att = self.attn_drop(att)
        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side, (B, T, C)
        att = att.mean(dim=1, keepdim=False) # (B, T, T)

        # output projection
        y = self.resid_drop(self.proj(y))
        return y, att   


class FullAttention(nn.Module):
    """
    A vanilla multi-head masked self-attention layer with a projection at the end.
    It is possible to use torch.nn.MultiheadAttention here but I am including an
    explicit implementation here to show that there is nothing too scary here.
    """

    def __init__(self,
                 n_embd, # the embed dim
                 n_head, # the number of heads
                 seq_len=None, # the max length of sequence
                 attn_pdrop=0.1, # attention dropout prob
                 resid_pdrop=0.1, # residual attention dropout prob
                 causal=True,
    ):
        super().__init__()
        assert n_embd % n_head == 0
        # key, query, value projections for all heads
        self.key = nn.Linear(n_embd, n_embd)
        self.query = nn.Linear(n_embd, n_embd)
        self.value = nn.Linear(n_embd, n_embd)
        # regularization
        self.attn_drop = nn.Dropout(attn_pdrop)
        self.resid_drop = nn.Dropout(resid_pdrop)
        # output projection
        self.proj = nn.Linear(n_embd, n_embd)

        self.n_head = n_head
        self.causal = causal



    def forward(self, x, encoder_output, mask=None):
        """
        x: B x T x C
        mask: None or tensor B x T, bool type. For values with False, no attention should be attened
        """
        B, T, C = x.size()
        # import pdb; pdb.set_trace()
        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        
        # print(q.shape, k.shape)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) # (B, nh, T, T)

        att = F.softmax(att, dim=-1) # (B, nh, T, T)
        att = self.attn_drop(att)
        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side, (B, T, C)
        att = att.mean(dim=1, keepdim=False) # (B, T, T)

        # output projection
        y = self.resid_drop(self.proj(y))
        return y, att

class CrossAttention(nn.Module):
    """
    A vanilla multi-head masked self-attention layer with a projection at the end.
    It is possible to use torch.nn.MultiheadAttention here but I am including an
    explicit implementation here to show that there is nothing too scary here.
    """

    def __init__(self,
                 condition_seq_len,
                 n_embd, # the embed dim
                 n_head, # the number of heads
                 seq_len=None, # the max length of sequence
                 attn_pdrop=0.1, # attention dropout prob
                 resid_pdrop=0.1, # residual attention dropout prob
                 causal=True,
    ):
        super().__init__()
        assert n_embd % n_head == 0
        # key, query, value projections for all heads
        self.key = nn.Linear(n_embd, n_embd)
        self.query = nn.Linear(n_embd, n_embd)
        self.value = nn.Linear(n_embd, n_embd)
        # regularization
        self.attn_drop = nn.Dropout(attn_pdrop)
        self.resid_drop = nn.Dropout(resid_pdrop)
        # output projection
        self.proj = nn.Linear(n_embd, n_embd)

        self.n_head = n_head
        self.causal = causal

        # causal mask to ensure that attention is only applied to the left in the input sequence
        if self.causal:
            self.register_buffer("mask", torch.tril(torch.ones(seq_len, seq_len))
                                        .view(1, 1, seq_len, seq_len))


    def forward(self, x, encoder_output, mask=None):
        """
        x: B x T x C
        encoder_output: B x T x C
        mask: None or tensor B x T, bool type. For values with False, no attention should be attened
        """
        B, T, C = x.size()
        B, T_E, C = encoder_output.size()
        # import pdb; pdb.set_trace()
        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        k = self.key(encoder_output).view(B, T_E, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = self.value(encoder_output).view(B, T_E, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        
        # print(q.shape, k.shape)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) # (B, nh, T, T)

        # if self.causal:
            # print(att.shape, self.mask.shape, T)
        #     att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))
        
        # if mask is not None:
        #     mask = mask.view(B, 1, 1, T)
        #     att.masked_fill(~mask, float('-inf'))

        att = F.softmax(att, dim=-1) # (B, nh, T, T)
        att = self.attn_drop(att)
        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side, (B, T, C)
        att = att.mean(dim=1, keepdim=False) # (B, T, T)

        # output projection
        y = self.resid_drop(self.proj(y))
        return y, att

class GELU2(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return x * F.sigmoid(1.702 * x)

class Block(nn.Module):
    """ an unassuming Transformer block """

    def __init__(self,
                 condition_seq_len,
                 n_embd,
                 n_head,
                 seq_len,
                 attn_pdrop=0.1,
                 resid_pdrop=0.1,
                 causal=True,
                 mlp_hidden_times=4,
                 activate='GELU',
                 attn_type='full',
                 if_upsample=False,
                 upsample_type='bilinear',
                 content_spatial_size=None, # H , W
                 conv_attn_kernel_size=None, # only need for dalle_conv attention
                 ):
        super().__init__()
        self.if_upsample = if_upsample
        if self.if_upsample == True:
            self.upsample = Upsample(scale=2, upsample_type=upsample_type, dim=n_embd)

        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)
        if attn_type == 'full':
            self.attn = FullAttention(
                n_embd=n_embd,
                n_head=n_head,
                seq_len=seq_len,
                attn_pdrop=attn_pdrop,
                resid_pdrop=resid_pdrop,
                causal=causal
            )
        elif attn_type == 'cross':
            self.attn = CrossAttention(
                condition_seq_len,
                n_embd=n_embd,
                n_head=n_head,
                seq_len=seq_len,
                attn_pdrop=attn_pdrop,
                resid_pdrop=resid_pdrop,
                causal=causal
            )
        elif attn_type == 'selfcross':
            self.attn1 = FullAttention(
                    n_embd=n_embd,
                    n_head=n_head,
                    seq_len=seq_len,
                    attn_pdrop=attn_pdrop, 
                    resid_pdrop=resid_pdrop,
                    causal=causal
                    )
            self.attn2 = CrossAttention(
                    condition_seq_len,
                    n_embd=n_embd,
                    n_head=n_head,
                    seq_len=seq_len,
                    attn_pdrop=attn_pdrop,
                    resid_pdrop=resid_pdrop,
                    causal=causal
                    )
            self.if_selfcross = True
            self.ln1_1 = nn.LayerNorm(n_embd)
        else:
            self.attn = SparseAttention(
                n_embd=n_embd,
                n_head=n_head,
                seq_len=seq_len,
                attn_pdrop=attn_pdrop,
                resid_pdrop=resid_pdrop,
                causal=causal,
                attn_type=attn_type,
                content_spatial_size=content_spatial_size, # H , W
                conv_attn_kernel_size=conv_attn_kernel_size, # only need for dalle_conv attention
            )
        assert activate in ['GELU', 'GELU2']
        act = nn.GELU() if activate == 'GELU' else GELU2()
        self.mlp = nn.Sequential(
            nn.Linear(n_embd, mlp_hidden_times * n_embd),
            act,
            nn.Linear(mlp_hidden_times * n_embd, n_embd),
            nn.Dropout(resid_pdrop),
        )

    def forward(self, x, encoder_output, mask=None):    
        if self.if_upsample == True:
            x = self.upsample(x)
        if self.if_selfcross == False:
            a, att = self.attn(self.ln1(x), encoder_output, mask=mask)
            x = x + a 
        else:
            a, att = self.attn1(self.ln1(x), encoder_output, mask=mask)
            x = x + a
            a, att = self.attn2(self.ln1_1(x), encoder_output, mask=mask)
            x = x + a
        x = x + self.mlp(self.ln2(x))

        return x, att

class Upsample(nn.Module):
    def __init__(self, scale, upsample_type, dim):
        super().__init__()
        assert scale == 2
        if upsample_type == 'pixel_shuffle':
            self.upsample = nn.PixelShuffle(2)
            self.reduction = nn.Linear(dim, dim*4, bias=False)
        elif upsample_type == 'bilinear':
            self.upsample = nn.Upsample(scale_factor=2, mode='bilinear')
        else:
            print("upsample_type error")
        self.dim = dim
        self.upsample_type = upsample_type

    def forward(self, x):
        """
        x: B, H*W, C
        """
        B, L, C = x.shape
        H = int(math.sqrt(L))
        W = int(math.sqrt(L))
        assert L == H * W, "input feature has wrong size"
        assert C == self.dim, "wrong in PatchMerging"

        # x = x.view(B, H, W, C)
        # x = x.permute(0, 3, 1, 2).contiguous()   # B,C,H,W
        # x = self.upsample(x)
        if self.upsample_type == 'pixel_shuffle':
            x = self.reduction(x)
            x = x.view(B, H, W, C*4)
            x = x.permute(0, 3, 1, 2).contiguous()   # B,C,H,W
            x = self.upsample(x)
            x = x.permute(0, 2, 3, 1).contiguous().view(B, L*4, C)
            # x = x.permute(0, 2, 3, 1).contiguous().view(B, L*4, C//4)
            # x = self.reduction(x)
        elif self.upsample_type == 'bilinear':
            x = x.view(B, H, W, C)
            x = x.permute(0, 3, 1, 2).contiguous()    # B,C,H,W
            x = self.upsample(x) 
            x = x.permute(0, 2, 3, 1).contiguous().view(B, L*4, C)
        else:
            print("upsample_type error")

        return x






class NAR2Transformer(BaseTransformer):
    def __init__(
        self,
        *,
        n_layer, # number of layers in transformer
        scale_n_layer,
        condition_seq_len, # length of conditoin sequences
        condition_num_embed = 16640, # number of condition embedding, 16640 for dalle
        content_seq_len, # length of content sequences
        embd_pdrop=0., # embedding dropout prob

        n_embd, # the embed dim
        n_head, # the number of heads
        attn_pdrop=0.1, # attention dropout prob
        resid_pdrop=0.1, # residual attention dropout prob
        causal=True,
        block_activate='GELU',
        mlp_hidden_times=4, # the times of hidden dimension in the MLP of attetntion block

        attn_with_mask = False, # perform attention with mask
        attn_condition_with_mask=False,
        attn_content_with_mask=False,
        attn_type='full',
        content_spatial_size=None, # H , W
        conv_attn_kernel_size=None, # only need for dalle_conv attention

        predict_condition=False, # prediction with condition tokens 
        condition_loss_weight=1/8.0,
        content_loss_weight=7/8.0,
        content_emb_config=None,
        condition_emb_config=None,

        condition_ignore_token=-100,
        content_ignore_token=-100,

        input_type='text',    # noise or text or pos_emb
        upsample_type='bilinear' # pixel_reshuffle

    ):
        super().__init__()

        self.input_type = input_type
        self.condition_num_embed = condition_num_embed
        # embeddings for condition and content
        self.content_emb = instantiate_from_config(content_emb_config)
        if condition_emb_config is None:
            # share the condition embed with content embed
            self.condition_emb = None
            assert not predict_condition, 'If want to predict condition token, please provide condition embed config'
        else:
            # for condition and config, we learn a seperate embedding
            self.condition_emb = instantiate_from_config(condition_emb_config)
            assert self.condition_emb.embed_dim == self.content_emb.embed_dim
            self.condition_num_embed == self.condition_emb.num_embed

        self.condition_ignore_token = condition_ignore_token
        self.content_ignore_token = content_ignore_token

        # drop for embedding
        if embd_pdrop > 0:
            self.drop = nn.Dropout(embd_pdrop)
        else:
            self.drop = None
        
        # transformer
        if isinstance(attn_type, str):
            if attn_type == 'full' or 'selfcross':
                all_attn_type = [attn_type] * n_layer
            elif attn_type == 'cross':
                all_attn_type = []
                for index in range(0,n_layer):
                    if index%2==0:
                        all_attn_type.append('full')
                    else:
                        all_attn_type.append('cross')
            else:
                raise NotImplementedError
                
        elif isinstance(attn_type, dict):
            start_attn_type = attn_type.get('start', []) # list
            end_attn_type = attn_type.get('end', []) # list
            middle_attn_type = attn_type['middle'] # list
            middle_layers = n_layer - len(start_attn_type) - len(end_attn_type)
            middle_attn_type = middle_attn_type * int(middle_layers / len(middle_attn_type))
            all_attn_type = start_attn_type + middle_attn_type + end_attn_type
            assert len(all_attn_type) == n_layer, 'number of attn type not equal to number of layers!'
        else:
            raise NotImplementedError
        
        if content_spatial_size is None:
            s = int(math.sqrt(content_seq_len))
            assert s * s == content_seq_len
            content_spatial_size = (s, s)

        # self.input_proj = nn.Linear(512, 1024*4)

        blocks = []
        scale_number = 4  # log2(32)-log2(4)+1
        # scale_n_layer = [7,7,7,7]
        assert sum(scale_n_layer) == n_layer
        assert scale_number == len(scale_n_layer)
        seq_len = 16
        block_count = 0
        for i in range(scale_number):
            if i == 0:
                if_upsample = False
            else:
                if_upsample = True
                seq_len = seq_len * 4
            for j in range(scale_n_layer[i]):
                blocks.append(Block(
                    condition_seq_len,
                    n_embd=n_embd,
                    n_head=n_head,
                    seq_len=seq_len,
                    attn_pdrop=attn_pdrop,
                    resid_pdrop=resid_pdrop,
                    causal=causal,
                    mlp_hidden_times=mlp_hidden_times,
                    activate=block_activate,
                    attn_type=all_attn_type[block_count],
                    if_upsample = if_upsample,
                    upsample_type = upsample_type,
                    ))
                block_count += 1
                if_upsample = False


        self.blocks = nn.Sequential(*blocks)
                
        out_cls = self.content_emb.num_embed

        self.to_logits = nn.Sequential(
            nn.LayerNorm(n_embd),
            nn.Linear(n_embd, out_cls),
        )
        
        self.causal = causal
        self.condition_seq_len = condition_seq_len
        self.content_seq_len = content_seq_len
        self.condition_loss_weight = condition_loss_weight
        self.content_loss_weight = content_loss_weight
        self.predict_condition = predict_condition
        self.attn_with_mask = attn_with_mask
        self.attn_condition_with_mask = attn_condition_with_mask
        self.attn_content_with_mask = attn_content_with_mask

        if self.attn_with_mask:
            print('Warning: attn_with_mask is suppressed!')
            self.attn_condition_with_mask = True
            self.attn_content_with_mask = True

        self.apply(self._init_weights)

        if condition_emb_config['target'] == 'image_synthesis.modeling.embeddings.clip_text_embedding.CLIPTextEmbedding':
            self.condition_emb = instantiate_from_config(condition_emb_config)
            assert self.condition_emb.embed_dim == self.content_emb.embed_dim

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    @property
    def device(self):
        return self.to_logits[-1].weight.device

    def parameters(self, recurse=True, name=None):
        """
        Following minGPT:
        This long function is unfortunately doing something very simple and is being very defensive:
        We are separating out all parameters of the model into two buckets: those that will experience
        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).
        We are then returning the PyTorch optimizer object.
        """
        # return super().parameters(recurse=True)
        if name is None or name == 'none':
            return super().parameters(recurse=recurse)
        else:
            # separate out all parameters to those that will and won't experience regularizing weight decay
            print("GPTLikeTransformer: get parameters by the overwrite method!")
            decay = set()
            no_decay = set()
            whitelist_weight_modules = (torch.nn.Linear, )
            blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
            for mn, m in self.named_modules():
                for pn, p in m.named_parameters():
                    fpn = '%s.%s' % (mn, pn) if mn else pn # full param name

                    if pn.endswith('bias'):
                        # all biases will not be decayed
                        no_decay.add(fpn)
                    elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
                        # weights of whitelist modules will be weight decayed
                        decay.add(fpn)
                    elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
                        # weights of blacklist modules will NOT be weight decayed
                        no_decay.add(fpn)
            # special case the position embedding parameter as not decayed
            module_name = ['condition_emb', 'content_emb']
            pos_emb_name = ['pos_emb', 'width_emb', 'height_emb', 'pad_emb', 'token_type_emb']
            for mn in module_name:
                if hasattr(self, mn) and getattr(self, mn) is not None:
                    for pn in pos_emb_name:
                        if hasattr(getattr(self, mn), pn):
                            if isinstance(getattr(getattr(self, mn), pn), torch.nn.Parameter):
                                no_decay.add('{}.{}'.format(mn, pn))

            # validate that we considered every parameter
            param_dict = {pn: p for pn, p in self.transformer.named_parameters()}# if p.requires_grad} 
            inter_params = decay & no_decay
            union_params = decay | no_decay
            assert len(inter_params) == 0, "parameters %s made it into both decay/no_decay sets!" % (str(inter_params), )
            assert len(param_dict.keys() - union_params) == 0, "parameters %s were not separated into either decay/no_decay set!" \
                                                        % (str(param_dict.keys() - union_params), )

            # create the pytorch optimizer object
            optim_groups = [
                {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": 0.01},
                {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
            ]
            return optim_groups

    def forward(
            self, 
            input, 
            return_loss=False, 
            return_logits=True, 
            return_att_weight=False,
            **kwargs):
        if self.causal: # for causal transformer, the condition token must not be None
            batch_size = input['condition_token'].shape[0]
        else: # for causal transformer, the content token must not be None
            batch_size = input['content_token'].shape[0]
        
        # 1) get embeddding for condition and content
        
        assert input.get('condition_token', None) != None  
        cont_mask = input['content_mask'].view(batch_size, -1) if 'content_mask' in input \
                    else torch.ones(batch_size, input['content_token'].shape[1]).to(self.device).bool()
        cont_emb = self.content_emb(input['content_token'], mask=cont_mask, token_type=input.get('content_token_type', None))

        if input.get('condition_mask', None) is not None:
            cond_mask = input['condition_mask'].view(batch_size, -1)
        else:
            cond_mask = torch.ones(batch_size, input['condition_token'].shape[1]).to(self.device).bool()
            
        if self.condition_emb is not None:  # do this
            if self.condition_emb.additional_last_embedding == True:
                with torch.no_grad():
                    cond_emb, last_embedding = self.condition_emb(input['condition_token'], mask=cond_mask, token_type=input.get('condition_token_type', None)) # B x Ld x D   #256*1024
                    last_embedding = last_embedding.type_as(cont_emb)
            else:
                with torch.no_grad():
                    cond_emb = self.condition_emb(input['condition_token'], mask=cond_mask, token_type=input.get('condition_token_type', None)) # B x Ld x D   #256*1024

            cond_emb = cond_emb.type_as(cont_emb)
            
        else: # share condition embeding with content
            cond_emb = input['condition_embed_token'].to(cont_emb.dtype)
            
  
        ############# prepare noise input or text input here ####################
        if self.input_type == 'noise':
            noise_input = torch.randn(batch_size, 512).cuda()
        elif self.input_type == 'text':
            noise_input = last_embedding.detach()
        elif self.input_type == '16text':
            noise_input = cond_emb[:,0:16,:].type_as(cont_emb).detach()
        else:
            print('Not implemented')

        # emb = self.input_proj(noise_input).reshape(batch_size,4,1024)
        emb = noise_input

        # 2) forward in transformer
        # emb = self.blocks(noise_input, cond_emb)
        for block_idx in range(len(self.blocks)):
            emb, att_weight = self.blocks[block_idx](emb, cond_emb)


        # 3) get logits
        logits = self.to_logits(emb) # B x (Ld+Lt) x n


        # 4) get output, especially loss
        out = {}
        if return_logits:
            out['logits'] = logits
        if return_att_weight:
            out['attention_weight'] = att_weight

        if return_loss:
            logtis = logits.contiguous()
            target = input['content_token'] # B x Lt
            loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), target.view(-1), ignore_index=self.content_ignore_token)

            out['loss'] = loss 
        return out

    def sample(
            self, 
            condition_token, 
            condition_mask, 
            condition_embed,
            content_token=None,
            filter_ratio = 0.5,
            temperature = 1.0,
            return_att_weight=False,
            return_logits=False,
            **kwargs):
        this_input = {'condition_token': condition_token,
                'content_token': content_token, 
                'condition_mask': condition_mask,
                'condition_embed_token': condition_embed,
                }
        with torch.no_grad():
            trans_out = self.forward(this_input, return_loss=False, return_logits=True, return_att_weight=return_att_weight)
            logits = trans_out['logits']
            # sample 
            batch_size = logits.shape[0]
            content_token = torch.zeros((batch_size,0)).to(condition_token)
            for index in range(logits.size()[1]):
                sample = self._get_sample_from_logits(logits[:,index:index+1,:], filter_ratio=filter_ratio, temperature=temperature)  # B x 1
                content_token = torch.cat((content_token, sample), dim=1)

        output = {'content_token': content_token}
        if return_logits:
            output['logits'] = logits
        return output










