from random import random
from tkinter import Y
import torch as tch
from transformers import *
from random import *


class SentTokenizer(tch.nn.Module):
    def __init__(self) -> None:
        import spacy
        super().__init__()
        BERT_NAME = 'bert-base-uncased'
        self.tok = BertTokenizer.from_pretrained(BERT_NAME)
        self.dummy = tch.nn.parameter.Parameter(tch.tensor(0.0))
        self.en_grammar = spacy.load('en_core_web_sm')

    @staticmethod
    def unifySent(sentence):
        for tok in ['[@@)', '(@@]', '[@|)', '(|@]']:
            sentence = ''.join(sentence.split(tok))
        return sentence

    def maskNoun(self, sentence):
        doc = self.en_grammar(sentence)
        for chunk in doc.noun_chunks:
            text = chunk.text
            if random() > 0.8:
                sentence = (sentence.replace(text, '[||]', 1)
                            .replace('[||][||]', '[||]'))
        return sentence

    def getLabel(self, tokList, tokLabeledList):
        label = tch.zeros(len(tokList), dtype=tch.int8,
                          device=self.dummy.device)
        idx, offset = 0, 0
        flagMaskIsDefinition = False
        flagMaskIsApplication = False
        while idx < len(tokList):
            idx_ = idx + offset
            if idx_ + 4 > len(tokLabeledList):
                idx += 1
                continue
            flag = ''.join(tokLabeledList[idx_:idx_+4])
            offset += 4
            if flag == '[@@)':
                flagMaskIsDefinition = True
            elif flag == '(@@]':
                flagMaskIsDefinition = False
            elif flag == '[@|)':
                flagMaskIsApplication = True
            elif flag == '(|@]':
                flagMaskIsApplication = False
            else:
                offset -= 4
            if flagMaskIsApplication:
                label[idx] = 2
            elif flagMaskIsDefinition:
                label[idx] = 1
            idx += 1
        return label

    @staticmethod
    def maskTokList(tokList, tokMasked):
        off = 0
        for idx, elem in enumerate(tokList):
            idx_ = idx + off
            flag = ''.join(tokMasked[idx_: idx_+4])
            if flag != '[||]':
                continue
            if tokMasked[idx_+4] == elem:
                off += 4
                continue
            tokList[idx] = '[MASK]'
            off -= 1
        return tokList

    @staticmethod
    def randMaskConcept(tokList, label):
        for idx, elem in enumerate(tokList):
            if label[idx] != 0 and random() < 0.1:
                tokList[idx] = '[MASK]'
        return tokList

    def forward(self, sentence):
        # label[i] = concept | rely
        uniSent = self.unifySent(sentence)
        tokLabeledList = self.tok.tokenize(sentence)
        tokList = self.tok.tokenize(uniSent)

        label = self.getLabel(tokList, tokLabeledList)
        tokList = self.randMaskConcept(tokList, label)

        # tokList = ['[CLS]'] + tokList + ['[SEP]']

        tokIdList = self.tok.convert_tokens_to_ids(tokList)
        tokIdList = tch.tensor([tokIdList], device=self.dummy.device)

        # return tokIdList, tokList[1:-1], label
        return tokIdList, tokList, label


class Trapezoid(tch.nn.Module):
    def __init__(self, in_features, out_features, layers):
        super().__init__()
        self.dummy = tch.nn.parameter.Parameter(tch.tensor(0.0))
        dim_diff = out_features - in_features
        self.layer = tch.nn.Sequential(*[
            tch.nn.Sequential(
                tch.nn.Linear(
                    in_features + i * dim_diff // layers,
                    in_features + (i + 1) * dim_diff // layers
                ),
                tch.nn.LeakyReLU())
            for i in range(layers)
        ])

    def forward(self, x):
        return self.layer(x)


class LinAttention(tch.nn.Module):
    def __init__(self, in_features, out_features, attention_features):
        super().__init__()
        self.dimDown = \
            tch.nn.Linear(in_features, attention_features)
        self.matKQ = \
            tch.nn.Linear(attention_features, attention_features, bias=False)
        self.matV = \
            tch.nn.Linear(in_features, out_features)
        self.leakyRELU = \
            tch.nn.LeakyReLU()

    def forward(self, x, y):
        xp = self.dimDown(x)
        attention = tch.einsum(
            '...ij, ...kj -> ...ik',
            self.matKQ(xp), xp
        )
        attention = self.leakyRELU(attention)
        return tch.einsum(
            '...ik, ...kj -> ...ij',
            attention, self.matV(y)
        )


class KLAttention(tch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, p, x):
        # attention[i0, i1] = \sum_j p[i0, j] (\log p[i0, j] - \log p[i1, j])
        # 表示 p[i1] 丢失了多少 p[i0] 当中的信息
        EPS = 1e-40
        plog = -(p + EPS).log()
        crs_entropy = tch.einsum('...ij, ...kj -> ...ik', p, plog)
        uni_entropy = (tch.einsum('...kj, ...kj -> ...k', p, plog)
                       .unsqueeze(-1))
        # 把因为浮点数运算不准而产生的负数变成 0
        attention = (crs_entropy - uni_entropy).relu()
        return tch.einsum('...ik, ...kj -> ...ij', attention, x)


class KLTransformer(tch.nn.Module):
    def __init__(self):
        super().__init__()
        self.attention_layer = KLAttention()
        self.mlp_layer0 = Trapezoid(768, 768, 5)
        self.mlp_layer1 = Trapezoid(768, 768, 5)

    def forward(self, x):
        p = self.mlp_layer0(x).softmax(-1)
        attended_x = self.attention_layer(p, x)
        transformed_x = self.mlp_layer1(attended_x)
        return x + transformed_x


class DefDiscriminator(tch.nn.Module):
    def __init__(self):
        super().__init__()
        self.transform = KLTransformer()
        self.mlp_layer = Trapezoid(768, 3, 5)

    def forward(self, x):
        y = self.transform(x)
        y = self.mlp_layer(y).softmax(-1)
        return y


class NERModel(tch.nn.Module):
    def __init__(self):
        super().__init__()
        self.dummy = tch.nn.Parameter(tch.tensor(0.0))
        BERT_NAME = 'bert-base-uncased'
        self.bert = BertModel.from_pretrained(BERT_NAME)
        self.head = DefDiscriminator()
        self.type_cnt = tch.nn.Parameter(
            tch.tensor([1.0, 1.0, 1.0],
                       dtype=tch.double,
                       requires_grad=False)
        )

    def criterion(self, y, label, giveRate):
        ys = [y[..., label == i, i] for i in range(3)]

        def uni_criterion(t):
            eps = 1e-2 / (t.shape[-2] + 1)
            randMask = tch.rand(size=t.shape, device=self.dummy.device) < 0.8
            clip = (t < eps) * randMask
            return ~clip * (t < 1-eps) * t.log()

        loss = [uni_criterion(ys[i]) for i in range(3)]

        with tch.no_grad():
            self.type_cnt += \
                tch.tensor([(label == i).sum() for i in range(3)],
                           device=self.dummy.device)
            tok_cnt = self.type_cnt.sum()

        tot_loss = (
            - loss[0].sum() * (tok_cnt / self.type_cnt[0]).to(float)
            - loss[1].sum() * (tok_cnt / self.type_cnt[1]).to(float)
            - loss[2].sum() * (tok_cnt / self.type_cnt[2]).to(float)
        )

        if giveRate:
            cnt = (ys[1] > 1/3).sum().item() + (ys[2] > 1/3).sum().item()
            label_cnt = (label > 0).sum().item()
            tot_rate = (cnt + (ys[0] > 1/3).sum().item()) / label.shape[-1]
            if label_cnt > 0:
                rate = cnt / label_cnt
            else:
                rate = -1
            return tot_loss, rate, tot_rate
        else:
            return tot_loss

    def forward(self, x, label=None, giveRate=True):
        with tch.no_grad():
            y = self.bert(x)[0]
        y = self.head(y)
        if label is None:
            return y
        else:
            return self.criterion(y, label, giveRate)


if __name__ == '__main__':
    import pandas as pd
    from tqdm import tqdm
    import matplotlib.pyplot as plt

    df = pd.read_csv('train.csv', sep='\n')['0']

    device = tch.device('cuda:0')

    # 模型
    tokenizer = SentTokenizer().to(device)
    model = NERModel().to(device)

    try:
        with open('NER.model', 'rb') as f:
            print('find model, load state dict')
            model.load_state_dict(tch.load(f))
            print('load model state dict success')
    except:
        pass

    # 训练配置
    optimizer = tch.optim.RMSprop(model.head.parameters(), lr=1e-5)
    BATCH_SIZE = 15
    SAVE_ONCE = 5000

    try:
        with open('NER.optimizer', 'rb') as f:
            print('find optimizer, load state dict')
            optimizer.load_state_dict(tch.load(f))
            print('load optimizer state dict success')
    except:
        pass

    # 可视化
    running_loss = 0.0
    running_rate = 0.0
    running_tot_rate = 0.0
    history_loss = []
    history_rate = []
    history_tot_rate = []
    skippedIter = 0

    plt.ion()

    for epoch in range(1, 5):
        dataset_with_progress_bar = tqdm(
            enumerate(df.sample(frac=1)), total=len(df))
        skippedIter = 0
        for i, sentence in dataset_with_progress_bar:
            tokIdList, _, label = tokenizer(sentence)
            if tokIdList.shape[-1] > 512 or\
                    (int((label == 1).sum()) == 0):
                skippedIter += 1
            else:
                loss, rate, tot_rate = model(tokIdList, label)

                running_loss = \
                    9e-1*running_loss + 1e-1*loss.item() if running_loss > 0.0 \
                    else loss.item()
                running_rate = \
                    99e-2*running_rate + 1e-2*rate if rate > 0.0 \
                    else running_rate
                running_tot_rate = \
                    99e-2*running_tot_rate + 1e-2*tot_rate

                dataset_with_progress_bar\
                    .set_description(
                        'loss[%-1.5f] rate[%-2.2f%%] '
                        'tot_rate[%-2.2f%%] sent_len[%-3d] '
                        'skipped[%d] epoch[%d] '
                        % (running_loss, running_rate * 100,
                        running_tot_rate * 100, tokIdList.shape[-1],
                        skippedIter, epoch)
                    )

                loss.backward()

                if (i - skippedIter) % BATCH_SIZE == 0:
                    optimizer.step()
                    optimizer.zero_grad()
                    # 记录训练历史
                    history_loss.append(running_loss)
                    history_rate.append(running_rate)
                    history_tot_rate.append(running_tot_rate)
                    if len(history_loss) > 20:
                        history_loss = history_loss[-20:]
                    if len(history_rate) > 20:
                        history_rate = history_rate[-20:]
                    if len(history_tot_rate) > 20:
                        history_tot_rate = history_tot_rate[-20:]
                    # 画图
                    plt.clf()
                    plt.subplot(1, 2, 1)
                    plt.plot(range(len(history_loss)),
                            history_loss, c='red',
                            label='loss (cross entropy loss)')
                    plt.legend()
                    plt.subplot(1, 2, 2)
                    plt.plot(range(len(history_rate)),
                            history_rate, c='blue', label='rate (only concepts)')
                    plt.plot(range(len(history_tot_rate)),
                            history_tot_rate, c='green', label='rate (all)')
                    plt.legend()
                    plt.draw()
                    plt.pause(0.01)
                    BATCH_SIZE = randint(15, 25)

            if i % SAVE_ONCE == 0:
                # 保存模型
                dataset_with_progress_bar\
                    .set_description('saving ! ')
                tch.save(optimizer.state_dict(), 'NER.optimizer')
                tch.save(model.state_dict(), 'NER.model')
                dataset_with_progress_bar\
                    .set_description('done ! ')

        dataset_with_progress_bar\
            .set_description('saving ! ')
        tch.save(optimizer.state_dict(), 'NER.optimizer')
        tch.save(model.state_dict(), 'NER.model')
        dataset_with_progress_bar\
            .set_description('done ! ')
