# 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_vdst import QGenModel
from src.tools.optimizer import Optimizer
from src.tools.tokenizer import GW_Tokenizer, BERT_Tokenizer
from src.data.image_features_reader import numpyReader as image_features_reader
from src.data.qgen_vdst import QGenDataset, collate_fn


NUM_LOG_TEXT_SAMPLES = 5 # must < len(valid_set)
NUM_TURNS_FOR_VALIDATION = 2

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

        self.entropy_loss_weight = 0
        if 'entropy_loss_weight' in self.config['hparas']:
            self.entropy_loss_weight = self.config['hparas']['entropy_loss_weight']

    def load_dataloader(self, img_feat_readers, tokenizer, splits):
        # 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 3*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.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, splits)
        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, qs, tf_input, answers, q_len, obj_feats = data
        # Teacher forcing
        tf_input = torch.cat([torch.zeros(qs.size(0), qs.size(1), 1).long(), tf_input], dim=-1)
        tf_input[:, 0, 0] = self.tokenizer.sos_id
        tf_input[:, 1:, 0] = answers
        return (
            game,
            qs.to(self.device),
            tf_input.to(self.device),
            answers.to(self.device),
            q_len.to(self.device),
            obj_feats.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:
        # while self.step < 1e100:
            # Validate every epoch
            self.validate(self.valid_set)
            self.timer.set()
            for data in self.train_set:
                game, qs, tf_input, answers, q_len, obj_feats = self.fetch_data(data)
                self.timer.cnt('rd')
                # Forward
                self.optimizer.pre_step(self.step)
                # (batch_size, max_num_turns, max_q_len, num_classes)             
                pred, _, entropy = self.model.forward_dialog(tf_input, q_len, answers, obj_feats)
                loss = self.loss(pred.reshape(-1, pred.size(-1)), qs.reshape(-1)) + self.entropy_loss_weight * entropy
                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} | Grad. norm - {:.2f} | {}".format(
                        loss.item(), grad_norm, self.timer.show()))            
                    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
        for val_step, data in enumerate(specified_set):
            with torch.no_grad():
                game, qs, tf_input, answers, q_len, obj_feats = self.fetch_data(data)
                pred, _, _ = self.model.forward_dialog(tf_input, q_len, answers, obj_feats)
                loss = self.loss(pred.reshape(-1, pred.size(-1)), qs.reshape(-1))
                total_loss += loss
                if (val_step == 0) or ((val_step+1) % self._progress_step == 0):
                    self.progress("Dev stat. | Loss - {:.4f}".format(
                        total_loss/float(val_step+1)))
                # Log
                if self.mode == 'train':
                    if val_step < NUM_LOG_TEXT_SAMPLES:
                        batch_size = pred.size(0)
                        num_turns = pred.size(1)
                        device = pred.device
                        dial_pred = ""
                        dial_tgt  = ""
                        for t in range(num_turns):
                            dial_pred += self.tokenizer.decode(pred[0][t].argmax(dim=-1).cpu().numpy())
                            dial_tgt  += self.tokenizer.decode(qs[0][t].cpu().numpy())
                        self.write_log('text', '(teacher-forcing)-pred-%d' % val_step, dial_pred)
                        if self.step == 0:
                            self.write_log('text', '(teacher-forcing)-target-%d' % val_step, dial_tgt)


                        last_wrd = torch.zeros(batch_size).fill_(self.tokenizer.sos_id).long().to(device)
                        end_of_dialog = torch.zeros(batch_size).bool().to(device)
                        # for t in range(NUM_TURNS_FOR_VALIDATION):
                        q_tokens, actual_length, _, _, _ = self.model.generate_sentence(
                            last_wrd, obj_feats, self.tokenizer.eoq_id, self.tokenizer.eod_id, end_of_dialog, 
                            max_q_len=12, pi=None, last_state=None, greedy=True)

                        # Take the first one
                        pred_sent = self.tokenizer.decode(q_tokens[0].cpu().numpy())
                        self.write_log('text', '(free-run)-pred-%d' % val_step, pred_sent)
                        
        avg_loss = total_loss / float(len(specified_set))
        self.write_log('scalars', 'loss', {'dev': avg_loss})
        score = -avg_loss
        if score > self.best_score and self.mode == 'train':
            self.save_checkpoint('best.pth', score)
            self.best_score = score

        self.verbose(["Val stat. @ step {} | Loss - {:.4f}"
                      .format(self.step, avg_loss)])

        self.model.train()
