import math
import numpy as np
import random
import torch
import torch.nn as nn
from transformers import BartModel
import torch.nn.functional as F
from huggingface_hub import PyTorchModelHubMixin
import pickle
from transformers import BartConfig


class Embeddings(nn.Module):
    def __init__(self, n_token, d_model):
        super().__init__()
        self.lut = nn.Embedding(n_token, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)


class PianoBart(nn.Module):
    def __init__(self, bartConfig, e2w, w2e):
        super().__init__()

        self.bart = BartModel(bartConfig)
        self.hidden_size = bartConfig.d_model
        self.bartConfig = bartConfig

        # token types: 0 Measure（第几个Bar（小节））, 1 Position（Bar中的位置）, 2 Program（乐器）, 3 Pitch（音高）, 4 Duration（持续时间）, 5 Velocity（力度）, 6 TimeSig（拍号）, 7 Tempo（速度）
        self.n_tokens = []  # 每个属性的种类数
        self.classes = ['Bar', 'Position', 'Instrument', 'Pitch', 'Duration', 'Velocity', 'TimeSig', 'Tempo']
        for key in self.classes:
            self.n_tokens.append(len(e2w[key]))
        self.emb_sizes = [256] * 8
        self.e2w = e2w
        self.w2e = w2e

        # for deciding whether the current input_ids is a <PAD> token
        self.bar_pad_word = self.e2w['Bar']['Bar <PAD>']
        self.mask_word_np = np.array([self.e2w[etype]['%s <MASK>' % etype] for etype in self.classes], dtype=np.int64)
        self.pad_word_np = np.array([self.e2w[etype]['%s <PAD>' % etype] for etype in self.classes], dtype=np.int64)
        self.sos_word_np = np.array([self.e2w[etype]['%s <SOS>' % etype] for etype in self.classes], dtype=np.int64)
        self.eos_word_np = np.array([self.e2w[etype]['%s <EOS>' % etype] for etype in self.classes], dtype=np.int64)


        # word_emb: embeddings to change token ids into embeddings
        self.word_emb = []
        for i, key in enumerate(self.classes):  # 将每个特征都Embedding到256维，Embedding参数是可学习的
            self.word_emb.append(Embeddings(self.n_tokens[i], self.emb_sizes[i]))
        self.word_emb = nn.ModuleList(self.word_emb)

        # linear layer to merge embeddings from different token types
        self.encoder_linear = nn.Linear(np.sum(self.emb_sizes), bartConfig.d_model)
        self.decoder_linear = self.encoder_linear
        self.decoder_emb=None
        #self.decoder_linear= nn.Linear(np.sum(self.emb_sizes), bartConfig.d_model)

    def forward(self, input_ids_encoder, input_ids_decoder=None, encoder_attention_mask=None, decoder_attention_mask=None, output_hidden_states=True, generate=False):
        encoder_embs = []
        decoder_embs = []
        for i, key in enumerate(self.classes):
            encoder_embs.append(self.word_emb[i](input_ids_encoder[..., i]))
            if self.decoder_emb is None and input_ids_decoder is not None:
                decoder_embs.append(self.word_emb[i](input_ids_decoder[..., i]))
        if self.decoder_emb is not None and input_ids_decoder is not None:
            decoder_embs.append(self.decoder_emb(input_ids_decoder))
        encoder_embs = torch.cat([*encoder_embs], dim=-1)
        emb_linear_encoder = self.encoder_linear(encoder_embs)
        if input_ids_decoder is not None:
            decoder_embs = torch.cat([*decoder_embs], dim=-1)
            emb_linear_decoder = self.decoder_linear(decoder_embs)
        # feed to bart
        if input_ids_decoder is not None:
            y = self.bart(inputs_embeds=emb_linear_encoder, decoder_inputs_embeds=emb_linear_decoder, attention_mask=encoder_attention_mask, decoder_attention_mask=decoder_attention_mask, output_hidden_states=output_hidden_states) #attention_mask用于屏蔽<PAD> (PAD作用是在结尾补齐长度)
        else:
            y=self.bart.encoder(inputs_embeds=emb_linear_encoder,attention_mask=encoder_attention_mask)
        return y

    def get_rand_tok(self):
        rand=[0]*8
        for i in range(8):
            rand[i]=random.choice(range(self.n_tokens[i]))
        return np.array(rand)

    def change_decoder_embedding(self,new_embedding,new_linear=None):
        self.decoder_emb=new_embedding
        if new_linear is not None:
            self.decoder_linear=new_linear


class PianoBartLM(nn.Module):
    def __init__(self, pianobart: PianoBart):
        super().__init__()
        self.pianobart = pianobart
        self.mask_lm = MLM(self.pianobart.e2w, self.pianobart.n_tokens, self.pianobart.hidden_size)

    def forward(self,input_ids_encoder, input_ids_decoder=None, encoder_attention_mask=None, decoder_attention_mask=None,generate=False,device_num=-1):
        if not generate:
            x = self.pianobart(input_ids_encoder, input_ids_decoder, encoder_attention_mask, decoder_attention_mask)
            return self.mask_lm(x)
        else:
            if input_ids_encoder.shape[0] !=1:
                print("ERROR")
                exit(-1)
            if device_num==-1:
                device=torch.device('cpu')
            else:
                device=torch.device('cuda:'+str(device_num))
            pad=torch.from_numpy(self.pianobart.pad_word_np)
            input_ids_decoder=pad.repeat(input_ids_encoder.shape[0],input_ids_encoder.shape[1],1).to(device)
            result=pad.repeat(input_ids_encoder.shape[0],input_ids_encoder.shape[1],1).to(device)
            decoder_attention_mask=torch.zeros_like(encoder_attention_mask).to(device)
            input_ids_decoder[:,0,:] = torch.tensor(self.pianobart.sos_word_np)
            decoder_attention_mask[:,0] = 1
            for i in range(input_ids_encoder.shape[1]):
            # pbar = tqdm.tqdm(range(input_ids_encoder.shape[1]), disable=False)
            # for i in pbar:
                x = self.mask_lm(self.pianobart(input_ids_encoder, input_ids_decoder, encoder_attention_mask, decoder_attention_mask))
                # outputs = []
                # for j, etype in enumerate(self.pianobart.e2w):
                #     output = np.argmax(x[j].cpu().detach().numpy(), axis=-1)
                #     outputs.append(output)
                # outputs = np.stack(outputs, axis=-1)
                # outputs = torch.from_numpy(outputs)
                # outputs=self.sample(x)
                # if i!=input_ids_encoder.shape[1]-1:
                #     input_ids_decoder[:,i+1,:]=outputs[:,i,:]
                #     decoder_attention_mask[:,i+1]+=1
                # result[:,i,:]=outputs[:,i,:]
                current_output=self.sample(x,i)
                # print(current_output)
                if i!=input_ids_encoder.shape[1]-1:
                    input_ids_decoder[:,i+1,:]=current_output
                    decoder_attention_mask[:,i+1]+=1
                # 为提升速度，提前终止生成
                if (current_output>=pad).any():
                    break
                result[:,i,:]=current_output
            return result

    def sample(self,x,index): # Adaptive Sampling Policy in CP Transformer
        # token types: 0 Measure（第几个Bar（小节））, 1 Position（Bar中的位置）, 2 Program（乐器）, 3 Pitch（音高）, 4 Duration（持续时间）, 5 Velocity（力度）, 6 TimeSig（拍号）, 7 Tempo（速度）
        t=[1.2,1.2,5,1,2,5,5,1.2]
        p=[1,1,1,0.9,0.9,1,1,0.9]
        result=[]
        for j, etype in enumerate(self.pianobart.e2w):
            y=x[j]
            y=y[:,index,:]
            y=sampling(y,p[j],t[j])
            result.append(y)
        return torch.tensor(result)


# -- nucleus -- #
def nucleus(probs, p):
    probs /= (sum(probs) + 1e-5)
    sorted_probs = np.sort(probs)[::-1]
    sorted_index = np.argsort(probs)[::-1]
    cusum_sorted_probs = np.cumsum(sorted_probs)
    after_threshold = cusum_sorted_probs > p
    if sum(after_threshold) > 0:
        last_index = np.where(after_threshold)[0][0] + 1
        candi_index = sorted_index[:last_index]
    else:
        candi_index = sorted_index[0:1]
    candi_probs = [probs[i] for i in candi_index]
    candi_probs /= sum(candi_probs)
    word = np.random.choice(candi_index, size=1, p=candi_probs)[0]
    return word


def sampling(logit, p=None, t=1.0):
    logit = logit.squeeze()
    probs = torch.softmax(logit/t,dim=-1)
    probs=probs.cpu().detach().numpy()
    cur_word = nucleus(probs, p=p)
    return cur_word


class MLM(nn.Module):
    def __init__(self, e2w, n_tokens, hidden_size):
        super().__init__()
        self.proj = []
        for i, etype in enumerate(e2w):
            self.proj.append(nn.Linear(hidden_size, n_tokens[i]))
        self.proj = nn.ModuleList(self.proj)
        self.e2w = e2w

    def forward(self, y):
        y = y.last_hidden_state
        ys = []
        for i, etype in enumerate(self.e2w):
            ys.append(self.proj[i](y))
        return ys


class SelfAttention(nn.Module):
    def __init__(self, input_dim, da, r):
        '''
        Args:
            input_dim (int): batch, seq, input_dim
            da (int): number of features in hidden layer from self-attn
            r (int): number of aspects of self-attn
        '''
        super(SelfAttention, self).__init__()
        self.ws1 = nn.Linear(input_dim, da, bias=False)
        self.ws2 = nn.Linear(da, r, bias=False)

    def forward(self, h):
        attn_mat = F.softmax(self.ws2(torch.tanh(self.ws1(h))), dim=1)
        attn_mat = attn_mat.permute(0,2,1)
        return attn_mat


class SequenceClassification(nn.Module):
    def __init__(self, pianobart, class_num, hs, da=128, r=4):
        super().__init__()
        self.pianobart = pianobart
        self.attention = SelfAttention(hs, da, r)
        self.classifier = nn.Sequential(
            nn.Dropout(0.1),
            nn.Linear(hs*r, 256),
            nn.ReLU(),
            nn.Linear(256, class_num)
        )

    def forward(self, input_ids_encoder, encoder_attention_mask=None):
        # y_shift = torch.zeros_like(input_ids_encoder)
        # y_shift[:, 1:, :] = input_ids_encoder[:, :-1, :]
        # y_shift[:, 0, :] = torch.tensor(self.pianobart.sos_word_np)
        # attn_shift = torch.zeros_like(encoder_attention_mask)
        # attn_shift[:, 1:] = encoder_attention_mask[:, :-1]
        # attn_shift[:, 0] = encoder_attention_mask[:, 0]
        # x = self.pianobart(input_ids_encoder=input_ids_encoder,input_ids_decoder=y_shift,encoder_attention_mask=encoder_attention_mask,decoder_attention_mask=attn_shift)

        x = self.pianobart(input_ids_encoder=input_ids_encoder,input_ids_decoder=input_ids_encoder,encoder_attention_mask=encoder_attention_mask,decoder_attention_mask=encoder_attention_mask)

        x = x.last_hidden_state
        attn_mat = self.attention(x)
        m = torch.bmm(attn_mat, x)
        flatten = m.view(m.size()[0], -1)
        res = self.classifier(flatten)
        return res


class TokenClassification(nn.Module):
    def __init__(self, pianobart, class_num, hs):
        super().__init__()
        self.pianobart = pianobart
        self.classifier = nn.Sequential(
            nn.Dropout(0.1),
            nn.Linear(hs, 256),
            nn.ReLU(),
            nn.Linear(256, class_num)
        )

    def forward(self, input_ids_encoder, input_ids_decoder, encoder_attention_mask=None, decoder_attention_mask=None):
        x = self.pianobart(input_ids_encoder, input_ids_decoder, encoder_attention_mask, decoder_attention_mask)
        x = x.last_hidden_state
        res = self.classifier(x)
        return res


class PianoBART(
    nn.Module,
    PyTorchModelHubMixin
):
    def __init__(self, max_position_embeddings=1024, hidden_size=1024, layers=8, heads=8, ffn_dims=2048):
        super().__init__()
        with open("./Octuple.pkl", 'rb') as f:
            self.e2w, self.w2e = pickle.load(f)
        self.config = BartConfig(max_position_embeddings=max_position_embeddings,
                               d_model=hidden_size,
                               encoder_layers=layers,
                               encoder_ffn_dim=ffn_dims,
                               encoder_attention_heads=heads,
                               decoder_layers=layers,
                               decoder_ffn_dim=ffn_dims,
                               decoder_attention_heads=heads
                               )
        self.model = PianoBart(bartConfig=self.config, e2w=self.e2w, w2e=self.w2e)


    def forward(self, input_ids_encoder, input_ids_decoder=None, encoder_attention_mask=None, decoder_attention_mask=None, output_hidden_states=True, generate=False):
        return self.model(input_ids_encoder,input_ids_decoder,encoder_attention_mask,decoder_attention_mask,output_hidden_states,generate=False)