# -*- coding:utf-8 -*-
# Author: hankcs
# Date: 2021-04-28 17:33
import datetime
import functools
import logging
import os
from typing import Union, List, Callable

import torch
from torch.utils.data import DataLoader
from transformers import get_constant_schedule_with_warmup, T5ForConditionalGeneration
from transformers.models.bart.modeling_bart import BartForConditionalGeneration

from hanlp.common.dataset import SamplerBuilder, SortingSamplerBuilder, PadSequenceDataLoader
from hanlp.common.structure import History
from hanlp.common.torch_component import TorchComponent
from hanlp.common.vocab import Vocab
from hanlp.components.amr.seq2seq.dataset.dataset import AMRDataset, dfs_linearize_tokenize
from hanlp.components.amr.seq2seq.dataset.penman import AMRGraph
from hanlp.components.amr.seq2seq.dataset.tokenization_bart import PENMANBartTokenizer
from hanlp.components.amr.seq2seq.dataset.tokenization_t5 import PENMANT5Tokenizer
from hanlp.components.amr.seq2seq.evaluation import write_predictions, compute_smatch
from hanlp.components.amr.seq2seq.optim import RAdam
from hanlp.layers.transformers.pt_imports import PretrainedConfig, AutoConfig_
from hanlp.layers.transformers.resource import get_model_mirror, get_tokenizer_mirror
from hanlp.metrics.amr.smatch_eval import smatch_eval
from hanlp.metrics.mtl import MetricDict
from hanlp.utils.time_util import CountdownTimer
from hanlp_common.constant import IDX
from hanlp_common.util import merge_locals_kwargs, reorder


class Seq2seq_AMR_Parser(TorchComponent):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._transformer_config: PretrainedConfig = None
        self._tokenizer: PENMANBartTokenizer = None
        self.model: BartForConditionalGeneration = None

    def build_dataloader(self, data, batch_size,
                         gradient_accumulation=1,
                         shuffle=False,
                         sampler_builder: SamplerBuilder = None,
                         device=None,
                         logger: logging.Logger = None,
                         **kwargs) -> DataLoader:
        dataset = self.build_dataset(data, not shuffle)
        if self.vocabs.mutable:
            self.build_vocabs(dataset, logger)
        self.finalize_dataset(dataset, logger)
        if isinstance(data, str):
            dataset.purge_cache()
            timer = CountdownTimer(len(dataset))
            max_num_tokens = 0
            # lc = Counter()
            for each in dataset:
                max_num_tokens = max(max_num_tokens, len(each['text_token_ids']))
                # lc[len(each['text_token_ids'])] += 1
                timer.log(f'Preprocessing and caching samples (longest sequence {max_num_tokens})'
                          f'[blink][yellow]...[/yellow][/blink]')
            # print(lc.most_common())
            if self.vocabs.mutable:
                self.vocabs.lock()
                self.vocabs.summary(logger)

        if not sampler_builder:
            sampler_builder = SortingSamplerBuilder(batch_max_tokens=500)
        sampler = sampler_builder.build([len(x['text_token_ids']) for x in dataset], shuffle,
                                        gradient_accumulation if dataset.cache else 1)
        return self._create_dataloader(dataset, batch_size, device, sampler, shuffle)

    def _create_dataloader(self, dataset, batch_size, device, sampler, shuffle):
        return PadSequenceDataLoader(dataset, batch_size, shuffle, device=device, batch_sampler=sampler,
                                     pad=self._get_pad_dict())

    def _get_pad_dict(self):
        return {'text_token_ids': self._transformer_config.pad_token_id,
                'graph_token_ids': self._transformer_config.pad_token_id}

    def finalize_dataset(self, dataset, logger: logging.Logger = None):
        dataset.append_transform(functools.partial(dfs_linearize_tokenize, tokenizer=self._tokenizer,
                                                   remove_space='chinese' in self.config.transformer))

    def build_dataset(self, data, generate_idx):
        dataset = AMRDataset(data, generate_idx=generate_idx)
        return dataset

    def collect_additional_tokens(self, additional_tokens, dataset):
        pred_min = self.config.pred_min
        frames = dataset.get_frames()
        for token, freq in frames.items():
            if freq >= pred_min:
                additional_tokens.add(token)
        for token, freq in dataset.get_roles().items():
            additional_tokens.add(token)
        additional_tokens.update(self.config.additional_tokens)

    def build_tokenizer(self, additional_tokens) -> PENMANBartTokenizer:
        transformer = self.config.transformer
        if 't5-' in transformer:
            cls = PENMANT5Tokenizer
        elif 'bart-' in transformer:
            cls = PENMANBartTokenizer
        else:
            raise NotImplemented(f'Unsupported transformer {transformer}')
        transformer = get_tokenizer_mirror(transformer)
        self._tokenizer = cls.from_pretrained(
            transformer,
            collapse_name_ops=self.config.collapse_name_ops,
            use_pointer_tokens=self.config.use_pointer_tokens,
            raw_graph=self.config.raw_graph,
            additional_tokens=additional_tokens,
            recategorization_tokens=self.config.recategorization_tokens,
            config=self._transformer_config,
        )
        return self._tokenizer

    def build_optimizer(self, trn, lr, epochs, gradient_accumulation, warmup_steps, weight_decay, **kwargs):
        num_training_steps = len(trn) * epochs // gradient_accumulation
        if isinstance(warmup_steps, float):
            warmup_steps = int(num_training_steps * warmup_steps)
        optimizer = RAdam(
            self.model.parameters(),
            lr=lr,
            weight_decay=weight_decay)
        scheduler = get_constant_schedule_with_warmup(
            optimizer,
            num_warmup_steps=warmup_steps)
        return optimizer, scheduler

    def build_criterion(self, **kwargs):
        pass

    def build_metric(self, **kwargs):
        pass

    def execute_training_loop(self, trn: DataLoader, dev: DataLoader, epochs, criterion, optimizer, metric, save_dir,
                              logger: logging.Logger, devices, ratio_width=None, dev_data=None, eval_after=None,
                              **kwargs):
        best_epoch, best_metric = 0, -1
        if isinstance(eval_after, float):
            eval_after = int(epochs * eval_after)
        timer = CountdownTimer(epochs)
        history = History()
        for epoch in range(1, epochs + 1):
            logger.info(f"[yellow]Epoch {epoch} / {epochs}:[/yellow]")
            self.fit_dataloader(trn, criterion, optimizer, metric, logger, history=history, ratio_width=ratio_width,
                                **self.config)
            if epoch > eval_after:
                dev_metric = self.evaluate_dataloader(dev, criterion, logger=logger, ratio_width=ratio_width,
                                                      output=os.path.join(save_dir, 'dev.pred.txt'),
                                                      input=dev_data, use_fast=True)
            timer.update()
            report = f"{timer.elapsed_human} / {timer.total_time_human} ETA: {timer.eta_human}"
            if epoch > eval_after:
                if dev_metric > best_metric:
                    best_epoch, best_metric = epoch, dev_metric
                    self.save_weights(save_dir)
                    report += ' [red](saved)[/red]'
                else:
                    report += f' ({epoch - best_epoch})'
                # if epoch - best_epoch >= patience:
                #     report += ' early stop'
            logger.info(report)
            # if epoch - best_epoch >= patience:
            #     break
        if not best_epoch:
            self.save_weights(save_dir)
        elif best_epoch != epoch:
            self.load_weights(save_dir)
        logger.info(f"Max score of dev is {best_metric} at epoch {best_epoch}")
        logger.info(f"Average time of each epoch is {timer.elapsed_average_human}")
        logger.info(f"{timer.elapsed_human} elapsed")
        return best_metric

    def fit_dataloader(self, trn: DataLoader, criterion, optimizer, metric, logger: logging.Logger,
                       history: History = None, gradient_accumulation=1, ratio_percentage=None, **kwargs):
        optimizer, scheduler = optimizer
        self.model.train()
        timer = CountdownTimer(history.num_training_steps(len(trn), gradient_accumulation=gradient_accumulation))
        total_loss = 0
        for batch in trn:
            output_dict = self.feed_batch(batch)
            loss = output_dict['loss']
            if gradient_accumulation and gradient_accumulation > 1:
                loss /= gradient_accumulation
            loss.backward()
            total_loss += loss.item()
            if history.step(gradient_accumulation):
                self._step(optimizer, scheduler)
                timer.log(self.report_metrics(total_loss / (timer.current + 1)),
                          ratio_percentage=ratio_percentage, logger=logger)
            del loss
            del output_dict
        return total_loss / max(timer.total, 1)

    def _step(self, optimizer, scheduler):
        if self.config.grad_norm:
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.grad_norm)
        optimizer.step()
        if scheduler:
            scheduler.step()
        optimizer.zero_grad()

    def report_metrics(self, loss):
        return f'loss: {loss:.4f}'

    def feed_batch(self, batch):
        input_ids, labels = batch['text_token_ids'], batch.get('graph_token_ids')
        attention_mask = input_ids.ne(self.model.config.pad_token_id).to(torch.long)
        if labels is not None:
            decoder_input_ids = labels[:, :-1]
            labels = labels[:, 1:].contiguous()
        else:
            decoder_input_ids = None
        return self.model(input_ids=input_ids, attention_mask=attention_mask, decoder_input_ids=decoder_input_ids,
                          labels=labels)

    @torch.no_grad()
    def evaluate_dataloader(self, data: DataLoader, criterion: Callable, metric=None, output=False, ratio_width=None,
                            logger=None, input=None, use_fast=False,
                            **kwargs):
        self.model.eval()
        timer = CountdownTimer(len(data))
        graphs = []
        orders = []
        smatch = 0
        for idx, batch in enumerate(data):
            graphs_per_batch = self.predict_amrs(batch)
            graphs_per_batch = [x[0] for x in graphs_per_batch]
            # Copy meta data from gold graph
            for gp, gg in zip(graphs_per_batch, batch['amr']):
                metadata = gg.metadata.copy()
                metadata['annotator'] = f'{self.config.transformer}-amr'
                metadata['date'] = str(datetime.datetime.now())
                if 'save-date' in metadata:
                    del metadata['save-date']
                gp.metadata = metadata
            graphs.extend(graphs_per_batch)
            orders.extend(batch[IDX])
            if idx == timer.total - 1:
                graphs = reorder(graphs, orders)
                write_predictions(output, self._tokenizer, graphs)
                try:
                    if use_fast:
                        smatch = compute_smatch(output, input)
                    else:
                        smatch = smatch_eval(output, input, use_fast=False)
                except:
                    pass
                timer.log(smatch.cstr() if isinstance(smatch, MetricDict) else f'{smatch:.2%}', ratio_percentage=False,
                          logger=logger)
            else:
                timer.log(ratio_percentage=False, logger=logger)

        return smatch

    def predict_amrs(self, batch, beam_size=1):
        out = self._model_generate(batch, beam_size)
        tokens = []
        for i1 in range(0, out.size(0), beam_size):
            tokens_same_source = []
            tokens.append(tokens_same_source)
            for i2 in range(i1, i1 + beam_size):
                tokk = out[i2].tolist()
                tokens_same_source.append(tokk)
        tokens = [t for tt in tokens for t in tt]
        graphs = []
        tokenizer = self._tokenizer
        for i1 in range(0, len(tokens), beam_size):
            graphs_same_source = []
            graphs.append(graphs_same_source)
            for i2 in range(i1, i1 + beam_size):
                tokk = tokens[i2]
                graph, status, (lin, backr) = tokenizer.decode_amr(tokk, restore_name_ops=False)
                graph.status = status
                graph.nodes = lin
                graph.backreferences = backr
                graph.tokens = tokk
                graphs_same_source.append(graph)
            graphs_same_source[:] = \
                tuple(zip(*sorted(enumerate(graphs_same_source), key=lambda x: (x[1].status.value, x[0]))))[1]

        return graphs

    def _model_generate(self, batch, beam_size):
        input_ids = batch['text_token_ids']
        attention_mask = input_ids.ne(self.model.config.pad_token_id).to(torch.long)
        out = self.model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_length=1024,
            decoder_start_token_id=0,
            num_beams=beam_size,
            num_return_sequences=beam_size)
        return out

    def build_model(self, training=True, **kwargs) -> torch.nn.Module:
        # noinspection PyTypeChecker
        transformer = self.config.transformer
        cls = self._get_model_cls(transformer)
        transformer = get_model_mirror(self.config.transformer)
        model: cls = cls.from_pretrained(
            transformer,
            config=self._transformer_config) if training else cls(self._transformer_config)
        if not training:
            self.build_tokenizer(self.vocabs['additional_tokens'])
        tokenizer = self._tokenizer
        model.resize_token_embeddings(len(tokenizer.encoder))
        if training:
            self._init_new_embeddings(model if cls == T5ForConditionalGeneration else model.model, tokenizer)
        return model

    def _get_model_cls(self, transformer: str):
        if 't5-' in transformer:
            cls = T5ForConditionalGeneration
        elif 'bart-' in transformer:
            cls = BartForConditionalGeneration
        else:
            raise NotImplemented(f'Unsupported transformer {transformer}')
        return cls

    @staticmethod
    def _init_new_embeddings(model, tokenizer):
        modified = 0
        encoder = tokenizer.encoder
        for tok, idx in encoder.items():
            tok = tok.lstrip(tokenizer.INIT)

            if idx < tokenizer.old_enc_size:
                continue

            elif tok.startswith('<pointer:') and tok.endswith('>'):
                tok_split = ['pointer', str(tok.split(':')[1].strip('>'))]

            elif tok.startswith('<'):
                continue

            elif tok.startswith(':'):
                if tok.startswith(':op'):
                    tok_split = ['relation', 'operator', str(int(tok[3:]))]

                elif tok.startswith(':snt'):
                    tok_split = ['relation', 'sentence', str(int(tok[4:]))]

                elif tok.startswith(':ARG'):
                    tok_split = ['relation', 'argument', str(int(tok[4:]))]
                else:
                    tok_split = ['relation'] + tok.lstrip(':').split('-')
            else:
                tok_split = tok.split('-')

            tok_split_ = tok_split
            tok_split = []
            for s in tok_split_:
                s_ = s + tokenizer.INIT
                if s_ in encoder:
                    tok_split.append(s_)
                else:
                    tok_split.extend(tokenizer._tok_bpe(s))

            vecs = []
            for s in tok_split:
                idx_split = encoder.get(s, -1)
                if idx_split > -1:
                    vec_split = model.encoder.embed_tokens.weight.data[idx_split].clone()
                    vecs.append(vec_split)

            if vecs:
                vec = torch.stack(vecs, 0).mean(0)
                noise = torch.empty_like(vec)
                noise.uniform_(-0.1, +0.1)
                model.encoder.embed_tokens.weight.data[idx] = vec + noise
                modified += 1

    def input_is_flat(self, data):
        return isinstance(data, str)

    def predict(self, data: Union[str, List[str]], beautiful_amr_graph=True, **kwargs):
        flat = self.input_is_flat(data)
        if flat:
            data = [data]
        dataloader = self.build_dataloader([{'text': x} for x in data], **self.config, device=self.device)
        orders = []
        results = []
        for batch in dataloader:
            graphs = self.predict_amrs(batch)
            graphs = [x[0] for x in graphs]
            if beautiful_amr_graph:
                graphs = [AMRGraph(x.triples, x.top, x.epidata, x.metadata) for x in graphs]
            results.extend(graphs)
            orders.extend(batch[IDX])
        results = reorder(results, orders)
        if flat:
            results = results[0]
        return results

    def fit(self, trn_data, dev_data, save_dir, batch_size=32, epochs=30,
            transformer='facebook/bart-base',
            lr=5e-05,
            grad_norm=2.5,
            weight_decay=0.004,
            warmup_steps=1,
            dropout=0.25,
            attention_dropout=0.0,
            pred_min=5,
            eval_after=0.5,
            collapse_name_ops=False,
            use_pointer_tokens=True,
            raw_graph=False,
            gradient_accumulation=1,
            recategorization_tokens=(
                    'PERSON', 'COUNTRY', 'QUANTITY', 'ORGANIZATION', 'DATE_ATTRS', 'NATIONALITY', 'LOCATION', 'ENTITY',
                    'CITY',
                    'MISC', 'ORDINAL_ENTITY', 'IDEOLOGY', 'RELIGION', 'STATE_OR_PROVINCE', 'URL', 'CAUSE_OF_DEATH', 'O',
                    'TITLE', 'DATE', 'NUMBER', 'HANDLE', 'SCORE_ENTITY', 'DURATION', 'ORDINAL', 'MONEY', 'SET',
                    'CRIMINAL_CHARGE', '_1', '_2', '_3', '_4', '_2', '_5', '_6', '_7', '_8', '_9', '_10', '_11', '_12',
                    '_13',
                    '_14', '_15'),
            additional_tokens=(
                    'date-entity', 'government-organization', 'temporal-quantity', 'amr-unknown', 'multi-sentence',
                    'political-party', 'monetary-quantity', 'ordinal-entity', 'religious-group', 'percentage-entity',
                    'world-region', 'url-entity', 'political-movement', 'et-cetera', 'at-least', 'mass-quantity',
                    'have-org-role-91', 'have-rel-role-91', 'include-91', 'have-concession-91', 'have-condition-91',
                    'be-located-at-91', 'rate-entity-91', 'instead-of-91', 'hyperlink-91', 'request-confirmation-91',
                    'have-purpose-91', 'be-temporally-at-91', 'regardless-91', 'have-polarity-91', 'byline-91',
                    'have-manner-91', 'have-part-91', 'have-quant-91', 'publication-91', 'be-from-91', 'have-mod-91',
                    'have-frequency-91', 'score-on-scale-91', 'have-li-91', 'be-compared-to-91', 'be-destined-for-91',
                    'course-91', 'have-subevent-91', 'street-address-91', 'have-extent-91', 'statistical-test-91',
                    'have-instrument-91', 'have-name-91', 'be-polite-91', '-00', '-01', '-02', '-03', '-04', '-05',
                    '-06',
                    '-07', '-08', '-09', '-10', '-11', '-12', '-13', '-14', '-15', '-16', '-17', '-18', '-19', '-20',
                    '-21',
                    '-22', '-23', '-24', '-25', '-26', '-27', '-28', '-29', '-20', '-31', '-32', '-33', '-34', '-35',
                    '-36',
                    '-37', '-38', '-39', '-40', '-41', '-42', '-43', '-44', '-45', '-46', '-47', '-48', '-49', '-50',
                    '-51',
                    '-52', '-53', '-54', '-55', '-56', '-57', '-58', '-59', '-60', '-61', '-62', '-63', '-64', '-65',
                    '-66',
                    '-67', '-68', '-69', '-70', '-71', '-72', '-73', '-74', '-75', '-76', '-77', '-78', '-79', '-80',
                    '-81',
                    '-82', '-83', '-84', '-85', '-86', '-87', '-88', '-89', '-90', '-91', '-92', '-93', '-94', '-95',
                    '-96',
                    '-97', '-98', '-of'),
            devices=None,
            logger=None,
            seed=None,
            finetune: Union[bool, str] = False,
            eval_trn=True,
            _device_placeholder=False,
            **kwargs):
        """

        Args:
            trn_data:
            dev_data:
            save_dir:
            batch_size:
            epochs:
            transformer:
            lr:
            grad_norm:
            weight_decay:
            warmup_steps:
            dropout:
            attention_dropout:
            pred_min:
            eval_after:
            collapse_name_ops: ``True`` to merge name ops.
            use_pointer_tokens: ``True`` to use pointer tokens to represent variables.
            raw_graph: ``True`` to use the raw graph as input and skip all pre/post-processing steps.
            gradient_accumulation:
            recategorization_tokens: Tokens used in re-categorization. They will be added to tokenizer too but do not
            put them into ``additional_tokens``.
            additional_tokens: Tokens to be added to the tokenizer vocab.
            devices:
            logger:
            seed:
            finetune:
            eval_trn:
            _device_placeholder:
            **kwargs:

        Returns:

        """
        return super().fit(**merge_locals_kwargs(locals(), kwargs))

    def on_config_ready(self, **kwargs):
        super().on_config_ready(**kwargs)
        config = AutoConfig_.from_pretrained(self.config.transformer)
        config.output_past = False
        config.no_repeat_ngram_size = 0
        config.prefix = " "
        # config.output_attentions = True
        config.dropout = self.config.dropout
        config.attention_dropout = self.config.attention_dropout
        self._transformer_config = config

    def evaluate(self, tst_data, save_dir=None, logger: logging.Logger = None, batch_size=None, output=True,
                 cache=None, ret_speed=False, **kwargs):
        return super().evaluate(tst_data, save_dir, logger, batch_size, output, cache, ret_speed, **kwargs)

    def build_vocabs(self, trn: torch.utils.data.Dataset, logger: logging.Logger):
        additional_tokens = set()
        self.collect_additional_tokens(additional_tokens, trn)
        additional_tokens = sorted(additional_tokens)
        self.build_tokenizer(additional_tokens)
        self.vocabs['additional_tokens'] = Vocab(idx_to_token=list(additional_tokens))