# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0
import os
import sys
import abc
import math
import torch
from shutil import copyfile
# from tensorboardX import SummaryWriter
from torch.utils.tensorboard import SummaryWriter
from datetime import datetime
from solver.utils import Timer, human_format

INFO_MARK = "[INFO]"
TB_FLUSH_FREQ = 180
PROGRESS_STEP = 20
CLIP_GRAD_NORM = 5.0


class BaseSolver(object):
    def __init__(self, config, args, mode):
        self.config = config
        self.args = args
        if self.args.gpu and torch.cuda.is_available():
            self.device = torch.device('cuda', args.local_rank)
        else:
            self.device = torch.device('cpu')
        print('the self.device is using')
        print(self.device)
        self.exp_name = args.name
        if self.exp_name is None:
            self.exp_name = '-'.join(
                [args.config.split('/')[-1].replace('.yaml',''),'sd'+str(args.seed)])
        self.distributed = args.distributed
        self.main_proc = args.local_rank == 0
        # Filepath setup
        self.ckptdir = os.path.join(args.result, self.exp_name)
        self.logdir = os.path.join(args.logdir, self.exp_name)
        print('self.logdir')
        print(self.logdir)
        self._clip_grad_norm = self.config['hparas'].get('clip_grad_norm', CLIP_GRAD_NORM)
        self._progress_step = PROGRESS_STEP
        self.mode = mode
        if mode == 'train':
            os.makedirs(self.ckptdir, exist_ok=True)
            copyfile(args.config, os.path.join(self.ckptdir, 'config.yaml'))
            self.logger = SummaryWriter(self.logdir, flush_secs=TB_FLUSH_FREQ)
            self.timer = Timer()
            # Hyperparameters
            self.step = 0
            self.max_epoch = config['hparas']['max_epoch']

        self.verbose(['Experiment: %s' % self.exp_name])

    @abc.abstractmethod
    def load_data(self):
        raise NotImplementedError

    @abc.abstractmethod
    def set_model(self):
        raise NotImplementedError

    @abc.abstractmethod
    def exec(self):
        raise NotImplementedError

    @abc.abstractclassmethod
    def fetch_data(self, data):
        raise NotImplementedError

    def verbose(self, msg):
        '''
        Verbose function for print information to stdout
            <list> msg - List of <str> for verbosity
        '''
        if self.args.verbose:
            self._clean_line()
            t = datetime.now().strftime("%H:%M:%S")
            prefix = INFO_MARK + "[%s]" % t
            if type(msg) is str:
                print(prefix, msg)
            else:
                for m in msg:
                    print(prefix, m)

    def progress(self, msg):
        '''
        Verbose function for updating progress on stdout
            <str> msg - String to show as progress bar, do not include "new line"
        '''
        if self.args.verbose:
            self._clean_line()
            print('[{}] {}'.format(human_format(self.step), msg), end='\r')

    def _clean_line(self):
        sys.stdout.write("\033[K")

    def save_checkpoint(self, filename, score):
        ''''
        Ckpt saver
            <str> filename  - the name of ckpt file (w/o prefix) to store, overwrite if existed
            <float> score - The value of metric used to evaluate model
        '''
        ckpt_path = os.path.join(self.ckptdir, filename)
        full_dict = {
            #"model": self.model.module.state_dict() if self.distributed else self.model.state_dict(),
            "model": self.model.state_dict(),
            "optimizer": self.optimizer.get_opt_state_dict(),
            "global_step": self.step,
        }
        torch.save(full_dict, ckpt_path)
        self.verbose("Saved checkpoint (step = {}, score = {:.4f}) and status @ {}".\
                                       format(human_format(self.step), score, ckpt_path))

    def backward(self, loss):
        '''
        Standard backward step with self.timer and debugger
            <torch> loss - the loss to perform loss.backward()
        '''
        self.timer.set()
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(
            self.model.parameters(), self._clip_grad_norm)
        if math.isnan(grad_norm):
            self.verbose('Error : grad norm is NaN @ step ' + str(self.step))
        else:
            self.optimizer.step()
        return grad_norm

    def write_log(self, log_type, name, vals):
        if log_type == 'text':
            self.logger.add_text(name, vals, self.step)
        else:
            self.logger.add_scalars(name, vals, self.step)
        self.logger.flush()
        
