'''


'''
import logging
import os
import shutil
from copy import deepcopy as dcopy

import numpy as np

import torch
import torch.optim as optim
import torch.utils.data as D

import tools # since T is used as a variable in this module, refrain from "import tools as T"
import tools.torch
import tools.modules

from . import utils as U
from . import evaluation as E
from . import model as M

# __all__ = [
# 'Trainer'
# ]

# %%
log = logging.getLogger(__name__)

# %%
def validate_cfg_val(cfg_val):
    cfg_val = dcopy(dict(cfg_val)) if cfg_val is not None else {}
    if 'forward_f' not in cfg_val:
        cfg_val['forward_f'] = None
    return cfg_val

class Trainer:
    def __init__(self, network, train_dataset=None, cfg_train={}, criterion=None, network_dir='network', cfg_val=None, val_dataset=None, val_metrics=None, name='default', verbose=True, amp=False, reproduce=True):
        """
        - train network with early stopping, given hyperparameters
        - save/load network with necessary preprocessing pipelin

        Parameters
        ----------
        network : torch.nn.Module object
            Network to train

        train_dataset: torch.utils.data.Dataset object
            This dataset is used for training.
            Recommended return format from dataset is dictionary form, although tuples or raw tensors are supported.
            Example: {'x': x, 'y': y}
            
            train_dataset is fed into torch.utils.data.DataLoader object, where it is batched then passed to self.forward(), with the following pattern depending on train_dataset return type:
                torch.Tensor: self.forward(data)
                tuple or list: self.forward(*data)
                dict: self.forward(**data)

        val_dataset: torch.utils.data.Dataset object or dict of such objects
            Dataset used for validation (Early stopping)
            forward mechanism from val_dataset may be different from train_dataset, which could be defined in cfg_val['forward_f']        

        cfg_train: dict-like object with following keys:
            lr: float
            batch_size: int
            weight_decay: bool (optional)
            cv_step: int (optional, recommended when val_dataset is given)
        
        cfg_val: (required if val_dataset is given) dict-like object with following keys:
            increase_better: Bool
            patience: int
            batch_size: int

            forward_f: callable, forward pass that might be different from training forward pass (optional)
            target_dataset: str (optional, if val_dataset is dict for multiple datasets)
            criterion: str (optional, if val_metrics is dict for multiple metrics) 
        
        criterion: nn.Module object
            Defines the loss function.
            If parameters exist within criterion, it is moved to the corresponding device where network lies.
        
        network_dir: str or Pathlike object, default='network'
            The directory to save the nn.
            If early stopping is used, a subdirectory will be created that keeps the models during training.

        Examples
        --------
        # Recommended way of making hs.Trainer object is like the following:

        kwargs = {'network': Network(), 'train_dataset': train_dataset, 'validation': None, 'cfg_train': cfg.train,
                'criterion': criterion, 'network_dir': path['network'], 'verbose': True, 'amp': True}
        trainer = hs.Trainer(**kwargs)

        # To make a custom Trainer for different neural network applications, just define the forward function:

        class MyTrainer(hs.Trainer):
            def forward(self, x, y, z):
                w = self.network(x,y)
                loss = self.criterion(w, z)
                return loss
        """

        # Store configurations
        self.network = network
        self.train_dataset = train_dataset
        self.cfg_train = dcopy(dict(cfg_train))
        self.criterion = criterion
        self.network_dir = network_dir

        self.cfg_val = validate_cfg_val(cfg_val)
        self.val_dataset = val_dataset
        self.val_metrics = val_metrics
        self.earlystopper = EarlyStopper(increase_better=cfg_val['increase_better'], patience=cfg_val['patience'], save_dir=os.path.join(self.network_dir, 'network_temp')) if self.val_dataset is not None else None
        if 'target_dataset' in self.cfg_val:
            if type(self.val_dataset)==dict:
                assert self.cfg_val['target_dataset'] in self.val_dataset, f'[Validation setup] target_dataset: {self.cfg_val["target_dataset"]} not found in val_dataset: {self.val_dataset.keys()}'
            assert ~((self.val_dataset is not None) ^ (self.val_metrics is not None)), f'val_dataset and val_metrics must be given together or neither. val_dataset: {self.val_dataset}, val_metrics: {self.val_metrics}'

        self.name = name
        self.verbose = verbose
        self.amp = amp
        self.reproduce = reproduce

        # Initializations
        if self.network_dir is not None:
            if os.path.exists(self.network_dir): log.warning(f'path: {self.network_dir} exists. Be careful')
            os.makedirs(self.network_dir, exist_ok=True)

        self.train_meter = tools.modules.AverageMeter() # Tracks loss per epoch
        self.loss_tracker = tools.modules.ValueTracker() # Tracks loss over all training
        self.set_model()
        self.reset()

        self._targets_type = self.infer_targets_type()

    def reset(self):
        self.iter = 0
        self.n_batch=0

    def print(self, content):
        if self.verbose:
            print(content)
        else:
            log.info(content)

    def set_model(self):
        '''Set miscellaneous variables'''
        self.model = M.BaseModel(self.network, amp=self.amp)
        self.misc = tools.TDict()
        if self.train_dataset is not None:
            if hasattr(self.train_dataset, 'get_f'): # get_f equal to preprocessing pipeline before feeding to network
                self.print('get_f found in train_dataset')
                self.misc.get_f = self.train_dataset.get_f

    def validate(self):
        patience_end = False
        self.print(f'[Node: {self.name}] Validation')
        if self.val_dataset is not None:
            if type(self.val_dataset)==dict:
                scores = {}
                for name, dataset in self.val_dataset.items():
                    results = E.forward_model(self.model, dataset, batch_size=self.cfg_val['batch_size'], forward_f=self.cfg_val['forward_f'], amp=self.amp)
                    isnan = any([np.isnan(v).any() for v in results.values()]) if isinstance(results, dict) else np.isnan(np.concatenate(results)).any()
                    if isnan:
                        log.warning(f'[Node: {self.name}][validate] NaN found in results. Skipping validation...')
                        score = {self.cfg_val['criterion']: self.earlystopper.best_score}
                    else:
                        score = E.evaluate(results, self.val_metrics)
                    scores[name] = score
                    self.print(f'[dataset_name: {name}] Validation Score: {score}')
                score = scores[self.cfg_val['target_dataset']]
            else:
                results = E.forward_model(self.model, self.val_dataset, batch_size=self.cfg_val['batch_size'], forward_f=self.cfg_val['forward_f'], amp=self.amp)
                isnan = any([np.isnan(v).any() for v in results.values()]) if isinstance(results, dict) else np.isnan(np.concatenate(results)).any()
                if isnan:
                    log.warning(f'[Node: {self.name}][validate] NaN found in results. Skipping validation...')
                    score = {self.cfg_val['criterion']: self.earlystopper.best_score}
                else:
                    score = E.evaluate(results, self.val_metrics)
                self.print(f'Validation Score: {score}')
            
            if self.earlystopper is not None:
                if type(score) is dict:
                    score_ = score[self.cfg_val['criterion']]
                else:
                    score_ = score
                    
                patience_end = self.earlystopper.step(self, score_)
                if patience_end:
                    self.print('patience met, flushing earlystopper history')
                    self.earlystopper.history.reset()
                    
        torch.cuda.empty_cache()
        self.network.train()
        return patience_end

    def generate_loader(self):
        if self.train_dataset is not None:
            reproduce_kwargs = U.reproducible_worker_dict() if self.reproduce else {}
            drop_last = len(self.train_dataset) % self.cfg_train['batch_size'] == 1 # To avoid batch normalization layers from raising exceptions
            self.loader = D.DataLoader(self.train_dataset, batch_size=self.cfg_train['batch_size'], shuffle=True, drop_last=drop_last, **reproduce_kwargs)
        else:
            log.warning('No train_dataset found. Skipping generate_loader()')

    def train(self, epoch=None, new_op=True, no_val=False, step=None, reset_loss_tracker=False, device=None):
        """
        Trains the network with the specified duration.

        Parameters
        ----------
        epoch : int, default=None
            Number of epochs to train the network.
            If None, the value is taken from self.cfg_train['epoch']
            Must be mutually exclusive with "step" argument.

        new_op : bool, default=True
            Whether to create a new optimizer.
            If False, the optimizer is not initialized, keeping the previous state.

        no_val : bool, default=False
            Whether to skip validation during training.
            If True, the validation is skipped even when val_dataset is given.
        
        step : int, default=None
            Number of steps to train the network.
            Must be mutually exclusive with "epoch" argument.
        
        reset_loss_tracker : bool, default=False
            Whether to reset the loss_tracker before training.

        device : torch.device, default=None
            Explicity Device to train the network, recommended to default to None. 
            If None, the device is taken from the network.

        Returns
        -------
        loss_tracker : tools.modules.ValueTracker

        """
        assert epoch is None or step is None, f'only one of epoch or step can be specified. received epoch: {epoch}, step: {step}'
        if step is None:
            horizon = 'epoch'
            if epoch is None: # When neither epoch or step are specified
                assert 'epoch' in self.cfg_train, 'key "epoch" must be provided in cfg_train, or the argument "epoch" must be provided \
                                                when argument "step" is not specified.'
                epoch = self.cfg_train['epoch']
            log.debug(f'[Node: {self.name}] train for {epoch} epochs')
        else:
            assert epoch is None, f'Either epoch or step must be specified. Received epoch: {epoch}, step: {step}'
            horizon = 'step'
            step = step
            log.debug(f'[Node: {self.name}] train for {step} steps')

        self._device = device if device is not None else tools.torch.get_device(self.network)
        self.criterion = self.criterion.to(self._device)
        if reset_loss_tracker: self.loss_tracker.reset()
        self.generate_loader()
        '''
        There may be one or more loaders, but self.loader is the standard of synchronization with batches.
        Either return multiple values from dataset, or modify self.forward to use other loaders.
        '''

        if self.val_dataset is not None:
            if 'cv_step' not in self.cfg_train:
                self.print('val_dataset is given but "cv_step" not specified in cfg_train. Defaults to 1 epoch')
                self.cfg_train['cv_step'] = len(self.loader)
            self.validate() # initial testing

        # Make new optimizer
        if new_op or not hasattr(self, 'op'):
            if not hasattr(self, 'op') and not new_op:
                log.warning("new_op=False when there's no pre-existing optimizer. Ignoring new_op...")
            
            l_kwargs = ['lr', 'weight_decay', 'amsgrad']
            kwargs = {key: self.cfg_train[key] for key in l_kwargs if key in self.cfg_train}
            self.op = optim.Adam(self.network.parameters(), **kwargs)

        if horizon == 'epoch':
            self._update_epoch(T=epoch, no_val=no_val)
        elif horizon=='step':
            self._update_step(T=step, no_val=no_val)

        # TODO: Return criterion back to its original device, meaning we have to store its previous device info
        self.criterion = self.criterion.cpu()
        return self.loss_tracker

    def _update_epoch(self, T, no_val=False):

        _iter = 0
        for epoch in range(1, T+1):
            self.train_meter.reset()
            self.epoch_f()
            for batch_i, data in enumerate(self.loader, 1):
                self.iter += 1
                _iter += 1
                loss = self.update(data)

                self.print(f'[Node: {self.name}][iter_sum: {self.iter}][Epoch: {epoch}/{T}][Batch: {batch_i}/{len(self.loader)}][Loss: {loss:.7f} (Avg: {self.train_meter.avg:.7f})]')
                self.loss_tracker.step(self.iter, loss)

                # Validation
                if (self.val_dataset is not None) and (not no_val) and (_iter % self.cfg_train['cv_step']==0):
                    patience_end = self.validate()
                    if patience_end: # If patience has reached, stop training
                        self.print('Patience met, stopping training')
                        return None # return None since double break is impossible in python

    def _update_step(self, T, no_val=False):
        if hasattr(self, '_loader_inst'): del self._loader_inst # Load data from the 1st batch, start fresh

        for _iter in range(1, T+1):
            # Get Data
            try:
                if hasattr(self, '_loader_inst'):
                    data = next(self._loader_inst)
                    self.n_batch += 1
                else:
                    raise StopIteration
            except StopIteration as e:
                self.train_meter.reset()
                self.epoch_f()
                self._loader_inst = iter(self.loader)
                data = next(self._loader_inst)
                self.n_batch = 1

            self.iter += 1
            loss = self.update(data)

            self.print(f'[iter_sum: {self.iter}][Iter: {_iter}/{T}][Batch: {self.n_batch}/{len(self.loader)}][Loss: {loss:.6f} (Avg: {self.train_meter.avg:.6f})]')
            self.loss_tracker.step(self.iter, loss)

            # Validation
            if (self.val_dataset is not None) and (not no_val) and (_iter % self.cfg_train['cv_step']==0):
            # if (not no_val) and (self.iter % self.cfg_train.cv_step==0):
                patience_end = self.validate()
                if patience_end: # If patience has reached, stop training
                    self.print('Patience met, stopping training')
                    return None # return None since double break is impossible in python

    def update(self, data):
        """
        - Perform single update (forward/backward pass + gradient descent step) with the given data.
        - Store loss in self.train_meter
        - This is where a lot of errors happen, so there's a pdb to save time.
          When there's error, use pdb to figure out the shape, device, tensor dtype, and more.

        Parameters
        ----------
        data : tuple, list, or dict of tensors (Batch of data)
            This is received from a torch.utils.Data.DataLoader
            Depending on the given format, the data is fed to the forwad pass

        Returns
        -------
        loss : float
        """
        # Automatic Mixed Precision (AMP, float32 -> float16) for acceleration
        if self.amp:
            with torch.autocast(device_type=self._device.type):
                loss, N = self._forward(data)
        else:
            loss, N = self._forward(data)

        self.op.zero_grad()
        loss.backward()
        self.op.step()

        loss = loss.item()
        self.train_meter.step(loss, N)
        return loss

    def _forward(self, data): # TODO: change to function decorator?
        '''
        Pseudo function to:
        1. Move data to self._device
        2. Support passing tuple or dict from a batch from dataloader to forward()
        '''
        datatype = type(data)
        # When data is given as a tuple/list
        if datatype is tuple or datatype is list:
            data = [x.to(self._device) for x in data]
            N = len(data[0]) # number of data in batch
            loss = self.forward(*data)

        # When data is given as a dict
        elif datatype is dict:
            data = {key: value.to(self._device) for key, value in data.items()}
            N = len(next(iter(data.values()))) # number of data in batch
            loss = self.forward(**data)

        elif datatype is torch.Tensor:
            data = data.to(self._device)
            N = len(data)
            loss = self.forward(data)
            
        else:
            raise Exception(f'return type from dataset must be one of [torch.Tensor, tuple, list, dict], received: {datatype}')
        return loss, N

    def forward(self, x, y):
        """
        Forward pass: Receive data and pass values to criterion.
        Define forward for custom trainers.

        May return tuple or dictionary, whichever will be feeded to criterion.

        Defaults to args[0] feeding to self.network (assuming x), and args[1] feeding to self.criterion (assuming y)
        """
        
        y_hat = self.network(x)
        loss = self.criterion(y_hat, y)

        return loss

    def epoch_f(self):
        # TODO: rename to forward hook? epoch_hook?
        '''function to call every other epoch. May be used in child class'''
        pass

    def infer_targets_type(self):
        if self.train_dataset is not None and hasattr(self.train_dataset, 'targets_type'):
            return self.train_dataset.targets_type
        elif self.val_dataset is not None:
            if type(self.val_dataset)==dict:
                for name, dataset in self.val_dataset.items():
                    if hasattr(dataset, 'targets_type'):
                        return dataset.targets_type
            elif hasattr(self.val_dataset, 'targets_type'):
                return self.val_dataset.targets_type
        else:
            return None

    def post_train(self, val_dataset=None):
        if val_dataset is None:
            assert self.val_dataset is not None, f'val_dataset must be given if self.val_dataset is None'
            if type(self.val_dataset)==dict:
                val_dataset = self.val_dataset[self.cfg_val['target_dataset']]
            else:
                val_datset = self.val_dataset

        if self._targets_type == 'binary':
            self.print(f'[Node: {self.name}][post_train] binary classfication: choose threshold based on validation set')
            self.threshold = E.binary_threshold(self, val_dataset)
        else:
            self.print(f'[Node: {self.name}][post_train] no post-train procedure for _targets_type: {self._targets_type}')

    '''
    Save/Load functions
    state_dict()
    load_state_dict()

    save()
    save_best()
    load()
    '''
    def state_dict(self):
        state_dict = {
        'loss_tracker': self.loss_tracker,
        }
        if 'get_f' in self.misc:
            state_dict['misc.get_f'] = self.misc.get_f
        if hasattr(self, 'threshold'):
            state_dict['threshold'] = self.threshold
        if self.earlystopper is not None:
            state_dict['earlystopper.history'] = self.earlystopper.history
        return state_dict

    def load_state_dict(self, state_dict):
        state_dict = dcopy(state_dict)

        if 'misc.get_f' in state_dict.keys():
            self.print(f'Updating get_f: {state_dict["misc.get_f"]}')
            self.misc.get_f = state_dict['misc.get_f']
            if self.train_dataset is not None:
                self.print('updating get_f to loader...')
                self.train_dataset.get_f = self.misc.get_f
            del state_dict['misc.get_f']

        if 'threshold' in state_dict.keys():
            self.print(f'Updating threshold: {state_dict["threshold"]}')
            self.threshold = state_dict['threshold']
            del state_dict['threshold']

        self.__dict__.update(state_dict) # update attributes
        
    def load_best_model(self):
        if self.earlystopper is not None:
            if self.earlystopper.best_network != None:
                self.print(f'Loading best network: {self.earlystopper.best_network}')
                self.network.load_state_dict(torch.load(self.earlystopper.best_network))
            else:
                self.print(f'No best_network found in self.earlystopper. No network loaded.')
        else:
            self.print('self.earlystopper not defined. No network loaded.')

    def save(self, path=None, best=True):
        '''
        Save the following:
        state_dict() -> path/state_dict.p
        network.state_dict() -> path/network.pt
        '''
        path = self.network_dir if path is None else path
        self.print(f'[Node: {self.name}]')

        state_dict_path = os.path.join(path, 'model.p')
        state_dict = self.state_dict()
        tools.save_pickle(state_dict, state_dict_path)
        self.print(f'[save] Saving Model info to: {state_dict_path}')
        self.print(f'[save] state_dict: {list(state_dict.keys())}')

        network_path = os.path.join(path, 'network.pt')
        if best:
            if self.earlystopper is None:
                self.print(f'[save][best: {best}] earlystopper not defined. Saving current network -> [{network_path}]')
                torch.save(self.network.state_dict(), network_path)
            else:
                if self.earlystopper.best_network != None:
                    self.print(f'[save][best: {best}] Saving [{self.earlystopper.best_network}] -> [{network_path}]')
                    shutil.copy(self.earlystopper.best_network, network_path)
                else:
                    self.print(f'[save][best: {best}] no best_network in self.earlystopper, Saving current network -> [{network_path}]')
                    torch.save(self.network.state_dict(), network_path)
        else:
            self.print(f'[save][best: {best}] Saving network to: {network_path}')
            torch.save(self.network.state_dict(), network_path)

    def load(self, path=None):
        '''
        load state_dict() and network:
        path/node.p -> state_dict
        path/network.pt -> network.state_dict
        '''
        path = self.network_dir if path is None else path
        self.print(f'[Node: {self.name}]')

        state_dict_path = os.path.join(path, 'node.p')
        self.print(f'[load] Loading from path: {state_dict_path}')
        state_dict = tools.load_pickle(state_dict_path)
        self.print(f'[load] Updating: {list(state_dict.keys())}')
        self.load_state_dict(state_dict)

        network_path = os.path.join(path, 'network.pt')
        self.print(f'[load] Loading network from: {network_path}')
        self.network.load_state_dict(torch.load(network_path))

class EarlyStopper():
    """
    Returns True if the metric has stopped improving for the given number of patience steps.
    Saves network in the given path and remembers which network is the best.

    Parameters
    ----------
    arg : array-like of shape (n_samples,), default=None
        Argument explanation.
        If ``None`` is given, those that appear at least once
        .. versionadded:: 0.18
    
    Functions
    -------
    step()
        returns True/False
    """

    def __repr__(self):
        return f'<EarlyStopper>\npatience: {self.patience}\nincrease_better: {self.increase_better}'

    def __init__(self, increase_better, patience, save_dir='network_temp', discard_best=True):
        self.increase_better = increase_better
        self.patience = patience
        self.save_dir = save_dir
        self.discard_best = discard_best

        os.makedirs(self.save_dir, exist_ok=True)
        self.history = tools.modules.ValueTracker() # For now, only track single score
        self.reset()

    def step(self, trainer, score, filename=None):
        filename = f'network_{trainer.iter}.pt' if filename is None else filename

        self.history.step(trainer.iter, score)

        # Save best network
        if isbetter(score_best=self.best_score, score=score, increase_better=self.increase_better):
            path = os.path.join(self.save_dir, filename)

            log.info(f'[EarlyStopping] New best score: {self.best_score} -> {score}')
            log.info(f'[EarlyStopping] Saved network: {path}')
            torch.save(trainer.network.state_dict(), path)
            if self.discard_best and self.best_network!=None:
                os.remove(self.best_network)

            self.best_score = score
            self.best_network = path

            # Clear patience tracking
            self.history.reset()

        # patience_end == True when best network did not appear for self.patience times
        patience_end = self.patience == len(self.history)
        log.info(f'[EarlyStopping][patience: {len(self.history)}/{self.patience}]')

        return patience_end

    def reset(self):
        # log.info('[EarlyStopper] Resetting...')
        self.history.reset()
        self.best_network = None
        if self.increase_better:
            self.best_score = -float('inf')
        else:
            self.best_score = float('inf')

def isbetter(score_best: float, score: float, increase_better: bool):
    return (not increase_better and score<score_best) or (increase_better and score>score_best)
