# -*- coding: utf-8 -*-
# Created by li huayong on 2019/9/28
import os
import re
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset
from tqdm import tqdm, trange
from abc import ABCMeta, abstractmethod
from typing import Dict, List, Tuple
from utils.data.conll_file import CoNLLFile
from utils.data.graph_vocab import GraphVocab
from utils.model.get_optimizer import get_optimizer
from utils.model.parser_funs import sdp_decoder, parse_semgraph
import utils.model.sdp_simple_scorer as sdp_scorer
from utils.best_result import BestResult
from utils.model.label_smoothing import label_smoothed_kl_div_loss
from PyToolkit.PyToolkit import get_logger
from PyToolkit.PyToolkit.seed import set_seed

try:
    from torch.utils.tensorboard import SummaryWriter
except ImportError:
    from tensorboardX import SummaryWriter


class BaseDependencyTrainer(metaclass=ABCMeta):
    def __init__(self, args, model):
        self.model = model
        self.optimizer = self.optim_scheduler = None
        self.graph_vocab = GraphVocab(args.graph_vocab_file)
        self.configs = args
        self.logger = get_logger(args.log_name)

    @abstractmethod
    def _unpack_batch(self, batch: TensorDataset) -> Dict:
        """
        拆分batch，得到encoder的输入和word mask，sentence length，以及dep ids,以及其他输入信息
        eg:
            dataset = TensorDataset(all_input_ids, all_input_mask,
                        all_segment_ids, all_start_pos,
                        all_end_pos, all_dep_ids,
                        all_pos_ids)

        Args:
            batch: 输入的单个batch,类型为TensorDataset(或者torchtext.dataset)，可用索引分别取值

        Returns:
            返回一个字典，[1]是inputs，类型为字典；[2]是word mask；[3]是sentence length,python 列表；[4]是dep ids,
            根据实际情况可能还包含其他输入信息
        """
        raise NotImplementedError('must implement in sub class')

    def _custom_train_operations(self, epoch: int):
        """
            某些模型在训练时可能需要一些定制化的操作，
            比如BERT类型的模型可能会在Training的时候动态freeze某些层
            为了支持这些操作同时不破坏BiaffineDependencyTrainer的普适性，我们加入这个方法
            BiaffineDependencyTrainer的子类可以选择重写该方法以支持定制化操作
            注意这个方法会在训练的每个epoch的开始调用一次
            本方法默认不会做任何事情
        :return:
        """
        pass

    def _update_and_predict(self, unlabeled_scores, labeled_scores, unlabeled_target, labeled_target, word_pad_mask,
                            label_loss_ratio=0.5, sentence_lengths=None,
                            calc_loss=True, update=True, calc_prediction=False,
                            pos_logits=None, pos_target=None, pos_loss_ratio=1.0,
                            summary_writer=None, global_step=None):
        """
            针对一个batch输入：计算loss，反向传播，计算预测结果
            :param word_pad_mask: 以word为单位，1为PAD，0为真实输入
        :return:
        """
        weights = torch.ones(word_pad_mask.size(0), self.configs.max_seq_len, self.configs.max_seq_len,
                             dtype=unlabeled_scores.dtype,
                             device=unlabeled_scores.device)
        # 将PAD的位置权重设为0，其余位置为1
        weights = weights.masked_fill(word_pad_mask.unsqueeze(1), 0)
        weights = weights.masked_fill(word_pad_mask.unsqueeze(2), 0)
        # words_num 记录batch中的单词数量
        # torch.eq(word_pad_mask, False) 得到word_mask
        words_num = torch.sum(torch.eq(word_pad_mask, False)).item()
        if calc_loss:
            assert label_loss_ratio
            assert unlabeled_target is not None and labeled_target is not None
            dep_arc_loss_func = nn.BCEWithLogitsLoss(weight=weights, reduction='sum')
            dep_arc_loss = dep_arc_loss_func(unlabeled_scores, unlabeled_target)

            dep_label_loss_func = nn.CrossEntropyLoss(ignore_index=-1, reduction='sum')
            dependency_mask = labeled_target.eq(0)
            labeled_target = labeled_target.masked_fill(dependency_mask, -1)
            labeled_scores = labeled_scores.contiguous().view(-1, len(self.graph_vocab.get_labels()))
            dep_label_loss = dep_label_loss_func(labeled_scores, labeled_target.view(-1))

            if self.configs.use_pos:
                assert pos_logits is not None
                assert pos_target is not None
                pos_loss_func = nn.CrossEntropyLoss(ignore_index=self.configs.pos_label_pad_idx)
                pos_loss = pos_loss_func(pos_logits.view(-1, self.configs.pos_label_num), pos_target.view(-1))

            loss = 2 * ((1 - label_loss_ratio) * dep_arc_loss + label_loss_ratio * dep_label_loss)

            if self.configs.use_pos:
                loss = loss + pos_loss_ratio * pos_loss

            if self.configs.average_loss_by_words_num:
                loss = loss / words_num

            if self.configs.scale_loss:
                loss = loss * self.configs.loss_scaling_ratio

            if self.configs.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training

            if summary_writer and global_step:
                summary_writer.add_scalar('train_loss/dep_arc_loss', dep_arc_loss, global_step)
                summary_writer.add_scalar('train_loss/dep_label_loss', dep_label_loss, global_step)
                if self.configs.use_pos:
                    summary_writer.add_scalar('train_loss/pos_loss', pos_loss, global_step)

            if update:
                # Ref:https://discuss.pytorch.org/t/model-zero-grad-or-optimizer-zero-grad/28426/5
                # Dmitry A. Konovalov:
                #     model.zero_grad() and optimizer.zero_grad() are the same IF all your model parameters
                #     are in that optimizer.
                #     I found it is safer to call model.zero_grad() to make sure all grads are zero,
                #     e.g. if you have two or more optimizers for one model.
                # Ref:https://discuss.pytorch.org/t/zero-grad-optimizer-or-net/1887/9
                # ptrblck：
                #      if you pass all parameters of your model to the optimizer, both calls will be equal.
                #      model.zero_grad() would clear all parameters of the model,
                #      while the optimizerX.zero_grad() call will just clean
                #      the gradients of the parameters that were passed to it.

                # loss.backward() **之前** 清空模型梯度
                self.model.zero_grad()
                # 参看上述注释，这里只需要model.zero_grad()
                # self.optimizer.zero_grad()
                loss.backward()
                if self.configs.max_grad_norm > 0:
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.configs.max_grad_norm)
                self.optimizer.step()
                if self.optim_scheduler:
                    self.optim_scheduler.step()  # Update learning rate schedule
            loss = loss.detach().cpu().item()
        else:
            loss = None
        if calc_prediction:
            assert sentence_lengths
            weights = weights.unsqueeze(3)
            head_probs = torch.sigmoid(unlabeled_scores).unsqueeze(3)
            label_probs = torch.softmax(labeled_scores, dim=3)
            batch_probs = head_probs * label_probs * weights
            batch_probs = batch_probs.detach().cpu().numpy()
            # debug_print(batch_probs)
            sem_graph = sdp_decoder(batch_probs, sentence_lengths)
            sem_sents = parse_semgraph(sem_graph, sentence_lengths)
            batch_prediction = self.graph_vocab.parse_to_sent_batch(sem_sents)
        else:
            batch_prediction = None
        return loss, batch_prediction

    def train(self, train_data_loader, dev_data_loader=None, dev_CoNLLU_file=None):
        self.optimizer, self.optim_scheduler = get_optimizer(self.configs, self.model)
        global_step = 0
        best_result = BestResult()
        self.model.zero_grad()
        set_seed(self.configs)  # Added here for reproductibility (even between python 2 and 3)
        train_stop = False
        if self.configs.local_rank in [-1, 0] and not self.configs.no_output:
            summary_writer = SummaryWriter(log_dir=self.configs.summary_dir)
        for epoch in range(1, self.configs.max_train_epochs + 1):
            epoch_ave_loss = 0
            train_data_loader = tqdm(train_data_loader, desc=f'Training epoch {epoch}',
                                     disable=self.configs.local_rank not in [-1, 0])
            # 某些模型在训练时可能需要一些定制化的操作，默认什么都不做
            # 具体参考子类中_custom_train_operations的实现
            self._custom_train_operations(epoch)
            for step, batch in enumerate(train_data_loader):
                batch = tuple(t.to(self.configs.device) for t in batch)
                self.model.train()
                # debug_print(batch)
                # word_mask:以word为单位，1为真实输入，0为PAD
                unpacked_batch = self._unpack_batch(batch)
                # word_pad_mask:以word为单位，1为PAD，0为真实输入
                word_pad_mask = torch.eq(unpacked_batch['word_mask'], 0)
                model_output = self.model(unpacked_batch['inputs'])
                unlabeled_scores, labeled_scores = model_output['unlabeled_scores'], model_output['labeled_scores']
                labeled_target = unpacked_batch['dep_ids']
                unlabeled_target = labeled_target.ge(1).to(unlabeled_scores.dtype)
                if self.configs.use_pos:
                    pos_logits = model_output['pos_logits']
                    pos_target = unpacked_batch['pos_ids']
                else:
                    pos_target = pos_logits = None
                # Calc loss and update:
                loss, _ = self._update_and_predict(unlabeled_scores, labeled_scores, unlabeled_target, labeled_target,
                                                   word_pad_mask,
                                                   # label_loss_ratio=self.model.module.label_loss_ratio if hasattr(self.model,'module') else self.model.label_loss_ratio,
                                                   calc_loss=True, update=True, calc_prediction=False,
                                                   pos_logits=pos_logits, pos_target=pos_target,
                                                   summary_writer=summary_writer if self.configs.local_rank in [-1,
                                                                                                                0] else None,
                                                   global_step=global_step)
                global_step += 1
                if loss is not None:
                    epoch_ave_loss += loss

                if global_step % self.configs.eval_interval == 0 and self.configs.local_rank in [-1, 0]:
                    if not self.configs.no_output:
                        summary_writer.add_scalar('train_loss/loss', loss, global_step)
                        # 记录学习率
                        for i, param_group in enumerate(self.optimizer.param_groups):
                            summary_writer.add_scalar(f'lr/group_{i}', param_group['lr'], global_step)
                    if dev_data_loader and self.configs.local_rank in [-1, 0]:
                        UAS, LAS = self.dev(dev_data_loader, dev_CoNLLU_file)
                        if not self.configs.no_output:
                            summary_writer.add_scalar('metrics/uas', UAS, global_step)
                            summary_writer.add_scalar('metrics/las', LAS, global_step)
                        if best_result.is_new_record(LAS=LAS, UAS=UAS,
                                                     epoch=epoch) and self.configs.local_rank in [-1, 0]:
                            self.logger.info(f"\n## NEW BEST RESULT in epoch {epoch} ##")
                            self.logger.info('\n' + str(best_result))
                            # 保存最优模型：
                            if not self.configs.no_output:
                                if hasattr(self.model, 'module'):
                                    # 多卡,torch.nn.DataParallel封装model
                                    self.model.module.save_pretrained(self.configs.output_model_dir)
                                else:
                                    self.model.save_pretrained(self.configs.output_model_dir)

                if self.configs.early_stop and epoch - best_result.best_LAS_epoch > self.configs.early_stop_epochs \
                        and self.configs.local_rank == -1:
                    # 当使用 torch.distributed 训练时无法使用 early stop
                    # todo fix bug [bug when use torch.distributed.launch !!]
                    self.logger.info(f'\n## Early stop in step:{global_step} ##')
                    train_stop = True
                    break
            if train_stop:
                break
            # print(f'\n- Epoch {epoch} average loss : {epoch_ave_loss / len(train_data_loader)}')
            if self.configs.local_rank in [-1, 0] and not self.configs.no_output:
                summary_writer.add_scalar('epoch_loss', epoch_ave_loss / len(train_data_loader), epoch)
        if self.configs.local_rank in [-1, 0] and not self.configs.no_output:
            with open(self.configs.dev_result_path, 'w', encoding='utf-8')as f:
                f.write(str(best_result) + '\n')
            self.logger.info("\n## BEST RESULT in Training ##")
            self.logger.info('\n' + str(best_result))
            summary_writer.close()

    def dev(self, dev_data_loader, dev_CoNLLU_file, input_conllu_path=None, output_conllu_path=None):
        if not isinstance(dev_CoNLLU_file, CoNLLFile):
            raise RuntimeError(f'dev_conllu_file type:{type(dev_CoNLLU_file)}')
        if input_conllu_path is None:
            input_conllu_path = os.path.join(self.configs.data_dir, self.configs.dev_file)
        if output_conllu_path is None:
            output_conllu_path = self.configs.dev_output_path if not self.configs.no_output else None
        dev_data_loader = tqdm(dev_data_loader, desc='Evaluation')
        predictions = []
        for step, batch in enumerate(dev_data_loader):
            self.model.eval()
            batch = tuple(t.to(self.configs.device) for t in batch)
            unpacked_batch = self._unpack_batch(batch)
            """
            unpacked_batch = {
                    'inputs': inputs,
                    'word_mask': word_mask,
                    'sent_len': sent_len,
                    'dep_ids': dep_ids,
                    'pos_ids': pos_ids,
                }
            """
            inputs, word_mask, sent_lens, dep_ids = unpacked_batch['inputs'], unpacked_batch['word_mask'], \
                                                    unpacked_batch['sent_len'], unpacked_batch['dep_ids']
            word_mask = torch.eq(word_mask, 0)
            model_output = self.model(inputs)
            unlabeled_scores, labeled_scores = model_output['unlabeled_scores'], model_output['labeled_scores']
            try:
                with torch.no_grad():
                    _, batch_prediction = self._update_and_predict(unlabeled_scores, labeled_scores, None, None,
                                                                   word_mask,
                                                                   # label_loss_ratio=self.model.module.label_loss_ratio if hasattr(self.model,'module') else self.model.label_loss_ratio,
                                                                   sentence_lengths=sent_lens,
                                                                   calc_loss=False, update=False, calc_prediction=True)
            except Exception as e:
                for b in batch:
                    print(b.shape)
                raise e
            predictions += batch_prediction
            # batch_sent_lens += sent_lens
        dev_CoNLLU_file.set(['deps'], [dep for sent in predictions for dep in sent])
        if output_conllu_path:
            dev_CoNLLU_file.write_conll(output_conllu_path)
        UAS, LAS = sdp_scorer.score(output_conllu_path, input_conllu_path)
        return UAS, LAS

    def inference(self, inference_data_loader, inference_CoNLLU_file, output_conllu_path):
        inference_data_loader = tqdm(inference_data_loader, desc='Inference')
        predictions = []
        for step, batch in enumerate(inference_data_loader):
            self.model.eval()
            unpacked_batch = self._unpack_batch(batch)
            """
            unpacked_batch = {
                    'inputs': inputs,
                    'word_mask': word_mask,
                    'sent_len': sent_len,
                    'dep_ids': dep_ids,
                    'pos_ids': pos_ids,
                }
            """
            inputs, word_mask, sent_lens, _ = unpacked_batch['inputs'], unpacked_batch['word_mask'], \
                                              unpacked_batch['sent_len'], unpacked_batch['dep_ids']
            word_mask = torch.eq(word_mask, 0)
            model_output = self.model(inputs)
            unlabeled_scores, labeled_scores = model_output['unlabeled_scores'], model_output['labeled_scores']
            with torch.no_grad():
                _, batch_prediction = self._update_and_predict(unlabeled_scores, labeled_scores, None, None, word_mask,
                                                               # label_loss_ratio=self.model.label_loss_ratio if not self.args.data_paralle else self.model.module.label_loss_ratio,
                                                               sentence_lengths=sent_lens,
                                                               calc_loss=False, update=False, calc_prediction=True)
            predictions += batch_prediction
        inference_CoNLLU_file.set(['deps'], [dep for sent in predictions for dep in sent])
        inference_CoNLLU_file.write_conll(output_conllu_path)
        return predictions


class TransformerBaseTrainer(BaseDependencyTrainer):
    def _unpack_batch(self, batch):
        pass


# class CharRNNBaseTrainer(BaseDependencyTrainer):
#     pass


if __name__ == '__main__':
    pass
