"""Provides an adapted version of the gunpowder torch Train node
"""
import logging
import os

import numpy as np

from gunpowder.array import ArrayKey, Array
from gunpowder.array_spec import ArraySpec
from gunpowder.ext import torch, NoSuchModule
from gunpowder.torch.nodes import Train

try:
    from torch.utils import tensorboard
except ImportError:
    tensorboard = NoSuchModule('tensorboard')


from typing import Dict, Union, Optional

logger = logging.getLogger(__name__)


class TorchTrainExt(Train):
    """extended Torch implementation of :class:`gunpowder.nodes.GenericTrain`.

    Args:

        model (subclass of ``torch.nn.Module``):

            The model to train.

        loss:

            The torch loss to use.

        optimizer:

            The torch optimizer to use.

        inputs (``dict``, ``string`` -> :class:`ArrayKey`):

            Dictionary from the names of input tensors (argument names of the
            ``forward`` method) in the model to array keys.

        loss_inputs (``dict``, ``string`` or ``int`` -> :class:`ArrayKey`):

            Dictionary with the names of input variables to the loss function
            as keys, and ArrayKeys containing the desired data as values. Keys
            can be either strings or integers. If the key is an integer, it
            will be treated as a positional argument to the loss function, a
            string will be used as a named argument

        outputs (``dict``, ``string`` or ``int`` -> :class:`ArrayKey`):

            Dictionary from the names of tensors in the network to array
            keys. If the key is a string, the tensor will be retrieved
            by checking the model for an attribute with they key as its name.
            If the key is an integer, it is interpreted as a tuple index of
            the outputs of the network.
            New arrays will be generated by this node for each entry (if
            requested downstream).

        array_specs (``dict``, :class:`ArrayKey` -> :class:`ArraySpec`,
                optional):

            Used to set the specs of generated arrays (at the moment only
            ``output``). This is useful to set the ``voxel_size``, for example,
            if they differ from the voxel size of the input arrays. Only fields
            that are not ``None`` in the given :class:`ArraySpec` will be used.

        checkpoint_basename (``string``, optional):

            The basename used for checkpoint files. Defaults to ``model``.

        save_every (``int``, optional):

            After how many iterations to create a checkpoint to store the
            learnt weights.

        log_dir (``string``, optional):

            Directory for saving tensorboard summaries.

        log_every (``int``, optional):

            After how many iterations to write out tensorboard summaries.

        spawn_subprocess (``bool``, optional):

            Whether to run the ``train_step`` in a separate process.
            Default is false.
    """

    def __init__(
        self,
        model,
        loss,
        optimizer,
        inputs: Dict[str, ArrayKey],
        outputs: Dict[Union[int, str], ArrayKey],
        loss_inputs: Dict[Union[int, str], ArrayKey],
        gradients: Dict[Union[int, str], ArrayKey] = {},
        array_specs: Optional[Dict[ArrayKey, ArraySpec]] = None,
        checkpoint_basename: str = "model",
        save_every: int = 2000,
        log_dir: str = None,
        log_every: int = 1,
        spawn_subprocess: bool = False,
        val_log_step: int = None,
        use_auto_mixed_precision: bool = False,
        use_swa: bool = False,
        swa_every_it: bool = False,
        swa_start_it: int = None,
        swa_freq_it: int = None,
    ):

        super(TorchTrainExt, self).__init__(
            model, loss, optimizer, inputs, outputs, loss_inputs,
            gradients=gradients, array_specs=array_specs,
            checkpoint_basename=checkpoint_basename, save_every=save_every,
            log_dir=log_dir, log_every=log_every,
            spawn_subprocess=spawn_subprocess
        )

        self.val_log_step = val_log_step
        self.use_auto_mixed_precision = use_auto_mixed_precision
        self.use_swa = use_swa
        self.swa_every_it = swa_every_it
        self.swa_start_it = swa_start_it
        self.swa_freq_it = swa_freq_it

        self.current_step = 0

        if not isinstance(tensorboard, NoSuchModule) and log_dir is not None:
            self.summary_writer = tensorboard.SummaryWriter(log_dir)
            if self.val_log_step is not None:
                self.summary_writer_val = tensorboard.SummaryWriter(
                    os.path.join(os.path.dirname(log_dir), "val"))
            self.log_every = log_every
        else:
            self.summary_writer = None
            if log_dir is not None:
                logger.warning(
                    "log_dir given, but tensorboard is not installed")

    def start(self):

        self.use_cuda = torch.cuda.is_available()
        self.device = torch.device("cuda" if self.use_cuda else "cpu")

        if self.use_auto_mixed_precision:
            self.scaler = torch.cuda.amp.GradScaler(init_scale=8192.0,
                                                    growth_interval=200)
            logger.info("using auto mixed precision")

        try:
            self.model = self.model.to(self.device)
        except RuntimeError as e:
            raise RuntimeError(
                "Failed to move model to device. If you are using a child "
                "process to run your model, maybe you already initialized CUDA"
                " by sending your model to device in the main process."
            ) from e
        if isinstance(self.loss, torch.nn.Module):
            self.loss = self.loss.to(self.device)

        if self.use_swa:
            self.swa_model = torch.optim.swa_utils.AveragedModel(self.model)
            self.swa_scheduler = torch.optim.swa_utils.SWALR(
                self.optimizer, anneal_strategy="linear",
                anneal_epochs=10, swa_lr=5e-6)

        checkpoint, self.iteration = self._get_latest_checkpoint(
            self.checkpoint_basename
        )
        self.current_step = int(self.iteration)

        if checkpoint is not None and self.iteration != 0:

            logger.info("Resuming training from iteration %d", self.iteration)
            logger.info("Loading %s", checkpoint)

            checkpoint = torch.load(checkpoint, map_location=self.device)
            self.model.load_state_dict(checkpoint["model_state_dict"])
            self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
            if self.use_swa:
                self.swa_model.load_state_dict(
                    checkpoint["swa_model_state_dict"])
                self.swa_scheduler.load_state_dict(
                    checkpoint["swa_scheduler_state_dict"])
            if self.use_auto_mixed_precision:
                if "scaler_state_dict" not in checkpoint:
                    logger.warning("no scaler state dict in checkpoint!")
                else:
                    self.scaler.load_state_dict(
                        checkpoint["scaler_state_dict"])
        else:

            logger.info("Starting training from scratch")

        logger.info("Using device %s", self.device)
        self.initialized = True

    def train_step(self, batch, request):

        # initialize before first step,
        if not self.initialized:
            self.start()

        inputs = self._collect_provided_inputs(batch)
        requested_outputs = self._collect_requested_outputs(request)

        # keys are argument names of model forward pass
        device_inputs = {
            k: torch.as_tensor(v, device=self.device)
            for k, v in inputs.items()
        }

        # get outputs.
        # Keys are tuple indices or model attr names as in self.outputs
        self.optimizer.zero_grad()
        model_outputs = self._get_model_output(device_inputs)
        if isinstance(model_outputs, tuple):
            outputs = {i: model_outputs[i] for i in range(len(model_outputs))}
        elif isinstance(model_outputs, torch.Tensor):
            outputs = {0: model_outputs}
        else:
            raise RuntimeError(
                "Torch train node only supports return types of tuple",
                "and torch.Tensor from model.forward(). not "
                f"{type(model_outputs)}"
            )
        outputs.update(self.intermediate_layers)

        # Some inputs to the loss should come from the batch, not the model
        provided_loss_inputs = self._collect_provided_loss_inputs(batch)

        device_loss_inputs = {
            k: torch.as_tensor(v, device=self.device)
            for k, v in provided_loss_inputs.items()
        }

        # Some inputs to the loss function should come from the outputs of
        # the model
        # Update device loss inputs with tensors from outputs if available
        flipped_outputs = {v: outputs[k] for k, v in self.outputs.items()}
        device_loss_inputs = {
            k: flipped_outputs.get(v, device_loss_inputs.get(k))
            for k, v in self.loss_inputs.items()
        }

        device_loss_args = []
        for i in range(len(device_loss_inputs)):
            if i in device_loss_inputs:
                device_loss_args.append(device_loss_inputs.pop(i))
            else:
                break
        device_loss_kwargs = {}
        for k, v in list(device_loss_inputs.items()):
            if isinstance(k, str):
                device_loss_kwargs[k] = device_loss_inputs.pop(k)
        assert len(device_loss_inputs) == 0, (
            "Not all loss inputs could be interpreted. Failed keys: "
            f"{device_loss_inputs.keys()}")

        self.retain_gradients(request, outputs)

        logger.debug(
            "model outputs: %s",
            {k: v.shape for k, v in outputs.items()})
        logger.debug(
            "loss inputs: %s %s",
            [v.shape for v in device_loss_args],
            {k: v.shape for k, v in device_loss_kwargs.items()})
        with torch.cuda.amp.autocast(enabled=self.use_auto_mixed_precision):
            loss, ciloss, pvloss, summaries, cisum = self.loss(
                *device_loss_args, **device_loss_kwargs)

        self._update_model_weights(loss)

        # add requested model outputs to batch
        for array_key, array_name in requested_outputs.items():
            spec = self.spec[array_key].copy()
            spec.roi = request[array_key].roi
            batch.arrays[array_key] = Array(
                outputs[array_name].cpu().detach().numpy(), spec
            )

        for array_name, array_key in self.gradients.items():
            if array_key not in request:
                continue
            if isinstance(array_name, int):
                tensor = outputs[array_name]
            elif isinstance(array_name, str):
                tensor = getattr(self.model, array_name)
            else:
                raise RuntimeError(
                    "only ints and strings are supported as gradients keys"
                )
            spec = self.spec[array_key].copy()
            spec.roi = request[array_key].roi
            batch.arrays[array_key] = Array(
                tensor.grad.cpu().detach().numpy(), spec
            )

        for array_key, array_name in requested_outputs.items():
            spec = self.spec[array_key].copy()
            spec.roi = request[array_key].roi
            batch.arrays[array_key] = Array(
                outputs[array_name].cpu().detach().numpy(), spec
            )

        batch.loss = loss.cpu().detach().numpy()
        self.iteration += 1
        batch.iteration = self.iteration
        self.current_step = int(batch.iteration)

        if batch.iteration % self.save_every == 0:

            self._save_checkpoint(batch.iteration)

        if self.summary_writer:
            self._write_summary(summaries, batch.iteration)

    def _do_swa_update(self):
        return (
            self.use_swa and
            (self.current_step) >= self.swa_start_it and
            (self.swa_every_it or
             (self.current_step % self.swa_freq_it == self.swa_freq_it - 1))
        )

    def _save_checkpoint(self, iteration):
        checkpoint_name = self._checkpoint_name(
            self.checkpoint_basename, iteration
        )

        logger.info("Creating checkpoint %s", checkpoint_name)

        data_to_save = {
            "model_state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
        }
        if self.use_swa:
            data_to_save['swa_model_state_dict'] = self.swa_model.state_dict()
            data_to_save['swa_scheduler_state_dict'] = \
                self.swa_scheduler.state_dict()
        if self.use_auto_mixed_precision:
            data_to_save['scaler_state_dict'] = self.scaler.state_dict()
        torch.save(
            data_to_save,
            checkpoint_name,
        )

    def _write_summary(self, summaries, iteration):
        for k, (v, f) in summaries.items():
            if self.val_log_step is not None and \
               (self.current_step-1) % self.val_log_step == 1:
                logger.debug("writing val summary %d", self.current_step)
                self.summary_writer_val.add_scalar(k, v, iteration)
            else:
                if (self.current_step-1) % f == 0:
                    self.summary_writer.add_scalar(k, v, iteration)

    def _get_model_output(self, device_inputs):
        with torch.cuda.amp.autocast(enabled=self.use_auto_mixed_precision):
            if (self.val_log_step is not None and
                    self.use_swa and
                    self.current_step >= self.swa_start_it and
                    self.current_step % self.val_log_step == 1):
                logger.debug("val with swa model")
                model_outputs = self.swa_model(**device_inputs)
            else:
                model_outputs = self.model(**device_inputs)

        return model_outputs

    def _update_model_weights(self, loss):
        if self.val_log_step is None or \
           int(self.current_step) % self.val_log_step != 1:
            if self.use_auto_mixed_precision:
                self.scaler.scale(loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                loss.backward()
                self.optimizer.step()

        if self._do_swa_update():
            logger.debug("updating swa")
            self.swa_model.update_parameters(self.model)
            self.swa_scheduler.step()

    def _collect_requested_outputs(self, request):

        array_outputs = {}

        for output_name, array_key in self.outputs.items():
            if array_key in request:
                array_outputs[array_key] = output_name

        return array_outputs

    def _collect_provided_inputs(self, batch):

        return self._collect_provided_arrays(
            {k: v for k, v in self.inputs.items()
             if k not in self.loss_inputs},
            batch
        )

    def _collect_provided_loss_inputs(self, batch):

        return self._collect_provided_arrays(
            self.loss_inputs, batch, expect_missing_arrays=True
        )

    def _collect_provided_arrays(
            self, reference, batch, expect_missing_arrays=False):

        arrays = {}

        for array_name, array_key in reference.items():
            if isinstance(array_key, ArrayKey):
                msg = (f"batch does not contain {array_key}, array "
                       f"{array_name} will not be set")
                if array_key in batch.arrays:
                    arrays[array_name] = batch.arrays[array_key].data
                elif not expect_missing_arrays:
                    logger.warn(msg)
                else:
                    logger.debug(msg)
            elif isinstance(array_key, np.ndarray):
                arrays[array_name] = array_key
            elif isinstance(array_key, str):
                arrays[array_name] = getattr(batch, array_key)
            else:
                raise Exception(
                    "Unknown network array key {}, can't be given to "
                    "network".format(array_key)
                )

        return arrays
