"""Utilities to train a model."""

from abc import ABCMeta, abstractmethod
import argparse
import os
import re
import signal

import torch as th

from .utils import get_logger


LOG = get_logger(__name__)


__all__ = ["ModelInterface", "Trainer", "Checkpointer", "BasicArgumentParser"]


class BasicArgumentParser(argparse.ArgumentParser):
    """A basic argument parser with commonly used training options."""

    def __init__(self, *args, **kwargs):
        super(BasicArgumentParser, self).__init__(*args, **kwargs)

        self.add_argument("--data", required=True, help="path to the training data.")
        self.add_argument("--val_data", help="path to the validation data.")
        self.add_argument("--config", help="path to a config file.")
        self.add_argument("--checkpoint_dir", required=True,
                          help="Output directory where checkpoints are saved")
        self.add_argument("--init_from", help="path to a checkpoint from which to try and initialize the weights.")

        self.add_argument("--lr", type=float, default=1e-4,
                          help="Learning rate for the optimizer")
        self.add_argument("--bs", type=int, default=4, help="Batch size")
        self.add_argument("--num_epochs", type=int,
                          help="Number of epochs to train for")
        self.add_argument("--num_worker_threads", default=4, type=int,
                          help="Number of threads that load data")

        # self.add_argument("--experiment_log",
        #                   help="csv file in which we log our experiments")

        self.add_argument("--cuda", action="store_true",
                          dest="cuda", help="Force GPU")
        self.add_argument("--no-cuda", action="store_false",
                          dest="cuda", help="Force CPU")

        self.add_argument("--server", help="Visdom server url")
        self.add_argument("--base_url", default="/", help="Visdom base url")
        self.add_argument("--env", default="main", help="Visdom environment")
        self.add_argument("--port", default=8097, type=int,
                          help="Visdom server port")

        self.add_argument('--debug', dest="debug", action="store_true")

        self.set_defaults(cuda=th.cuda.is_available(), debug=False)


class ModelInterface(metaclass=ABCMeta):
    """An adapter to run or train a model."""

    def __init__(self):
        pass

    @abstractmethod
    def training_step(self, batch):
        """Training step given a batch of data.

        This should implement a forward pass of the model, compute gradients,
        take an optimizer step and return useful metrics and tensors for
        visualization and training callbacks. 

        Args:
          batch (dict): batch of data provided by a data pipeline.

        Returns:
          train_step_data (dict): a dictionary of outputs.
        """
        return {}

    def init_validation(self):
        """Initializes the quantities to be reported during validation.

        The default implementation is a no-op

        Returns:
          data (dict): initialized values
        """
        LOG.warning("Running a ModelInterface validation initialization that was not overriden: this is a no-op.")
        data = {}
        return data

    def validation_step(self, batch, running_val_data):
        """Updates the running validataion with the current batch's results.

        The default implementation is a no-op

        Args:
          batch (dict): batch of data provided by a data pipeline.
          running_val_data (dict): current aggregates of the validation loop.

        Returns:
          updated_data (dict): new updated value for the running_val_data.
        """
        LOG.warning("Running a ModelInterface validation step that was not overriden: this is a no-op.")
        return {}

    def __repr__(self):
        return self.__class__.__name__


class Trainer(object):
    """Implements a simple training loop with hooks for callbacks.

    Args:
      interface (ModelInterface): adapter to run forward and backward
        pass on the model being trained.

    Attributes:
      callbacks (list of Callbacks): hooks that will be called while training
        progresses.
    """

    def __init__(self, interface):
        super(Trainer, self).__init__()
        self.callbacks = []
        self.interface = interface
        LOG.debug("Creating {}".format(self))

        signal.signal(signal.SIGINT, self.interrupt_handler)

        self._keep_running = True

    def interrupt_handler(self, signo, frame):
        """Stop the training process upon receiving a SIGINT (Ctrl+C)."""
        LOG.debug("interrupting run")
        self._keep_running = False

    def _stop(self):
        # Reset the run flag
        self._keep_running = True
        self.__training_end()

    def add_callback(self, callback):
        """Adds a callback to the list of training hooks.

        Args:
            callback(ttools.Callback): callback to add.
        """
        LOG.debug("Adding callback {}".format(callback))
        # pass an interface reference to the callback
        callback.model_interface = self.interface
        self.callbacks.append(callback)

    def train(self, dataloader, starting_epoch=None, num_epochs=None,
              val_dataloader=None):
        """Main training loop. This starts the training procedure.

        Args:
          dataloader (DataLoader): loader that yields training batches.
          starting_epoch (int, optional): index of the epoch we are starting from.
          num_epochs (int, optional): max number of epochs to run.
          val_dataloader (DataLoader, optional): loader that yields validation
            batches
        """
        self.__training_start(dataloader)
        if starting_epoch is None:
            starting_epoch = 0

        LOG.info("Starting taining from epoch %d", starting_epoch)

        epoch = starting_epoch

        while num_epochs is None or epoch < starting_epoch + num_epochs:
            self.__epoch_start(epoch)
            for batch_idx, batch in enumerate(dataloader):
                if not self._keep_running:
                    self._stop()
                    return
                self.__batch_start(batch_idx, batch)
                train_step_data = self.__training_step(batch)
                self.__batch_end(batch, train_step_data)
            self.__epoch_end()

            # TODO: allow validation at intermediate steps during one epoch

            # Validate
            if val_dataloader:
                with th.no_grad():
                    running_val_data = self.__validation_start(val_dataloader)
                    for batch_idx, batch in enumerate(val_dataloader):
                        if not self._keep_running:
                            self._stop()
                            return
                        self.__val_batch_start(batch_idx, batch)
                        running_val_data = self.__validation_step(batch, running_val_data)
                        self.__val_batch_end(batch, running_val_data)
                    self.__validation_end(running_val_data)

            epoch += 1

            if not self._keep_running:
                self._stop()
                return

        self._stop()

    def __repr__(self):
        return "Trainer({}, {} callbacks)".format(
            self.interface, len(self.callbacks))

    def __training_start(self, dataloader):
        for cb in self.callbacks:
            cb.training_start(dataloader)

    def __training_end(self):
        for cb in self.callbacks:
            cb.training_end()

    def __epoch_start(self, epoch_idx):
        for cb in self.callbacks:
            cb.epoch_start(epoch_idx)

    def __epoch_end(self):
        for cb in self.callbacks:
            cb.epoch_end()

    def __batch_start(self, batch_idx, batch):
        for cb in self.callbacks:
            cb.batch_start(batch_idx, batch)

    def __batch_end(self, batch, train_step_data):
        for cb in self.callbacks:
            cb.batch_end(batch, train_step_data)

    def __val_batch_start(self, batch_idx, batch):
        for cb in self.callbacks:
            cb.val_batch_start(batch_idx, batch)

    def __val_batch_end(self, batch, running_val_data):
        for cb in self.callbacks:
            cb.val_batch_end(batch, running_val_data)

    def __validation_start(self, dataloader):
        for cb in self.callbacks:
            cb.validation_start(dataloader)
        return self.interface.init_validation()

    def __validation_end(self, running_val_data):
        for cb in self.callbacks:
            cb.validation_end(running_val_data)

    def __training_step(self, batch):
        return self.interface.training_step(batch)

    def __validation_step(self, batch, running_val_data):
        return self.interface.validation_step(batch, running_val_data)


class Checkpointer(object):
    """Save and restore model and optimizer variables.

    Args:
      root (string): path to the root directory where the files are stored.
      model (torch.nn.Module):
      meta (dict): a dictionary of training or configuration parameters useful
          to initialize the model upon loading the checkpoint again.
      optimizers (single or list of torch.optimizer): optimizers whose parameters will
        be checkpointed together with the model.
      schedulers (single or list of
      torch.optim.lr_scheduler): schedulers whose
          parameters will be checkpointed together with
          the model.
      prefix (str): unique prefix name in case several models are stored in the
        same folder.
    """

    EXTENSION = ".pth"

    def __init__(self, root, model=None, meta=None, optimizers=None,
                 schedulers=None, prefix=None):
        self.root = root
        self.model = model
        self.meta = meta

        # TODO(mgharbi): verify the prefixes are unique.

        if optimizers is None:
            LOG.info("No optimizer state will be stored in the "
                        "checkpointer")
        else:
            # if we have only one optimizer, make it a list
            if not isinstance(optimizers, list):
                optimizers = [optimizers]
        self.optimizers = optimizers
        if schedulers is not None:
            if not isinstance(schedulers, list):
                schedulers = [schedulers]
        self.schedulers = schedulers

        LOG.debug(self)

        self.prefix = ""
        if prefix is not None:
            self.prefix = prefix

    def __repr__(self):
        return "Checkpointer with root at \"{}\"".format(self.root)

    def __path(self, path, prefix=None):
        if prefix is None:
            prefix = ""
        return os.path.join(self.root, prefix+os.path.splitext(path)[0] + ".pth")

    def save(self, path, extras=None):
        """Save model, metaparams and extras to relative path.

        Args:
          path (string): relative path to the file being saved (without extension).
          extras (dict): extra user-provided information to be saved with the model.
        """

        if self.model is None:
            model_state = None
        else:
            LOG.debug("Saving model state dict")
            model_state = self.model.state_dict()

        opt_dicts = []
        if self.optimizers is not None:
            for opt in self.optimizers:
                opt_dicts.append(opt.state_dict())

        sched_dicts = []
        if self.schedulers is not None:
            for s in self.schedulers:
                sched_dicts.append(s.state_dict())

        filename = self.__path(path, prefix=self.prefix)
        os.makedirs(self.root, exist_ok=True)
        th.save({'model': model_state,
                 'meta': self.meta,
                 'extras': extras,
                 'optimizers': opt_dicts,
                 'schedulers': sched_dicts,
                 }, filename)
        LOG.debug("Checkpoint saved to \"{}\"".format(filename))

    def try_and_init_from(self, path):
        """Try to initialize the models's weights from an external checkpoint.

        Args:
            path(str): full path to the checkpoints to load model parameters
                from.
        """
        LOG.info("Loading weights from foreign checkpoint {}".format(path))
        if not os.path.exists(path):
            raise ValueError("Checkpoint {} does not exist".format(path))

        chkpt = th.load(path, map_location=th.device("cpu"))
        if "model" not in chkpt.keys() or chkpt["model"] is None:
            raise ValueError("{} has no model saved".format(path))

        mdl = chkpt["model"]
        for n, p in self.model.named_parameters():
            if n in mdl:
                p2 = mdl[n]
                if p2.shape != p.shape:
                    LOG.warning("Parameter {} ignored, checkpoint size does not match: {}, should be {}".format(n, p2.shape, p.shape))
                    continue
                LOG.debug("Parameter {} copied".format(n))
                p.data.copy_(p2)
            else:
                LOG.warning("Parameter {} ignored, not found in source checkpoint.".format(n))

        LOG.info("Weights loaded from foreign checkpoint {}".format(path))

    def load(self, path):
        """Loads a checkpoint, updates the model and returns extra data.

        Args:
          path (string): path to the checkpoint file, relative to the root dir.

        Returns:
          extras (dict): extra information passed by the user at save time.
          meta (dict): metaparameters of the model passed at save time.
        """

        filename = self.__path(path, prefix=None)
        chkpt = th.load(filename, map_location="cpu")  # TODO: check behavior

        if self.model is not None and chkpt["model"] is not None:
            LOG.debug("Loading model state dict")
            self.model.load_state_dict(chkpt["model"])

        if "optimizers" in chkpt.keys():
            if self.optimizers is not None and chkpt["optimizers"] is not None:
                try:
                    for opt, state in zip(self.optimizers,
                                          chkpt["optimizers"]):
                        LOG.debug("Loading optimizers state dict for %s", opt)
                        opt.load_state_dict(state)
                except:
                    # We do not raise an error here, e.g. in case the user simply
                    # changes optimizer
                    LOG.warning("Could not load optimizer state dicts, "
                                "starting from scratch")

        if "schedulers" in chkpt.keys():
            if self.schedulers is not None and chkpt["schedulers"] is not None:
                try:
                    for s, state in zip(self.schedulers,
                                          chkpt["schedulers"]):
                        LOG.debug("Loading scheduler state dict for %s", s)
                        s.load_state_dict(state)
                except:
                    LOG.warning("Could not load scheduler state dicts, "
                                "starting from scratch")

        LOG.debug("Loaded checkpoint \"{}\"".format(filename))
        return tuple(chkpt[k] for k in ["extras", "meta"])

    def load_latest(self):
        """Try to load the most recent checkpoint, skip failing files.

        Returns:
          extras (dict): extra user-defined information that was saved in the
              checkpoint.
          meta (dict): metaparameters of the model passed at save time.
        """
        all_checkpoints = self.sorted_checkpoints()

        extras = None
        meta = None
        for f in all_checkpoints:
            try:
                extras, meta = self.load(f)
                return extras, meta
            except Exception as e:
                LOG.debug(
                    "Could not load checkpoint \"{}\", moving on ({}).".format(f, e))
        LOG.debug("No checkpoint found to load.")
        return extras, meta

    def sorted_checkpoints(self):
        """Get list of all checkpoints in root directory, sorted by creation date.

        Returns:
            chkpts (list of str): sorted checkpoints in the root folder.
        """
        reg = re.compile(r"{}.*\{}".format(self.prefix, Checkpointer.EXTENSION))
        if not os.path.exists(self.root):
            all_checkpoints = []
        else:
            all_checkpoints = [f for f in os.listdir(
                self.root) if reg.match(f)]
        mtimes = []
        for f in all_checkpoints:
            mtimes.append(os.path.getmtime(os.path.join(self.root, f)))

        mf = sorted(zip(mtimes, all_checkpoints))
        chkpts = [m[1] for m in reversed(mf)]
        LOG.debug("Sorted checkpoints {}".format(chkpts))
        return chkpts

    def delete(self, path):
        """Delete checkpoint at path.

        Args:
            path(str): full path to the checkpoint to delete.
        """
        if path in self.sorted_checkpoints():
            os.remove(os.path.join(self.root, path))
        else:
            LOG.warning("Trying to delete a checkpoint that does not exists.")

    @staticmethod
    def load_meta(root, prefix=None):
        """Fetch model metadata without touching the saved parameters.

        This loads the metadata from the most recent checkpoint in the root
        directory.

        Args:
            root(str): path to the root directory containing the checkpoints
            prefix(str): unique prefix for the checkpoint to be loaded (e.g. if
                multiple models are saved in the same location)
        """
        chkptr = Checkpointer(root, model=None, meta=None, prefix=prefix, 
                              optimizers=[])
        LOG.debug("checkpoints: %s", chkptr.sorted_checkpoints())
        _, meta = chkptr.load_latest()
        return meta
