# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0
import os
import torch
import torch.nn as nn
from functools import partial
from matplotlib.image import imread
from solver.solver import BaseSolver
from solver.utils import human_format, cal_hit
from torch.utils.data import DataLoader
from src.model.qgen import QGenModel
from src.tools.optimizer import Optimizer
from src.tools.tokenizer import GW_Tokenizer, BERT_Tokenizer
from src.data.image_features_reader import h5FeatureReader as image_features_reader
from src.data.qgen import QGenDataset, collate_fn


NUM_LOG_TEXT_SAMPLES = 5 # must < len(valid_set)


class QGenSolver(BaseSolver):
    def __init__(self, config, args, mode):
        super().__init__(config, args, mode)
        self.step = 0
        self.best_score = -1e10

    def load_dataloader(self, img_feat_readers, tokenizer):
        # Prepare self.train_set, self.valid_set, self.test_set
        dataroot = self.config['data']['dataroot']
        batch_size = self.config['data']['batch_size']
        splits = ['train', 'valid'] if self.mode == 'train' else ['test', 'valid']
        for split in splits:
            dataset = QGenDataset(
                dataroot, 
                split, 
                img_feat_readers[split], 
                tokenizer, 
                padding_index=tokenizer.pad_id)
            # Set self.XXX_set = torch.utils.data.Dataloader
            setattr(
                self,
                split+'_set',
                DataLoader(
                    dataset,
                    batch_size=batch_size if split == 'train' else 4*batch_size,
                    shuffle=(split=='train'),
                    drop_last=False,
                    collate_fn=partial(collate_fn, wrd_pad_id=tokenizer.pad_id),
                    num_workers=self.args.n_jobs,
                    pin_memory=self.args.pin_memory)
            )

    def load_data(self):
        self.verbose(['Loading data...'])
        config = self.config
        if config['data']['tokenizer'].lower() == 'bert':
            self.verbose(["Use Bert tokenizer"])
            tokenizer = BERT_Tokenizer.initialize()
            #tokenizer = BERT_Tokenizer.from_pretrained(
                #'bert-base-uncased', do_lower_case=True)
        else:
            tokenizer = GW_Tokenizer(config['data']['vocab_path'])
        
        feat_path = config['data']['features_path']
        splits = ['train', 'valid'] if self.mode == 'train' else ['test', 'valid']
        img_feat_readers = {
            split: image_features_reader(feat_path[split]) 
            for split in splits}
        self.load_dataloader(img_feat_readers, tokenizer)
        self.tokenizer = tokenizer
        if self.mode == 'train':
            self.steps_per_epoch = len(self.train_set)
            self.max_step = self.steps_per_epoch * self.max_epoch

    def fetch_data(self, data):
        game, qgen_in, qgen_in_len, qgen_tgt, qgen_tgt_len, img_feat, out_mask = data

        return (
            game,
            qgen_in.to(self.device),
            qgen_in_len.to(self.device),
            qgen_tgt.to(self.device),
            qgen_tgt_len.to(self.device),
            img_feat.to(self.device),
            out_mask.to(self.device),
        )

    def set_model(self):
        self.verbose(['Set model...'])
        self.model = QGenModel(
            num_wrds=len(self.tokenizer),
            wrd_pad_id=self.tokenizer.pad_id,
            **self.config['model']).to(self.device)
        self.optimizer = Optimizer(
            self.model.parameters(), **self.config['hparas'])
        # self.loss = nn.CrossEntropyLoss(reduction='sum')
        self.loss = nn.CrossEntropyLoss(ignore_index=self.tokenizer.pad_id)

        if self.args.load:
            ckpt = torch.load(self.args.load, map_location=self.device)
            self.model.load_state_dict(ckpt['model'])
            self.optimizer.load_opt_state_dict(ckpt['optimizer'])
            self.step = ckpt['global_step']
            self.verbose('Load ckpt from {}, restarting at step {}'.format(
                self.args.load, self.step))

    def exec(self):
        if self.mode == 'train':
            self.train()
        elif self.mode == 'test':
            self.verbose(["Evaluate on test set..."])
            self.validate(self.test_set)
            self.verbose(["Evaluate on valid set..."])
            self.validate(self.valid_set)
        elif self.mode == 'sanity-check':
            self.sanity_check(self.valid_set)
        

    def train(self):
        self.verbose(['Total training epoch/steps: {}/{}'.format(
            self.max_epoch, human_format(self.max_step))])
        self.verbose(['Number of steps per epoch: {}'.format(
            human_format(self.steps_per_epoch))])
        while self.step < self.max_step:
            # Validate every epoch
            self.validate(self.valid_set)
            self.timer.set()
            for data in self.train_set:
                game, qgen_in, qgen_in_len, qgen_tgt, qgen_tgt_len, img_feat, out_mask = self.fetch_data(data)
                self.timer.cnt('rd')
                # Forward
                self.optimizer.pre_step(self.step)
                pred = self.model(qgen_in, qgen_in_len, img_feat, mask=None)
                loss = self.loss(pred.view(-1, pred.size(-1)), qgen_tgt.view(-1))
                hit = cal_hit(pred, qgen_tgt)
                acc = hit / float(qgen_tgt.size(0) * qgen_tgt.size(1))
                hit_nopad = cal_hit(pred, qgen_tgt, pad_id=self.tokenizer.pad_id)
                acc_nopad = hit_nopad / float(qgen_tgt.size(0) * qgen_tgt.size(1))
                self.timer.cnt('fw')
                # Backward
                grad_norm = self.backward(loss)
                self.timer.cnt('bw')

                self.step += 1
                # Log
                if (self.step == 1) or (self.step % self._progress_step == 0):
                    self.progress("Tr stat. | Loss - {:.4f} | Acc.(pad/nopad) - {:.3f}/{:.3f} | Grad. norm - {:.2f} | {}".format(
                        loss.item(), acc, acc_nopad, grad_norm, self.timer.show()))            
                    self.write_log('scalars', 'accuracy', {'train': acc})
                    self.write_log('scalars', 'accuracy', {'train-nopad': acc_nopad})
                    self.write_log('scalars', 'loss', {'train': loss})

                # End of step
                self.timer.set()
                if self.step > self.max_step:
                    self.verbose("Reach max training step.")
                    self.logger.close()
                    break


    def validate(self, specified_set):

        self.model.eval()
        total_loss = 0
        total_hit = 0
        total_hit_nopad = 0
        for val_step, data in enumerate(specified_set):
            game, qgen_in, qgen_in_len, qgen_tgt, qgen_tgt_len, img_feat, out_mask = self.fetch_data(data)
            with torch.no_grad():
                pred = self.model(qgen_in, qgen_in_len, img_feat, mask=None)
                loss = self.loss(pred.view(-1, pred.size(-1)), qgen_tgt.view(-1)).item()
                hit = cal_hit(pred, qgen_tgt)
                hit /= qgen_tgt.size(0) * qgen_tgt.size(1)
                hit_nopad = cal_hit(pred, qgen_tgt, pad_id=self.tokenizer.pad_id)
                hit_nopad /= qgen_tgt.size(0) * qgen_tgt.size(1)

                total_hit += hit
                total_hit_nopad += hit_nopad
                total_loss += loss
                if (val_step == 0) or ((val_step+1) % self._progress_step == 0):
                    self.progress("Dev stat. | Loss - {:.4f} | Acc.(pad/nopad) - {:.3f}/{:.3f}".format(
                        total_loss/float(val_step+1), 
                        total_hit/float(val_step+1), 
                        total_hit_nopad/float(val_step+1)))
                # Log
                if self.mode == 'train':
                    if val_step < NUM_LOG_TEXT_SAMPLES:
                        pred = self.tokenizer.decode(pred[0].argmax(dim=-1).cpu().numpy())
                        tgt = self.tokenizer.decode(qgen_tgt[0].cpu().numpy())
                        self.write_log('text', '(teacher-forcing)-pred-%d' % val_step, pred)
                        if self.step == 0:
                            self.write_log('text', '(teacher-forcing)-target-%d' % val_step, tgt)

                        qgen_in_sc, qgen_in_len_sc, last_ans = self.make_sanity_check_input(qgen_in)
                        pred_sc = self.model.generate_from_dialog(
                            qgen_in_sc, qgen_in_len_sc, last_ans, img_feat,
                            self.tokenizer.eoq_id, self.tokenizer.eod_id, max_q_len=20)
                        # Take the first one
                        qgen_in_sc = self.tokenizer.decode(qgen_in_sc[0].cpu().numpy())
                        pred_sc = self.tokenizer.decode(pred_sc[0].cpu().numpy())
                        self.write_log('text', '(free-run)-input-%d' % val_step, qgen_in_sc)
                        self.write_log('text', '(free-run)-pred-%d' % val_step, pred_sc)
                        
        # Total hit is not actaul hit here
        avg_acc = total_hit / float(len(specified_set))
        avg_acc_nopad = total_hit_nopad / float(len(specified_set))
        avg_loss = total_loss / float(len(specified_set))
        self.write_log('scalars', 'accuracy', {'dev': avg_acc})
        self.write_log('scalars', 'accuracy', {'dev-nopad': avg_acc_nopad})
        self.write_log('scalars', 'loss', {'dev': avg_loss})
        score = -avg_loss
        if score > self.best_score and self.mode == 'train':
            #self.save_checkpoint('step_{}.pth'.format(self.step), score)
            self.save_checkpoint('best.pth', score)
            self.best_score = score

        self.verbose(["Val stat. @ step {} | Loss - {:.4f} | Acc.(pad/nopad) - {:.3f}/{:.3f}"
                      .format(self.step, avg_loss, avg_acc, avg_acc_nopad)])

        self.model.train()

    def sanity_check(self, specified_set):
        self.model.eval()
        out_file = open(self.args.result, 'w')
        out_file.write("game_id|predicted_question|dialog_history\n")
        out_format = "{}|{}|{}\n"

        for val_step, data in enumerate(specified_set):
            self.progress("[{}/{}] Generating questions for sanity check...".format(val_step+1, len(specified_set)))
            game, qgen_in, qgen_in_len, qgen_tgt, qgen_tgt_len, img_feat, out_mask = self.fetch_data(data)
            qgen_in_sc, qgen_in_len_sc, last_ans = self.make_sanity_check_input(qgen_in)
            with torch.no_grad():
                pred = self.model.generate_from_dialog(
                    qgen_in_sc, qgen_in_len_sc, last_ans, img_feat,
                    self.tokenizer.eoq_id, self.tokenizer.eod_id, max_q_len=20)
                
                for b in range(len(game)):
                    gid = game[b].id
                    q_pred = self.tokenizer.decode(pred[b].tolist())
                    dial_his = self.tokenizer.decode(qgen_in_sc[b][:qgen_in_len_sc[b]].tolist() + [last_ans[b].item()])
                    out_str = out_format.format(gid, q_pred, dial_his)
                    out_file.write(out_str)
            # break

        out_file.close()


    def make_sanity_check_input(self, qgen_in):
        import random
        from torch.nn.utils.rnn import pad_sequence
        batch_size = qgen_in.size(0)
        device = qgen_in.device
        qgen_in_sc = []
        qgen_in_len_sc = []
        last_ans = []
        for b in range(batch_size):
            potential = (qgen_in[b] == self.tokenizer.yes_id) |\
                        (qgen_in[b] == self.tokenizer.no_id)  |\
                        (qgen_in[b] == self.tokenizer.na_id)

            indices = potential.nonzero().view(-1)
            mask = torch.arange(1, len(qgen_in[b])+1).to(device)
            potential_input_length = mask[indices].tolist()
            input_length = random.sample(potential_input_length, k=1)[0]
            qgen_in_sc.append(qgen_in[b][:input_length-1])
            qgen_in_len_sc.append(input_length-1)
            last_ans.append(qgen_in[b][input_length-1])
        qgen_in_sc = pad_sequence(qgen_in_sc, batch_first=True, padding_value=self.tokenizer.pad_id)
        qgen_in_len_sc = torch.LongTensor(qgen_in_len_sc)
        last_ans = torch.stack(last_ans)
        return qgen_in_sc.to(device), qgen_in_len_sc.to(device), last_ans.to(device)