import logging
import numpy as np

from gunpowder.array import ArrayKey, Array
from gunpowder.array_spec import ArraySpec
from gunpowder.ext import torch, tensorboardX, NoSuchModule
from gunpowder.nodes.generic_train import GenericTrain

from typing import Dict, Union, Optional, List, Any, Tuple

logger = logging.getLogger(__name__)


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

    Args:

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

            The model to train.

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

            The heads to train.

        losses:

            The torch losses to use per head.

        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 (``List(dict)``, ``string`` or ``int`` -> :class:`ArrayKey`):

            List of dictionaries with one dictionary per head.
            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.
    """

    def __init__(
        self,
        model,
        heads: List[Tuple[str, Any]],
        losses: List[Any],
        optimizer: Any,
        inputs: Dict[str, ArrayKey],
        outputs: List[Dict[Union[int, str], ArrayKey]],
        head_outputs: List[Dict[Union[int, str], ArrayKey]],
        loss_inputs: List[Dict[Union[int, str], ArrayKey]],
        gradients: List[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,
    ):

        if not model.training:
            logger.warning(
                "Model is in evaluation mode during training. "
                "Consider using model.train()"
            )

        self.head_map = {head[0]: head[1] for head in heads}

        for head, head_output_dict in zip(heads, head_outputs):
            outputs.update({f"{head[0]}_{k}": v for k, v in head_output_dict.items()})

        self.all_loss_inputs = {}
        for head, loss_input_dict in zip(heads, loss_inputs):
            inputs.update(
                {
                    f"{head[0]}_{k}": v
                    for k, v in loss_input_dict.items()
                    if v not in outputs.values()
                }
            )
            self.all_loss_inputs.update(
                {f"{head[0]}_{k}": v for k, v in loss_input_dict.items()}
            )

        self.gradient_dict = gradients[0]
        for head, grad_dict in zip(heads, gradients[1:]):
            self.gradient_dict.update(
                {f"{head[0]}_{k}": v for k, v in grad_dict.items()}
            )

        super(Train, self).__init__(
            inputs,
            outputs,
            self.gradient_dict,
            array_specs,
            spawn_subprocess=spawn_subprocess,
        )

        self.model = model
        self.heads = heads
        self.losses = losses
        self.optimizer = optimizer
        self.loss_inputs = loss_inputs
        self.checkpoint_basename = checkpoint_basename
        self.save_every = save_every

        self.iteration = 0

        if not isinstance(tensorboardX, NoSuchModule) and log_dir is not None:
            self.summary_writer = tensorboardX.SummaryWriter(log_dir)
            self.log_every = log_every
        else:
            self.summary_writer = None
            if log_dir is not None:
                logger.warning("log_dir given, but tensorboardX is not installed")

        self.intermediate_layers = {}
        self.register_hooks()

    def register_hooks(self):
        for key in self.outputs:
            if isinstance(key, str):
                name_parts = key.split("_")
                if name_parts[0] in self.head_map.keys():
                    layer_name = "_".join(name_parts[1:])
                    try:
                        layer_name = int(layer_name)
                    except ValueError:
                        layer = getattr(self.head_map[name_parts[0]], layer_name)
                        layer.register_forward_hook(self.create_hook(key))
                else:
                    layer = getattr(self.model, key)
                    layer.register_forward_hook(self.create_hook(key))

    def create_hook(self, key):
        def save_layer(module, input, output):
            self.intermediate_layers[key] = output

        return save_layer

    def retain_gradients(self, request, outputs):
        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):
                name_parts = array_name.split("_")
                if name_parts[0] in self.head_map.keys():
                    layer_name = "_".join(name_parts[1:])
                    try:
                        layer_name = int(layer_name)
                        tensor = outputs[array_name]
                    except ValueError:
                        tensor = getattr(self.head_map[name_parts[0]], layer_name)
                else:
                    tensor = getattr(self.model, array_name)
            else:
                raise RuntimeError(
                    "only ints and strings are supported as gradients keys"
                )
            tensor.retain_grad()

    def start(self):

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

        try:
            self.model = self.model.to(self.device)
            for state in self.optimizer.state.values():
                for k, v in state.items():
                    if isinstance(v, torch.Tensor):
                        state[k] = v.cuda()
            for i, head in enumerate(self.heads):
                self.heads[i] = (head[0], head[1].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
        for i in range(len(self.losses)):
            if isinstance(self.losses[i], torch.nn.Module):
                self.losses[i] = self.losses[i].to(self.device)

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

        if checkpoint is not None:

            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"])
            for i, head in enumerate(self.heads):
                self.heads[i][1].load_state_dict(checkpoint[f"head_{i}"])
            self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

        else:

            logger.info("Starting training from scratch")

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

    def train_step(self, batch, request):

        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.model(**device_inputs)
        if isinstance(model_outputs, torch.Tensor):
            outputs = {0: model_outputs}
        else:
            raise RuntimeError(
                "Torch train node only supports return types of ",
                f"torch.Tensor from model.forward(). not {type(model_outputs)}",
            )
        outputs.update(self.intermediate_layers)

        head_outputs = []
        for i, head in enumerate(self.heads):
            head_output = head[1](model_outputs)
            head_outputs.append(head_output)
            assert isinstance(
                head_output, torch.Tensor
            ), "predictor did not return a torch tensor"
            outputs[f"{head[0]}_0"] = head_output

        # Add hooks to retain gradients on some outputs
        self.retain_gradients(request, outputs)

        flipped_outputs = {v: outputs[k] for k, v in self.outputs.items()}

        all_losses = []
        i = 0
        for loss, loss_inputs in zip(self.losses, self.loss_inputs):
            i += 1
            # Some inputs to the loss should come from the batch, not the model
            provided_loss_inputs = self.__collect_provided_loss_inputs(
                batch, loss_inputs
            )
            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
            device_loss_inputs = {
                k: flipped_outputs.get(v, device_loss_inputs.get(k))
                for k, v in loss_inputs.items()
            }
            assert all(
                v is not None for v in device_loss_inputs.values()
            ), "loss function input must be provided upstream or by the model"

            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, (
                f"Not all loss inputs could be interpreted. "
                f"Failed keys: {device_loss_inputs.keys()}"
            )

            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()},
            )

            loss = loss(*device_loss_args, **device_loss_kwargs)
            all_losses.append(loss)

        loss = sum(all_losses)
        loss.backward()
        self.optimizer.step()

        # 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):
                name_parts = array_name.split("_")
                if name_parts[0] in self.head_map.keys():
                    layer_name = "_".join(name_parts[1:])
                    try:
                        layer_name = int(layer_name)
                        tensor = outputs[array_name]
                    except ValueError:
                        tensor = getattr(self.head_map[name_parts[0]], layer_name)
                else:
                    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

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

            checkpoint_name = self._checkpoint_name(
                self.checkpoint_basename, batch.iteration
            )

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

            state_dicts = {
                "model_state_dict": self.model.state_dict(),
                "optimizer_state_dict": self.optimizer.state_dict(),
            }
            for i, head in enumerate(self.heads):
                state_dicts[f"head_{i}"] = head.state_dict()

            torch.save(
                state_dicts,
                checkpoint_name,
            )

        if self.summary_writer and batch.iteration % self.log_every == 0:
            self.summary_writer.add_scalar("loss", batch.loss, batch.iteration)

    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.all_loss_inputs},
            batch,
        )

    def __collect_provided_loss_inputs(self, batch, loss_inputs):

        return self.__collect_provided_arrays(
            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 {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
