import functools
import itertools
from abc import ABC
from enum import Enum
from typing import List, Optional, Tuple, Union

import torch
import torch.utils.hooks as hooks
from torch import Tensor
from torch.nn.modules.module import (
    _global_backward_hooks,
    _global_forward_hooks,
    _global_forward_pre_hooks,
)

__all__ = ["CoModule", "call_mode"]

# First element is a value buffer, while others are indexes
State = Tuple[Tensor, Tensor, Tensor]


class PaddingMode(Enum):
    REPLICATE = "replicate"
    ZEROS = "zeros"
    NEG_INF = "neg_inf"


# Not compatible with torch.jit:
# class CallMode(Enum):
#     FORWARD = "forward"
#     FORWARD_STEPS = "forward_steps"
#     FORWARD_STEP = "forward_step"

CALL_MODES = {
    "forward": torch.tensor(0),
    "forward_steps": torch.tensor(1),
    "forward_step": torch.tensor(2),
}


def _callmode(cm) -> torch.Tensor:
    """_summary_

    Args:
        cm (Union[str, int, torch.Tensor]): Identifier for call mode

    Returns:
        torch.Tensor: validated call_mode
    """
    if isinstance(cm, str):
        cm = CALL_MODES[cm.lower()]
    elif isinstance(cm, int):
        cm = torch.tensor(int)
    return cm


class _CallModeContext(object):
    """Context-manager which temporarily specifies a call_mode

    When the call_mode context is used, the ``__call__`` function of continual modules with be set accordingly

    Example::

        sequence = torch.randn(1, 3, 10)
        step = sequence[:, :, -1]

        forward_output = module(sequence)  # Calls `forward`

        with co.call_mode("forward_step"):
            forward_step_output = module(step)  # Calls `forward_step`

        forward_output = module(sequence)  # Calls `forward`
    """

    def __init__(self):
        self.cur = _callmode("forward")
        self.prev = None

    def __call__(self, value: Union[str, int, torch.Tensor]):
        self.prev = self.cur
        self.cur = _callmode(value)
        return self

    def __enter__(self):
        pass

    def __exit__(self, *args, **kwargs):
        self.cur = self.prev
        self.prev = None


call_mode = _CallModeContext()


def _clone_first(state: State) -> State:
    # return (state[0].clone(), state[1], state[2])
    return (state[0].clone(),) + state[1:]
    # return (state[0].clone(),) + state[1:]


class CoModule(ABC):
    """Base class for continual modules.
    Deriving from this class provides base-functionality and enforces
    the implementation of necessary methods.

    Attributes:
        receptive_field (int): Temporal receptive field of the module.
        delay (int): Number of step inputs to observe before the modules produces valid outputs.
        stride (Tuple[int,...]): (Spatio)-temporal stride.
        padding (Tuple[int,...]): (Spatio)-temporal padding.

    """

    receptive_field: int = 1
    stride: Tuple[int, ...] = (1,)
    padding: Tuple[int, ...] = (0,)
    _make_padding = torch.zeros_like
    _state_shape: int = 0
    _dynamic_state_inds: List[bool] = []
    _call_mode = _callmode("forward")

    def __init_subclass__(cls) -> None:
        CoModule._validate_class(cls)

    @staticmethod
    def _validate_class(cls):
        for fn, description in [
            ("forward_step", "forward computation for a single temporal step"),
            (
                "forward_steps",
                "forward computation for multiple temporal step",
            ),
            (
                "forward",
                "a forward computation which is identical to a regular non-continual forward.",
            ),
            ("get_state", "a retrieval of the internal state."),
            ("set_state", "an update of the internal state."),
            ("clean_state", "an internal state clean-up."),
        ]:
            assert callable(
                getattr(cls, fn, None)
            ), f"{cls} should implement a `{fn}` function which performs {description} to satisfy the CoModule interface."

        for prop in {"delay", "receptive_field"}:
            assert type(getattr(cls, prop, None)) in {
                int,
                torch.Tensor,
                property,
            }, f"{cls} should implement a `{prop}` property to satisfy the CoModule interface."

        for prop in {"stride", "padding", "_state_shape", "_dynamic_state_inds"}:
            assert type(getattr(cls, prop, None)) in {
                int,
                property,
                tuple,
                list,
            }, f"{cls} should implement a `{prop}` property to satisfy the CoModule interface."

    @staticmethod
    def is_valid(module):
        try:
            CoModule._validate_class(module)
        except AssertionError:
            return False
        return True

    def get_state(self) -> Optional[State]:
        """Get model state.

        Returns:
            Optional[State]: A State tuple if the model has been initialised and otherwise None.
        """
        ...  # pragma: no cover

    def set_state(self, state: State):
        """Set model state

        Args:
            state (State): State tuple to set as new internal internal state
        """
        ...  # pragma: no cover

    def clean_state(self):
        """Clean model state, resetting the network memory."""
        ...  # pragma: no cover

    @property
    def delay(self) -> int:
        return self.receptive_field - 1 - self.padding[0]

    def forward_step(
        self, input: Tensor, update_state: bool = True
    ) -> Optional[Tensor]:
        """Performs a forward computation for a single frame and (optionally) updates internal states accordingly.
        This function performs efficient continual inference.

        Illustration::

            O+S O+S O+S O+S   (O: output, S: updated internal state)
             ↑   ↑   ↑   ↑
             N   N   N   N    (N: network module)
             ↑   ↑   ↑   ↑
             I   I   I   I    (I: input frame)

        Args:
            input (Tensor): Layer input.
            update_state (bool): Whether internal state should be updated during this operation.

        Returns:
            Optional[Tensor]: Step output. This will be a placeholder while the module initializes and every (stride - 1) / stride.
        """
        state = self.get_state()
        if not update_state and state:
            state = _clone_first(state)
        output, state = self._forward_step(input, state)
        if update_state and state:
            self.set_state(state)
        return output

    def forward_steps(
        self, input: Tensor, pad_end: bool = False, update_state=True
    ) -> Optional[Tensor]:
        """Performs a forward computation across multiple time-steps while updating internal states for continual inference (if update_state=True).
        Start-padding is always accounted for, but end-padding is omitted per default in expectance of the next input step. It can be added by specifying pad_end=True. If so, the output-input mapping the exact same as that of forward.

        Illustration::

                    O            (O: output)
                    ↑
            -----------------    (-: aggregation)
            O  O+S O+S O+S  O    (O: output, S: updated internal state)
            ↑   ↑   ↑   ↑   ↑
            N   N   N   N   N    (N: network module)
            ↑   ↑   ↑   ↑   ↑
            P   I   I   I   P    (I: input frame, P: padding)

        Args:
            input (Tensor): Layer input.
            pad_end (bool): Whether results for temporal padding at sequence end should be included.
            update_state (bool): Whether internal state should be updated during this operation.

        Returns:
            Optional[Tensor]: Layer output
        """
        return self._forward_steps_impl(input, pad_end, update_state)

    def _forward_steps_impl(
        self, input: Tensor, pad_end: bool = False, update_state: bool = True
    ) -> Optional[Tensor]:
        """Forward computation for multiple steps with state initialisation

        Args:
            module (CoModule): Continual module.
            input (Tensor): Layer input.
            pad_end (bool): Whether results for temporal padding at sequence end should be included.
            update_state (bool): Whether internal state should be updated during this operation.

        Returns:
            Tensor: Layer output
        """
        outs = []
        state: Optional[State] = self.get_state()

        if not update_state and state is not None:
            state = _clone_first(state)

        for t in range(input.shape[2]):
            o, state = self._forward_step(input[:, :, t], state)
            if isinstance(o, Tensor):
                outs.append(o)

        if update_state and state is not None:
            self.set_state(state)

        if pad_end:
            # Don't save state for the end-padding
            opt_state = self.get_state()
            if opt_state is not None:
                state = _clone_first(opt_state)

            for t, i in enumerate(
                [self._make_padding(input[:, :, -1]) for _ in range(self.padding[0])]
            ):
                o, state = self._forward_step(i, state)
                if isinstance(o, Tensor):
                    outs.append(o)

        if len(outs) == 0:
            return None  # pragma: no cover

        return torch.stack(outs, dim=2)

    def forward(self, input: Tensor) -> Tensor:
        """Performs a forward computation over multiple time-steps.
        This function is identical to the corresponding module in _torch.nn_,
        ensuring cross-compatibility. Moreover, it's handy for efficient training
        on clip-based data.

        Illustration::

                    O            (O: output)
                    ↑
                    N            (N: network module)
                    ↑
            -----------------    (-: aggregation)
            P   I   I   I   P    (I: input frame, P: padding)

        Args:
            input (Tensor): Network input.
        """
        ...  # pragma: no cover

    def warm_up(self, step_shape: List[int]):
        """Warms up the model state with a dummy input.
        The initial `self.delay` steps will produce results, but they will be inexact.

        To warm up the model with a user-defined data, pass the data to forward_steps::

            net.forward_steps(user_data)

        Args:
            step_shape (Sequence[int]): input shape with which to warm the model up, including batch size.
        """
        step_shape = (*step_shape[:2], self.delay, *step_shape[2:])
        dummy = self._make_padding(torch.zeros(step_shape, dtype=torch.float))
        self.forward_steps(dummy)

    @property
    def call_mode(self) -> torch.Tensor:
        return self._call_mode

    @call_mode.setter
    def call_mode(self, value):
        self._call_mode = _callmode(value)
        if hasattr(self, "__len__"):
            for m in self:
                if hasattr(m, "call_mode"):
                    m.call_mode = self._call_mode

    def _call_impl(self, *input, **kwargs):  # noqa: C901  # pragma: no cover
        """Modified version torch.nn.Module._call_impl

        Returns:
            [type]: [description]
        """
        _call_mode = call_mode.cur if call_mode.prev is not None else self.call_mode
        forward_call = {
            (True, _callmode("forward")): self._slow_forward,
            (False, _callmode("forward")): self.forward,
            (False, _callmode("forward_steps")): self.forward_steps,
            (False, _callmode("forward_step")): self.forward_step,
        }[(bool(torch._C._get_tracing_state()), _call_mode)]

        # If we don't have any hooks, we want to skip the rest of the logic in
        # this function, and just call forward.
        if not (
            self._backward_hooks
            or self._forward_hooks
            or self._forward_pre_hooks
            or _global_backward_hooks
            or _global_forward_hooks
            or _global_forward_pre_hooks
        ):
            return forward_call(*input, **kwargs)
        # Do not call functions when jit is used
        full_backward_hooks, non_full_backward_hooks = [], []
        if self._backward_hooks or _global_backward_hooks:
            full_backward_hooks, non_full_backward_hooks = self._get_backward_hooks()
        if _global_forward_pre_hooks or self._forward_pre_hooks:
            for hook in itertools.chain(
                _global_forward_pre_hooks.values(), self._forward_pre_hooks.values()
            ):
                result = hook(self, input)
                if result is not None:
                    if not isinstance(result, tuple):
                        result = (result,)
                    input = result

        bw_hook = None
        if full_backward_hooks:
            bw_hook = hooks.BackwardHook(self, full_backward_hooks)
            input = bw_hook.setup_input_hook(input)

        result = forward_call(*input, **kwargs)
        if _global_forward_hooks or self._forward_hooks:
            for hook in itertools.chain(
                _global_forward_hooks.values(), self._forward_hooks.values()
            ):
                hook_result = hook(self, input, result)
                if hook_result is not None:
                    result = hook_result

        if bw_hook:
            result = bw_hook.setup_output_hook(result)

        # Handle the non-full backward hooks
        if non_full_backward_hooks:
            var = result
            while not isinstance(var, torch.Tensor):
                if isinstance(var, dict):
                    var = next((v for v in var.values() if isinstance(v, torch.Tensor)))
                else:
                    var = var[0]
            grad_fn = var.grad_fn
            if grad_fn is not None:
                for hook in non_full_backward_hooks:
                    wrapper = functools.partial(hook, self)
                    functools.update_wrapper(wrapper, hook)
                    grad_fn.register_hook(wrapper)
                self._maybe_warn_non_full_backward_hook(input, result, grad_fn)

        return result

    __call__ = _call_impl

    @staticmethod
    def build_from(
        module: torch.nn.Module,
        *args,
        **kwargs,
    ) -> "CoModule":
        """Copy parameters and weights from a non-continual module and
        build the corresponding continual version.

        Args:
            module (torch.nn.Module): Module from which to copy variables and weights

        Returns:
            CoModule: Continual Module with the parameters and weights of the passed module.
        """
