# 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 solver.solver import BaseSolver
from solver.utils import human_format
from torch.utils.data import DataLoader
from src.model.guesser import GuesserModel
from src.tools.optimizer import Optimizer
from src.tools.tokenizer import GW_Tokenizer, BERT_Tokenizer
# from src.data.image_features_reader import ImageFeaturesH5Reader
from src.data.guesser import GuesserDataset, collate_fn


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

    def load_dataloader(self, image_features_reader, 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 = GuesserDataset(
                dataroot, split, image_features_reader, tokenizer, padding_index=tokenizer.pad_id)
            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'])
        image_features_reader = None
        self.load_dataloader(image_features_reader, 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):
        dialog, dialog_len, cats, bboxs, bboxs_mask, label = data
        dialog = dialog.to(self.device)
        dialog_len = dialog_len.to(self.device)
        cats = cats.to(self.device)
        bboxs = bboxs.to(self.device)
        bboxs_mask = bboxs_mask.to(self.device)
        label = label.to(self.device)
        return dialog, dialog_len, cats, bboxs, bboxs_mask, label

    def set_model(self):
        self.verbose(['Set model...'])
        self.model = GuesserModel(
            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()

        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):
        # [TODO]: finish testing code
        if self.mode == 'train':
            self.train()
        else:
            self.verbose(["Evaluate on test set..."])
            self.validate(self.test_set)
            self.verbose(["Evaluate on valid set..."])
            self.validate(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:
                dialog, dialog_len, cats, bboxs, bboxs_mask, label = self.fetch_data(data)

                self.timer.cnt('rd')
                # Forward
                self.optimizer.pre_step(self.step)
                pred = self.model(dialog, dialog_len, cats, bboxs, bboxs_mask)
                loss = self.loss(pred, label)
                acc = (pred.argmax(dim=-1) == label).sum().item() / float(len(label))
                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. - {:.3f} | Grad. norm - {:.2f} | {}".format(
                        loss.item(), acc, grad_norm, self.timer.show()))

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


    def validate(self, specified_set):
        self.model.eval()
        total_loss = 0
        total_hit = 0
        cnt = 0
        for val_step, data in enumerate(specified_set):
            dialog, dialog_len, cats, bboxs, bboxs_mask, label = self.fetch_data(data)
            with torch.no_grad():
                pred = self.model(dialog, dialog_len, cats, bboxs, bboxs_mask)
                loss = self.loss(pred, label).item()
                hit = (pred.argmax(dim=-1) == label).sum().item()
                total_hit += hit
                total_loss += loss
                cnt += len(label)
                print(val_step)
                if (val_step == 0) or ((val_step+1) % self._progress_step == 0):
                    self.progress("Dev stat. | Loss - {:.4f} | Acc. - {:.4f}".format(
                        total_loss/float(cnt), total_hit/float(cnt)))


        score = total_hit / float(cnt)
        loss = total_loss / float(len(specified_set))
        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} | Acc. - {:.4f}"
                      .format(self.step, loss, score)])

        self.model.train()

    def write_log(self):
        pass










