from json import decoder
from operator import mod
import numpy as np
import time,pdb,os, random, math, copy
import pandas as pd
import pickle as pkl
from sklearn.metrics import f1_score

import matplotlib.lines as mlines
from scipy.stats import norm

import torch
from torch import nn, Tensor
from torchinfo import summary
from torch.utils.data import DataLoader,Dataset
from torch.nn import TransformerEncoder, TransformerEncoderLayer
        

def find_replace(seq, token, word):
    st = 0
    while st<len(word):
        if seq==tuple(word[st:st+len(seq)]):
            word = word[:st] + [token] + word[st+len(seq):]
        st+=1
    return word

def generateVocabdict(data):
    vocab = list(data.event.unique())
    vocabdict = dict((x, tuple([idx])) for idx, x in enumerate(vocab))
    vocabdict['OVER'] = tuple([-1])
    return vocabdict

def generateWordsVocab(data, vocabdict):
    word, words = [], []
    for gkey, gdata in data.groupby(['user', 'task', 'trial']):
        events = list(gdata.event.values) + ['OVER']
        word = [vocabdict[x] for x in list(events)]
        words.append(word)
    vocab = set([])
    for word in words:
        vocab = vocab.union(set(word))
    vocab = set(x for x in list(vocab))
    return words, vocab

def generateBPE(iters, vocab, vocabdict, words, pool, savefileprefix='./', savefreq=50):
    for i in range(iters):
        stackeddata = pool.starmap(getPair, [[x] for x in words])        
        if len(stackeddata)==0:
            break
        stackedsubwords = pool.starmap(uniqueSubwords, stackeddata)
        subwords = set([])
        for idx, subword in enumerate(stackedsubwords):
            subwords = subwords.union(subword)

        paircounter = dict((x, 0) for x in subwords)
        for pairs, word in stackeddata:
            for _, pair in enumerate(pairs):
                subword = []
                for x in list(pair):
                    if len(x)==1:
                        subword.append(tuple(x))
                    else:
                        subword.append(list(x))
                subword = tuple(map(tuple,subword))
                paircounter[subword] += 1
        if len(paircounter)==0:
            break
        if max(paircounter.values())==1:
            break
        targetpair = max(paircounter, key=paircounter.get)
        prelen = len(vocab)
        vocab.add(targetpair)

        for idx, word in enumerate(words):
            stidx = []
            for st in range(len(word)-len(targetpair)+1):
                if equalTuple(targetpair, tuple(word[st:st+len(targetpair)])):
                    stidx.append(st)
            pre = 0
            updated = []
            for st in stidx:
                if pre<st:
                    updated = updated + word[pre:st] + [targetpair]
                    pre = st+len(targetpair)
            words[idx] = updated + word[pre:]

        if (i+1)%savefreq==0:
            savefilename = savefileprefix+'%d.pkl'%(i)
            print('saving to =>', savefilename)
            with open(savefilename, 'wb') as f:
                pkl.dump([vocab, vocabdict, words, paircounter],f)
    return savefilename

def sort_by_key_len(dict):
    dict_len= {key: len(key) for key in dict.keys()}
    import operator
    sorted_key_list = sorted(dict_len.items(), key=operator.itemgetter(1), reverse=True)
    sorted_dict = [{item[0]: dict[item [0]]} for item in sorted_key_list]
    return sorted_dict

def getPair(word, N=2):
    word = np.array(word)
    slid = 1
    sub_windows = (
        np.expand_dims(np.arange(N), 0) +
        np.expand_dims(np.arange(0, word.shape[0]-N+1, slid), 0).T
    ).astype(int)
    return word[sub_windows], word

def flattenTuple(x):
    flatten = []
    for oo in list(x):
        if type(oo)==int:
            flatten.append(tuple([oo]))
        else:
            if len(oo)==1:
                flatten.append(oo)
            else:
                flatten = flatten + flattenTuple(oo)
    return flatten

def equalTuple(x, y):
    if set(x)==set(y):
        if flattenTuple(x)==flattenTuple(y):
            return True
    return False

def findTuple(target, candidates):
    for x in candidates:
        if equalTuple(target, x):
            return True
    return False

def uniqueSubwords(pairs, ignore):
    subwords = set([])
    for pair in pairs:                    
        subword = []
        for x in list(pair):
            if len(x)==1:
                subword.append(tuple(x))
            else:
                subword.append(list(x))
        subword = tuple(map(tuple,subword))
        subwords.add(subword)
    return subwords

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True

def poolsegmentTokenise(gkey, gdata, win_len, vocabdict, mark, BPEvocab=None):
    unknowntoken = vocabdict['unknown']
    paddingtoken = vocabdict['padding']
    slid = win_len
    windowTokendict, windowBPEdict = {}, {}
    events = np.array(list(gdata.event.values) + ['OVER'])
    sub_windows = (
        np.expand_dims(np.arange(win_len), 0) +
        np.expand_dims(np.arange(0, len(gdata)-win_len+1, slid), 0).T
    ).astype(int)
    windowOri = events[sub_windows].tolist()
    if len(sub_windows)==0:
        lastidx = 0
    else:
        lastidx = sub_windows[-1][-1]
    if lastidx<len(gdata)-1:
        windowOri.append(events[lastidx+1:].tolist())     

    windowToken = []
    for windowO in windowOri:
        word = []
        for x in list(windowO):
            if x in vocabdict.keys():
                word.append(vocabdict[x])
            else:
                assert mark!='train'
                word.append(unknowntoken)
        windowToken.append(word)
    windowTokendict[gkey] = windowToken
    if BPEvocab is None:
        BPEvocab = [{tuple([tuple([x-1])]):x} for x in range(len(vocabdict))]

    windowBPE = []
    for word in windowToken:
        prelen = len(word)
        for vocab in BPEvocab:
            seq = list(vocab.keys())[0]
            token = list(vocab.values())[0]
            word = find_replace(seq, token, word)
        assert len(word)<=prelen
        windowBPE.append(word)
    windowBPEdict[gkey] = windowBPE
    for key,values in windowBPEdict.items():
        oo = []
        for word in values:
            newword = []
            for x in word:
                if type(x)==tuple:
                    newword.append(x[0])
                else:
                    newword.append(x)
            oo.append(newword)
        windowBPEdict[key] = oo
    segments, labels, nextaction = np.array([]), [], []
    for gkey, groupWindowBPE in windowBPEdict.items():
        for idx, windowBPE in enumerate(groupWindowBPE):
            windowBPE = windowBPE + [paddingtoken[0] for x in range(win_len-len(windowBPE))]
            assert len(windowBPE)==win_len
            if idx<len(groupWindowBPE)-1:
                nextaction.append(windowTokendict[gkey][idx+1][0][0])
            else:
                nextaction.append(vocabdict['OVER'][0])
            labels.append(gkey[1])
            if segments.shape[0]==0:
                segments = np.array([windowBPE])
            else:
                segments = np.concatenate((segments, np.array([windowBPE])), axis=0)
    return segments, labels, unknowntoken[0], paddingtoken[0]


class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, dropout: float, max_len: int = 5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
        pe = torch.zeros(1, max_len, d_model)
        pe[0, :, 0::2] = torch.sin(position * div_term)
        pe[0, :, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)
    def forward(self, x: Tensor) -> Tensor:
        x = x + self.pe[:,:x.shape[1]]
        return self.dropout(x)


class SupervisedTransformerv2(nn.Module):
    def __init__(self, **params):
        super().__init__()
        self.params = params
        self.pos_encoder = PositionalEncoding(d_model=params['d_model'], dropout=params['dropout'])
        if 'ntokens' in params.keys():
            self.embedding = nn.Embedding(num_embeddings=params['ntokens'], embedding_dim=params['d_model'])
        self.transformer_encoder = TransformerEncoder(TransformerEncoderLayer(d_model=params['d_model'], nhead=params['nhead'], 
                                                                    dim_feedforward=params['d_model']*4, 
                                                                    dropout=params['dropout'], activation='relu', batch_first=True),
                                                     params['n_layers'])
        self.linear = nn.Linear(params['d_model']*params['win_len'], params['nlabels'])

    def forward(self, encoder_input, paddingmask):
        encoder_embed = self.embedding(encoder_input) * math.sqrt(self.params['d_model'])
        encoder_pos = self.pos_encoder(encoder_embed)
        encoder_output = self.transformer_encoder(src=encoder_pos, src_key_padding_mask=paddingmask)
        output = encoder_output.view(encoder_output.shape[0],-1)
        final_output = self.linear(output)
        return final_output

def evaluate(model, criterion, val_loader, device):
    model = model.to(device)
    model.eval()
    losses, f1s = [], []
    with torch.no_grad():
        for batchdata, batchlabel, batchmask in val_loader:
            predictions = model(batchdata.to(device), batchmask.to(device))
            loss = criterion(predictions, batchlabel.reshape(-1).to(device).long())
            if np.isnan(loss.item()):
                raise "Loss NaN!"
            losses.append(loss.item())
            pred_label = np.argmax(predictions.detach().cpu().numpy(), axis=1)
            f1 = f1_score(batchlabel.numpy(), pred_label, average='macro')
            f1s.append(f1)
    return np.mean(losses), np.mean(f1s)

class DatasetPadding(Dataset):
    def __init__(self, data, paddingtoken=None, label=None):        
        self.data = data
        self.label = label
        if paddingtoken is not None:
            self.mask = data==paddingtoken
    def __getitem__(self, idx):
        if self.label is None:
            return self.data[idx], self.mask[idx]
        return self.data[idx], self.label[idx], self.mask[idx]    
    def __len__(self):
        return len(self.data)

def fit_transformer(traindata, trainlabel, testdata, testlabel, args, device, paddingtoken, nLabels, savemodel='model', nTokens=None):
    if os.path.exists(savemodel):
        return

    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
    traindata = traindata.astype(int)

    params = {'nlabels': nLabels, 'batch_size':args.batch_size, 'd_model':args.d_model,
            'optimizer_name': args, 'nhead':args.nhead, 'dropout':args.dropout,
            'win_len': traindata.shape[1], 'lr':args.lr, 'n_layers':args.n_layers, 'ntokens': nTokens}
    trainset = DatasetPadding(data=traindata, paddingtoken=paddingtoken, label=trainlabel)
    testset = DatasetPadding(data=testdata, paddingtoken=paddingtoken, label=testlabel)
    trainloader = DataLoader(trainset, batch_size=params['batch_size'], shuffle=True, num_workers=0)
    testloader = DataLoader(testset, batch_size=params['batch_size'], shuffle=True, num_workers=0)        

    model = SupervisedTransformerv2(**params).to(device)
    model.train()
    
    optimizer = getattr(torch.optim, params['optimizer_name'])(model.parameters(), 
                    lr=params['lr'], betas=(0.9,0.999), weight_decay=0.01)
    if len(trainloader)>=20:
        LOG = int(len(trainloader)/20)
    else:
        LOG = 1
    trloss, valoss, trf1, vaf1 = [], [], [], []
    evaloss, evaf1 = 0,0
    for epoch in range(1, args.epochs+1):
        for batch, (batchdata, batchlabel, batchmask) in enumerate(trainloader):
            predictions = model(batchdata.to(device), batchmask.to(device))
            loss = criterion(predictions, batchlabel.reshape(-1).to(device).long())
            if np.isnan(loss.item()):
                raise "Loss NaN!"
            loss.requires_grad_(True)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            evaloss+=loss.item()
            pred_label = np.argmax(predictions.detach().cpu().numpy(), axis=1)
            f1 = f1_score(batchlabel.numpy(), pred_label, average='macro')
            evaf1 += f1
            if batch%(LOG)==0 or batch==len(trainloader)-1:
                cur_valoss, cur_vaf1 = evaluate(model, criterion, testloader, device)
                model.train()
                trloss.append(evaloss/LOG)
                valoss.append(cur_valoss)
                trf1.append(evaf1/LOG)
                vaf1.append(cur_vaf1)
                evaloss, evaf1 = 0,0
                print('Epoch [{}/{}], Batch [{}/{}], Train Loss: {:.4f}, Train F1: {:.4f}, Val Loss: {:.4f}, Val F1: {:.4f}'
                      .format(epoch, args.epochs, batch, len(trainloader), trloss[-1], trf1[-1], valoss[-1], vaf1[-1]))
        torch.save([model.cpu(), [trloss, valoss, trf1, vaf1]], savemodel+'%d.pkl'%(epoch))