"""
A trainer class to handle training and testing of models.
"""

import os
import sys
import numpy as np
from collections import Counter
import logging
import torch
from torch import nn
import torch.nn.init as init

import stanza.models.common.seq2seq_constant as constant
from stanza.models.common.doc import TEXT, UPOS
from stanza.models.common.foundation_cache import load_charlm
from stanza.models.common.seq2seq_model import Seq2SeqModel
from stanza.models.common.char_model import CharacterLanguageModelWordAdapter
from stanza.models.common import utils, loss
from stanza.models.lemma import edit
from stanza.models.lemma.vocab import MultiVocab
from stanza.models.lemma_classifier.base_model import LemmaClassifier

logger = logging.getLogger('stanza')

def unpack_batch(batch, device):
    """ Unpack a batch from the data loader. """
    inputs = [b.to(device) if b is not None else None for b in batch[:6]]
    orig_idx = batch[6]
    text = batch[7]
    return inputs, orig_idx, text

class Trainer(object):
    """ A trainer for training models. """
    def __init__(self, args=None, vocab=None, emb_matrix=None, model_file=None, device=None, foundation_cache=None, lemma_classifier_args=None):
        if model_file is not None:
            # load everything from file
            self.load(model_file, args, foundation_cache, lemma_classifier_args)
        else:
            # build model from scratch
            self.args = args
            if args['dict_only']:
                self.model = None
            else:
                self.model = self.build_seq2seq(args, emb_matrix, foundation_cache)
            self.vocab = vocab
            # dict-based components
            self.word_dict = dict()
            self.composite_dict = dict()
            self.contextual_lemmatizers = []

        self.caseless = self.args.get('caseless', False)

        if not self.args['dict_only']:
            self.model = self.model.to(device)
            if self.args.get('edit', False):
                self.crit = loss.MixLoss(self.vocab['char'].size, self.args['alpha']).to(device)
                logger.debug("Running seq2seq lemmatizer with edit classifier...")
            else:
                self.crit = loss.SequenceLoss(self.vocab['char'].size).to(device)
            self.optimizer = utils.get_optimizer(self.args['optim'], self.model, self.args['lr'])

    def build_seq2seq(self, args, emb_matrix, foundation_cache):
        charmodel = None
        charlms = []
        if args is not None and args.get('charlm_forward_file', None):
            charmodel_forward = load_charlm(args['charlm_forward_file'], foundation_cache=foundation_cache)
            charlms.append(charmodel_forward)
        if args is not None and args.get('charlm_backward_file', None):
            charmodel_backward = load_charlm(args['charlm_backward_file'], foundation_cache=foundation_cache)
            charlms.append(charmodel_backward)
        if len(charlms) > 0:
            charlms = nn.ModuleList(charlms)
            charmodel = CharacterLanguageModelWordAdapter(charlms)
        model = Seq2SeqModel(args, emb_matrix=emb_matrix, contextual_embedding=charmodel)
        return model

    def update(self, batch, eval=False):
        device = next(self.model.parameters()).device
        inputs, orig_idx, text = unpack_batch(batch, device)
        src, src_mask, tgt_in, tgt_out, pos, edits = inputs

        if eval:
            self.model.eval()
        else:
            self.model.train()
            self.optimizer.zero_grad()
        log_probs, edit_logits = self.model(src, src_mask, tgt_in, pos, raw=text)
        if self.args.get('edit', False):
            assert edit_logits is not None
            loss = self.crit(log_probs.view(-1, self.vocab['char'].size), tgt_out.view(-1), \
                    edit_logits, edits)
        else:
            loss = self.crit(log_probs.view(-1, self.vocab['char'].size), tgt_out.view(-1))
        loss_val = loss.data.item()
        if eval:
            return loss_val

        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args['max_grad_norm'])
        self.optimizer.step()
        return loss_val

    def predict(self, batch, beam_size=1, vocab=None):
        if vocab is None:
            vocab = self.vocab

        device = next(self.model.parameters()).device
        inputs, orig_idx, text = unpack_batch(batch, device)
        src, src_mask, tgt, tgt_mask, pos, edits = inputs

        self.model.eval()
        batch_size = src.size(0)
        preds, edit_logits = self.model.predict(src, src_mask, pos=pos, beam_size=beam_size, raw=text)
        pred_seqs = [vocab['char'].unmap(ids) for ids in preds] # unmap to tokens
        pred_seqs = utils.prune_decoded_seqs(pred_seqs)
        pred_tokens = ["".join(seq) for seq in pred_seqs] # join chars to be tokens
        pred_tokens = utils.unsort(pred_tokens, orig_idx)
        if self.args.get('edit', False):
            assert edit_logits is not None
            edits = np.argmax(edit_logits.data.cpu().numpy(), axis=1).reshape([batch_size]).tolist()
            edits = utils.unsort(edits, orig_idx)
        else:
            edits = None
        return pred_tokens, edits

    def postprocess(self, words, preds, edits=None):
        """ Postprocess, mainly for handing edits. """
        assert len(words) == len(preds), "Lemma predictions must have same length as words."
        edited = []
        if self.args.get('edit', False):
            assert edits is not None and len(words) == len(edits)
            for w, p, e in zip(words, preds, edits):
                lem = edit.edit_word(w, p, e)
                edited += [lem]
        else:
            edited = preds # do not edit
        # final sanity check
        assert len(edited) == len(words)
        final = []
        for lem, w in zip(edited, words):
            if len(lem) == 0 or constant.UNK in lem:
                final += [w] # invalid prediction, fall back to word
            else:
                final += [lem]
        return final

    def has_contextual_lemmatizers(self):
        return self.contextual_lemmatizers is not None and len(self.contextual_lemmatizers) > 0

    def predict_contextual(self, sentence_words, sentence_tags, preds):
        if len(self.contextual_lemmatizers) == 0:
            return preds

        # reversed so that the first lemmatizer has priority
        for contextual in reversed(self.contextual_lemmatizers):
            pred_idx = []
            pred_sent_words = []
            pred_sent_tags = []
            pred_sent_ids = []
            for sent_id, (words, tags) in enumerate(zip(sentence_words, sentence_tags)):
                indices = contextual.target_indices(words, tags)
                for idx in indices:
                    pred_idx.append(idx)
                    pred_sent_words.append(words)
                    pred_sent_tags.append(tags)
                    pred_sent_ids.append(sent_id)
            if len(pred_idx) == 0:
                continue
            contextual_predictions = contextual.predict(pred_idx, pred_sent_words, pred_sent_tags)
            for sent_id, word_id, pred in zip(pred_sent_ids, pred_idx, contextual_predictions):
                preds[sent_id][word_id] = pred
        return preds

    def update_contextual_preds(self, doc, preds):
        """
        Update a flat list of preds with the output of the contextual lemmatizers

        - First, it unflattens the preds based on the lengths of the sentences
        - Then it uses the contextual lemmatizers
        - Finally, it reflattens the preds into the format expected by the caller
        """
        if len(self.contextual_lemmatizers) == 0:
            return preds

        sentence_words = doc.get([TEXT], as_sentences=True)
        sentence_tags = doc.get([UPOS], as_sentences=True)
        sentence_preds = []
        start_index = 0
        for sent in sentence_words:
            end_index = start_index + len(sent)
            sentence_preds.append(preds[start_index:end_index])
            start_index += len(sent)
        preds = self.predict_contextual(sentence_words, sentence_tags, sentence_preds)
        preds = [lemma for sentence in preds for lemma in sentence]
        return preds

    def update_lr(self, new_lr):
        utils.change_lr(self.optimizer, new_lr)

    def train_dict(self, triples, update_word_dict=True):
        """
        Train a dict lemmatizer given training (word, pos, lemma) triples.

        Can update only the composite_dict (word/pos) in situations where
        the data might be limited from the tags, such as when adding more
        words at pipeline time
        """
        # accumulate counter
        ctr = Counter()
        ctr.update([(p[0], p[1], p[2]) for p in triples])
        # find the most frequent mappings
        for p, _ in ctr.most_common():
            w, pos, l = p
            if (w,pos) not in self.composite_dict:
                self.composite_dict[(w,pos)] = l
            if update_word_dict and w not in self.word_dict:
                self.word_dict[w] = l
        return

    def predict_dict(self, pairs):
        """ Predict a list of lemmas using the dict model given (word, pos) pairs. """
        lemmas = []
        for p in pairs:
            w, pos = p
            if self.caseless:
                w = w.lower()
            if (w,pos) in self.composite_dict:
                lemmas += [self.composite_dict[(w,pos)]]
            elif w in self.word_dict:
                lemmas += [self.word_dict[w]]
            else:
                lemmas += [w]
        return lemmas

    def skip_seq2seq(self, pairs):
        """ Determine if we can skip the seq2seq module when ensembling with the frequency lexicon. """

        skip = []
        for p in pairs:
            w, pos = p
            if self.caseless:
                w = w.lower()
            if (w,pos) in self.composite_dict:
                skip.append(True)
            elif w in self.word_dict:
                skip.append(True)
            else:
                skip.append(False)
        return skip

    def ensemble(self, pairs, other_preds):
        """ Ensemble the dict with statistical model predictions. """
        lemmas = []
        assert len(pairs) == len(other_preds)
        for p, pred in zip(pairs, other_preds):
            w, pos = p
            if self.caseless:
                w = w.lower()
            if (w,pos) in self.composite_dict:
                lemma = self.composite_dict[(w,pos)]
            elif w in self.word_dict:
                lemma = self.word_dict[w]
            else:
                lemma = pred
            if lemma is None:
                lemma = w
            lemmas.append(lemma)
        return lemmas

    def save(self, filename, skip_modules=True):
        model_state = None
        if self.model is not None:
            model_state = self.model.state_dict()
            # skip saving modules like the pretrained charlm
            if skip_modules:
                skipped = [k for k in model_state.keys() if k.split('.')[0] in self.model.unsaved_modules]
                for k in skipped:
                    del model_state[k]
        params = {
            'model': model_state,
            'dicts': (self.word_dict, self.composite_dict),
            'vocab': self.vocab.state_dict(),
            'config': self.args,
            'contextual': [],
        }
        for contextual in self.contextual_lemmatizers:
            params['contextual'].append(contextual.get_save_dict())
        save_dir = os.path.split(filename)[0]
        if save_dir:
            os.makedirs(os.path.split(filename)[0], exist_ok=True)
        torch.save(params, filename, _use_new_zipfile_serialization=False)
        logger.info("Model saved to {}".format(filename))

    def load(self, filename, args, foundation_cache, lemma_classifier_args=None):
        try:
            checkpoint = torch.load(filename, lambda storage, loc: storage, weights_only=True)
        except BaseException:
            logger.error("Cannot load model from {}".format(filename))
            raise
        self.args = checkpoint['config']
        if args is not None:
            self.args['charlm_forward_file'] = args.get('charlm_forward_file', self.args['charlm_forward_file'])
            self.args['charlm_backward_file'] = args.get('charlm_backward_file', self.args['charlm_backward_file'])
        self.word_dict, self.composite_dict = checkpoint['dicts']
        if not self.args['dict_only']:
            self.model = self.build_seq2seq(self.args, None, foundation_cache)
            # could remove strict=False after rebuilding all models,
            # or could switch to 1.6.0 torch with the buffer in seq2seq persistent=False
            self.model.load_state_dict(checkpoint['model'], strict=False)
        else:
            self.model = None
        self.vocab = MultiVocab.load_state_dict(checkpoint['vocab'])
        self.contextual_lemmatizers = []
        for contextual in checkpoint.get('contextual', []):
            self.contextual_lemmatizers.append(LemmaClassifier.from_checkpoint(contextual, args=lemma_classifier_args))
