import os
import logging

import numpy as np
import torch
import torch.optim as optim

from TakeBlipNer import utils, data, nermodel

import time


class LSTMCRFTrainer(object):
    def __init__(self, bilstmcrf_model: nermodel.LSTMCRF_NER, input_vocab, input_path,
                 label_vocab, postag_model, postag_label_vocab, save_dir, tensorboard,
                 ckpt_period, val, val_period, samples, pad_string, unk_string,
                 batch_size, shuffle, label_column, encoding, separator,
                 use_pre_processing, learning_rate, learning_rate_decay,
                 max_patience, max_decay_num, patience_threshold, val_path=None,
                 epochs=5, optimizer=optim.Adam, n_iter=5, azureml_logger=None):

        self.model = bilstmcrf_model
        self.epochs = epochs
        self.optimizer_cls = optimizer
        self.optimizer = optimizer(self.model.parameters(), learning_rate)
        self.save_dir = save_dir
        self.tensorboard = tensorboard
        self.input_vocab = input_vocab
        self.postag_label_vocab = postag_label_vocab
        self.postag_model = postag_model
        self.inputs_vocab = [input_vocab] + [postag_label_vocab]
        self.label_vocab = label_vocab
        self.input_path = input_path
        self.val_path = val_path
        self.val = val
        self.writer = None
        self.samples = samples
        self.encoding = encoding
        self.label_column = label_column
        self.separator = separator
        self.use_pre_processing = use_pre_processing
        self.pad_string = pad_string
        self.unk_string = unk_string
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.n_iter = n_iter
        self.min_dev_loss = float('inf')
        self.patience = 0
        self.patience_threshold = patience_threshold
        self.max_patience = max_patience
        self.decay_num = 0
        self.max_decay_num = max_decay_num
        self.learning_rate_decay = learning_rate_decay
        self.logic_break = False
        self.azureml_logger = azureml_logger
        self.validation_number = 0

        self.repeatables = {}

        if self.val:
            self.repeatables[val_period] = self.validate

    @staticmethod
    def get_longest_word(vocab):
        lens = [(len(w), w) for w in vocab.f2i]
        return max(lens, key=lambda x: x[0])[1]

    def display_row(self, items, widths):
        assert len(items) == len(widths)

        padded = ['{{:>{}s}}'.format(w).format(item)
                  for item, w in zip(items, widths)]
        logging.info(' '.join(padded))

    def load_data(self, train):
        path = self.input_path if train else self.val_path

        dataset = data.MultiSentWordDataset(
            path=path,
            label_column=self.label_column,
            encoding=self.encoding,
            separator=self.separator,
            use_pre_processing=self.use_pre_processing)

        dataloader = data.MultiSentWordDataLoader(
            dataset=dataset,
            vocabs=[self.input_vocab]+[self.postag_label_vocab]+[self.label_vocab],
            pad_string=self.pad_string,
            unk_string=self.unk_string,
            batch_size=self.batch_size,
            shuffle=self.shuffle,
            tensor_lens=True,
            postag_model=self.postag_model)
        return dataloader

    def display_samples(self, inputs, targets, lstms, crfs):

        inputs = [[utils.lexicalize(s, v) for s in input]
                  for input, v in zip(inputs, self.inputs_vocab)]
        targets = [utils.lexicalize(s, self.label_vocab) for s in targets]
        lstms = [utils.lexicalize(s, self.label_vocab) for s in lstms]
        crfs = [utils.lexicalize(s, self.label_vocab) for s in crfs]
        transposed = list(zip(*(inputs + [lstms, crfs, targets])))
        col_names = [f'INPUT{i + 1:02d}' for i in range(len(inputs))] + \
                    ['LSTM', 'CRF', 'TARGETS']
        vocabs = self.inputs_vocab + [self.label_vocab] * 3
        col_widths = [max(len(self.get_longest_word(v)), len(c))
                      for v, c in zip(vocabs, col_names)]

        for i, sample in enumerate(transposed):
            rows = list(zip(*sample))

            logging.info('')
            logging.info(f'SAMPLE #{i + 1}')
            self.display_row(col_names, col_widths)
            for row in rows:
                self.display_row(row, col_widths)

    @staticmethod
    def random_idx(max_count, subset=None):
        idx = np.random.permutation(np.arange(max_count))

        if subset is not None:
            return idx[:subset]

        return idx

    @staticmethod
    def gather(lst, idx):
        return [lst[i] for i in idx]

    def save_model(self, steps=None):
        save_path = os.path.join(self.save_dir, 'nermodel.pkl')
        torch.save(self.model, save_path)
        if self.azureml_logger:
            os.system('cp %s %s' % ('nermodel.py', self.save_dir))
            self.azureml_logger.register_ner_model(
                training_mount_folder=self.save_dir,
                ner_registry='NerModel'
            )

    def on_epoch_complete(self, epoch_idx, global_iter, global_step):
        for period_checker, func in self.repeatables.items():
            if period_checker(epochs=epoch_idx,
                              iters=global_iter,
                              steps=global_step):
                func(steps=global_step)
                self.model.train(True)

    def train(self):
        self.model.train(True)
        global_step = 0
        global_iter = 0

        start = time.time()
        for epoch_idx in range(self.epochs):
            epoch_idx += 1
            train_data = self.load_data(train=True)
            for iteration_idx, (batch, lens, _) in enumerate(train_data):
                batch_size = batch[0].size(0)
                iteration_idx += 1
                global_step += batch_size
                global_iter += 1

                self.model.zero_grad()

                batch_sentences, batch_labels = batch[:-1], batch[-1]
                batch_sentences, batch_labels_var, lens_s = utils.prepare_batch(batch_sentences, batch_labels, lens)
                loglik = self.model.loglik(batch_sentences, batch_labels_var, lens_s)
                negative_loglik = -loglik.mean()
                negative_loglik.backward()
                self.optimizer.step()

                if global_step % (self.batch_size * self.n_iter) == 0:
                    logging.info(f'iteration = {global_iter}  negative_loglik={negative_loglik:.4f}')
                    if self.logic_break:
                        logging.info('Early stopping')
                        logging.info('Saving model ...')
                        self.save_model()
                        break
            train_data.dataset.close_file()
            self.on_epoch_complete(epoch_idx, global_iter, global_step)
            if self.logic_break:
                logging.info('Early stopping')
                logging.info('Saving model ...')
                self.save_model()
                break
        logging.info(f'time: {time.time()-start}')
        logging.info(f'iteration ={global_iter}')
        logging.info(f'negative_loglik={negative_loglik:.4f}')
        logging.info('Saving model ...')
        self.save_model()

    def check_loss(self, loss):
        if loss < self.min_dev_loss * self.patience_threshold:
            self.min_dev_loss = loss
            self.patience = 0
        else:
            self.patience += 1
            if self.patience == self.max_patience:
                self.decay_num += 1
                if self.decay_num == self.max_decay_num:
                    self.logic_break = True
                learning_rate = self.optimizer.param_groups[0]['lr'] * self.learning_rate_decay
                for param_group in self.optimizer.param_groups:
                    param_group['lr'] = learning_rate
                self.patience = 0

    def validate(self, steps=None):
        if not self.val:
            return

        path = os.path.join(self.save_dir, 'predict_validation.csv')
        utils.create_save_file(path)

        logging.info('...............................')
        logging.info('Validating...')
        self.model.train(False)
        data_size = 0
        sampled = False

        val_data = self.load_data(train=False)

        nb_classes = len(self.label_vocab)

        confusion_matrix = torch.zeros(nb_classes, nb_classes)

        loss = 0

        for i_idx, (batch, lens, _) in enumerate(val_data):
            i_idx += 1
            batch_size = batch.size(1)
            data_size += batch_size

            sequence_batch, label_batch = batch[:-1], batch[-1]
            with torch.no_grad():
                sequence_batch_wrapped, label_batch_wrapped, lens_s = utils.prepare_batch(sequence_batch, label_batch, lens)
                loglik, logits = self.model.loglik(sequence_batch_wrapped, label_batch_wrapped, lens_s, return_logits=True)
                preds, _, _ = self.model.predict(sequence_batch_wrapped, lens_s)
            loss += -loglik.sum()
            pred_list = preds.data.tolist()
            targets = label_batch_wrapped.data.tolist()
            lens_s = lens_s.data.tolist()
            sequence_batch = sequence_batch_wrapped.data.tolist()
            preds_all = utils.lexicalize_sequence(pred_list, lens_s, self.label_vocab)
            sequence_batch_all = [
                utils.lexicalize_sequence(sequence, lens_s, vocab) for sequence, vocab in zip(
                    utils.tighten(sequence_batch, lens_s), self.inputs_vocab)]
            utils.save_predict(path, sequence_batch_all, preds_all)

            for t, p in zip(label_batch_wrapped.view(-1), preds.view(-1)):
                confusion_matrix[t.long(), p.long()] += 1

            if not sampled and self.samples > 0:
                sample_idx = self.random_idx(batch_size, self.samples)
                sequence_batch = sequence_batch_wrapped.data.tolist()
                sequence_batch = [utils.tighten(x, lens_s) for x in sequence_batch]
                lstm = logits.max(2)[1]
                lstm = lstm.data.tolist()
                lstm = utils.tighten(lstm, lens_s)

                sequence_batch_smp = [self.gather(x, sample_idx) for x in sequence_batch]
                label_batch_smp = self.gather(targets, sample_idx)
                crf_smp = self.gather(pred_list, sample_idx)
                lstm_smp = self.gather(lstm, sample_idx)

                self.display_samples(sequence_batch_smp, label_batch_smp, lstm_smp, crf_smp)
                del sequence_batch, lstm, sequence_batch_smp, label_batch_smp, crf_smp, lstm_smp

                sampled = True

        val_data.dataset.close_file()
        loss = loss/data_size

        self.check_loss(loss)

        confusion_matrix = confusion_matrix[2:, 2:]
        print(confusion_matrix)
        logging.info(f'precision: {confusion_matrix.diag() / confusion_matrix.sum(dim=0)}')
        logging.info(f'recall: {confusion_matrix.diag() / confusion_matrix.sum(dim=1)}')
        logging.info(f'negative_loglik validate {loss:.4f}')
        if self.azureml_logger: 
            self.azureml_logger.save_confusion_matrix_from_tensor(
                confusion_matrix=confusion_matrix,
                labels=list(self.label_vocab.f2i.keys())[2:],
                current_epoch=self.validation_number
            )
            self.azureml_logger.save_metrics(
                confusion_matrix=confusion_matrix,
                labels=list(self.label_vocab.f2i.keys())[2:]
            )
            self.azureml_logger.save_validation_metrics(loss)
        
        self.validation_number += 1