# AUTOGENERATED FROM: babilim/core/module.ipynb

# Cell: 0
"""doc
# babilim.core.module

> An object which can have a state that is trainable or checkpointable. The core unit of babilim.
"""

# Cell: 1
from typing import Sequence, Any, Sequence, Callable, Dict, Iterable
from collections import defaultdict, OrderedDict
import inspect
import numpy as np

import babilim
from babilim import PYTORCH_BACKEND, TF_BACKEND
from babilim.core.annotations import RunOnlyOnce
from babilim.core.device import get_current_device_native_format, Device
from babilim.core.checkpoint import load_state, save_state
from babilim.core.itensor import ITensor
from babilim.core.tensor import Tensor, TensorWrapper
from babilim.core.logging import info, warn, DEBUG_VERBOSITY


@RunOnlyOnce
def _warn_once(msg):
    warn(msg)


# Cell: 2
_statefull_object_name_table = {}


# Cell: 3
class Module(object):
    def __init__(self):
        """
        A module is an object with variables that can be trainable and checkpointable.

        Furthermore every module is callable.
        A module can be used with native or babilim tensors when the callable api is used.
        It automatically wraps native tensors and calls the `call` function.
        
        Attributes:
        * `self.initialized_module`: A boolen storing if the module is already initialized. When not initialized loading state will fail.
        * `self.device`: Specifies the device on which this module is.
        """
        self._wrapper = TensorWrapper()
        self._training = True
        self.initialized_module = False
        self.device = get_current_device_native_format()

    def initialize(self, dataset):
        """
        Initializes your module by running a sample of your dataset through it.

        :param dataset: The dataset you want to use for initialization. (Must be of type babilim.data.Dataset)
        """
        if not self.initialized_module:
            if babilim.core.logging.DEBUG_VERBOSITY:
                info("Build Model")
            self.initialized_module = True
            dataloader = dataset.to_dataloader()
            features, _ = next(iter(dataloader))
            self(**features._asdict())

    def __call__(self, *args, **kwargs) -> Any:
        """
        Makes a module callable. Automatically wraps tensorflow or pytorch tensors to ITensors from babilim.

        ```python
        module = MyModule()
        module(*args, **kwargs)
        ```

        Warning: This function should not be overwritten. Instead overwrite `call` with no underscores.

        :param *args: All indexed parameters of your call function derivate.
        :param **kwargs: All named parameters of your call function derivate.
        """
        with Device(self.device):
            # ensure that call gets called with ITensor objects but the caller can use native tensors.
            args_wrapped = self._wrapper.wrap(args)
            kwargs_wrapped = self._wrapper.wrap(kwargs)
            wrapped = False
            if args_wrapped is not None and len(args) > 0:
                wrapped = True
                args = args_wrapped
            if kwargs_wrapped is not None and len(kwargs) > 0:
                wrapped = True
                kwargs = kwargs_wrapped

            self.build(*args, **kwargs)
            result = self.call(*args, **kwargs)
            parent_dict = inspect.stack()[1][0].f_locals
            if "self" in parent_dict:
                parent = parent_dict["self"]
                self._register_params(parent)
            if wrapped:
                return self._wrapper.unwrap(result)
            else:
                return result

    def build(self, *args, **kwargs) -> None:
        """
        This function will build your model and must be annotated with the RunOnlyOnce-Annotation.

        Allocating weight tensors should be done here.
        You can make use of the knowledge of your inputs to compute shapes for your weight tensors.
        This will make coding dimensions a lot easier.

        ```python
        @RunOnlyOnce
        def build(self, image: ITensor) -> None:
        ```

        :param *args: You must specify the exact same parameters as for your call.
        :param **kwargs: You must specify the exact same parameters as for your call.
        """
        pass

    def call(self, *args, **kwargs) -> Any:
        """
        Makes a module callable and contains the forward pass of your model.
        This should be pure computation and not allocate any weights.
        Allocating weights should be done in the `build` function.

        This function gets called by `__call__` and must be overwritten by any derived class.

        ```python
        def call(self, image: ITensor) -> NetworkOutput:
        ```

        :param *args: You can specify any parameters you want.
        :param **kwargs: You can specify any named parameters you want.
        """
        raise NotImplementedError("Every modules must implement this method.")

    def predict(self, **kwargs):
        """
        Pass in single training examples as numpy arrays.
        And predict the value without gradients.
        Should be used for testing and evaluation.

        If your network has eval modes you need to set them manually.

        The array must not have batch dimension.

        :param kwargs: The parameters to feed the network as a single example.
        :return: The output for a single example.
        """
        with Device(self.device):
            kwargs = {k: np.array([kwargs[k]]) for k in kwargs.keys() if isinstance(kwargs[k], np.ndarray)}
            kwargs = {k: Tensor(data=kwargs[k], trainable=False) for k in kwargs.keys()}

            preds = self.__call__(**kwargs)
            tmp = preds._asdict()
            tmp = {k: tmp[k].numpy()[0] for k in tmp.keys()}
            preds = type(preds)(**tmp)
            return preds

    @property
    def submodules(self):
        """
        A property to get all submodules.
        
        A submodule is a module stored in an attribute of a module.
        
        ```python
        module.submodules
        ```
        """
        modules = []
        for k in self.__dict__:
            v = self.__dict__[k]
            if isinstance(v, Module):
                modules.append(v)
                modules.append(v.submodules)
        return modules

    def modules(self):
        """
        Returns an iterator over all submodules in the module.
        
        A submodule is a module stored in an attribute of a module.
        """
        for name, module in self.named_modules():
            yield module

    def named_modules(self, memo=None, prefix=''):
        """
        A named list of all submodules.
        
        A submodule is a module stored in an attribute of a module.
        """
        modules = {}
        for k in self.__dict__:
            v = self.__dict__[k]
            if isinstance(v, Module):
                modules[prefix + "." + k] = v
                modules.update(**dict(v.named_modules(memo, prefix)))
        return list(modules.items())

    @RunOnlyOnce
    def _register_params(self, module):
        """
        Allows registration of the parameters with a native module.

        This makes the parameters of a babilim modules available to the native modules.
        When using a babilim modules in a native modules, use this function and pass the native module as a parameter.

        This function works by adding all trainable_variables to the module you pass.
        Warning: You need to build the babilim modules before calling this function. Building can be done by calling for example.

        Here is a pytorch example:

        ```python
        import torch
        from torch.nn import Module
        from babilim.modules import Linear

        class MyModule(Module):
            def __init__(self):
                super().__init__()
                self.linear = Linear(10)

            def forward(self, features):
                result = self.linear(features)
                self.linear.register_params(self)
                return result
        ```

        :param module: The native module on which parameters of this modules should be registered.
        """
        if babilim.is_backend(PYTORCH_BACKEND):
            from torch.nn import Module
            if isinstance(module, Module):
                myname = "_error_"
                for var in module.__dict__:
                    if module.__dict__[var] == self:
                        myname = var
                    if isinstance(module.__dict__[var], list) and self in module.__dict__[var]:
                        myname = "{}/{}".format(var, module.__dict__[var].index(self))

                # Register self as pytorch module.
                module._modules[myname] = self

                for name, param in self.named_variables.items():
                    if param.trainable:
                        module.register_parameter(myname + name, param.native)
                    else:
                        module.register_buffer(myname + name, param.native)
        else:
            if babilim.core.logging.DEBUG_VERBOSITY:
                _warn_once("babilim.model.module.Module:_register_params Not implemented for tf2 but I think it is not required.")
        
    @property
    def training(self) -> bool:
        """
        Property if the object is in training mode.
        
        ```python
        module.training
        ```
        
        :return: True if the object is in training mode.
        """
        return self._training

    @property
    def variables(self):
        """
        Property with all variables of the object.
        
        ```python
        module.variables
        ```
        
        :return: A list of the variables in this object.
        """
        return list(self.named_variables.values())

    @property
    def named_variables(self):
        """
        Property with all variables of the object.
        
        ```python
        module.named_variables
        ```
        
        :return: A dictionary of the variables in this object.
        """
        return dict(self.__variables_with_namespace())

    def __variables_with_namespace(self, namespace=""):
        all_vars = []
        extra_vars = []
        for member_name in self.__dict__:
            v = self.__dict__[member_name]
            if isinstance(v, str):
                pass
            elif isinstance(v, Dict):
                for i, (k, x) in enumerate(v.items()):
                    if not isinstance(k, str):
                        k = "{}".format(i)
                    name = namespace + "/" + member_name + "/" + k
                    if isinstance(x, Module):
                        all_vars.extend(x.__variables_with_namespace(name))
                    if isinstance(x, ITensor):
                        all_vars.append((name, x))
                    if self._wrapper.is_variable(x):
                        extra_vars.append((name, self._wrapper.wrap_variable(x)))
                    if isinstance(x, object):
                        extra_vars.extend(self._wrapper.vars_from_object(v, name))
            elif isinstance(v, Iterable):
                for i, x in enumerate(v):
                    name = namespace + "/" + member_name + "/{}".format(i)
                    if isinstance(x, Module):
                        all_vars.extend(x.__variables_with_namespace(name))
                    if isinstance(x, ITensor):
                        all_vars.append((name, x))
                    if self._wrapper.is_variable(x):
                        extra_vars.append((name, self._wrapper.wrap_variable(x)))
                    if isinstance(x, object):
                        extra_vars.extend(self._wrapper.vars_from_object(v, name))
            elif isinstance(v, Module):
                name = namespace + "/" + member_name
                all_vars.extend(v.__variables_with_namespace(name))
            elif isinstance(v, ITensor):
                name = namespace + "/" + member_name
                all_vars.append((name, v))
            # Native only counts if there is no babilim stuff.
            elif self._wrapper.is_variable(v):
                name = namespace + "/" + member_name
                extra_vars.append((name, self._wrapper.wrap_variable(v)))
            elif isinstance(v, object):
                name = namespace + "/" + member_name
                extra_vars.extend(self._wrapper.vars_from_object(v, name))
                for x in getattr(v, '__dict__', {}):
                    name = namespace + "/" + member_name + "/" + x
                    if isinstance(v.__dict__[x], Module):
                        all_vars.extend(v.__dict__[x].__variables_with_namespace(name))
                    if isinstance(v.__dict__[x], ITensor):
                        all_vars.append((name, v.__dict__[x]))
                    if self._wrapper.is_variable(v.__dict__[x]):
                        extra_vars.append((name, self._wrapper.wrap_variable(v.__dict__[x])))
        if len(all_vars) == 0:
            all_vars.extend(extra_vars)
        return all_vars

    @property
    def trainable_variables(self):
        """
        Property with trainable variables of the object.
        
        ```python
        module.trainable_variables
        ```
        
        :return: A list of the trainable variables in this object.
        """
        all_vars = self.variables
        train_vars = []
        for v in all_vars:
            if v.trainable:
                train_vars.append(v)
        return train_vars

    @property
    def named_trainable_variables(self):
        """
        Property with trainable variables of the object.
        
        ```python
        module.named_trainable_variables
        ```
        
        :return: A dictionary of the trainable variables in this object.
        """
        all_vars = self.named_variables
        train_vars = []
        for k, v in all_vars.items():
            if v.trainable:
                train_vars.append((k, v))
        return dict(train_vars)

    @property
    def untrainable_variables(self):
        """
        Property with not trainable variables of the object.
        
        ```python
        module.untrainable_variables
        ```
        
        :return: A list of not trainable variables in this object.
        """
        all_vars = self.variables
        train_vars = []
        for v in all_vars:
            if not v.trainable:
                train_vars.append(v)
        return train_vars

    @property
    def named_untrainable_variables(self):
        """
        Property with not trainable variables of the object.
        
        ```python
        module.named_untrainable_variables
        ```
        
        :return: A dictionary of not trainable variables in this object.
        """
        all_vars = self.named_variables
        train_vars = []
        for k, v in all_vars.items():
            if not v.trainable:
                train_vars.append((k, v))
        return dict(train_vars)

    @property
    def trainable_variables_native(self):
        """
        Property with not trainable variables of the object in native format.
        
        ```python
        module.trainable_variables_native
        ```
        
        :return: A list of trainable variables in this object in native format.
        """
        all_vars = self.trainable_variables
        train_vars = []
        for v in all_vars:
            train_vars.append(v.native)
        return train_vars

    @property
    def _parameters(self) -> OrderedDict:
        params = OrderedDict()
        params.update(self.named_trainable_variables)
        return params

    @property
    def _buffers(self) -> OrderedDict:
        params = OrderedDict()
        params.update(self.named_untrainable_variables)
        return params

    def state_dict(self) -> Dict:
        """
        Get the state of the object as a state dict (usable for checkpoints).
        
        :return: A dictionary containing the state of the object.
        """
        state = {}
        for name, var in self.named_variables.items():
            if babilim.is_backend(babilim.TF_BACKEND):
                state[name] = var.numpy()
            else:
                state[name] = var.numpy().T
        return state

    def load_state_dict(self, state_dict: Dict) -> None:
        """
        Load the state of the object from a state dict.
        
        Handy when loading checkpoints.
        
        :param state_dict: A dictionary containing the state of the object.
        """
        for name, var in self.named_variables.items():
            if name in state_dict:
                if babilim.is_backend(babilim.TF_BACKEND):
                    var.assign(state_dict[name])
                else:
                    var.assign(state_dict[name].T)
                if DEBUG_VERBOSITY:
                    info("  Loaded: {}".format(name))
            else:
                warn("  Variable {} not in checkpoint.".format(name))

    def eval(self):
        """
        Set the object into eval mode.
        
        ```python
        self.train(False)
        ```
        """
        self.train(False)

    def train(self, mode=True):
        """
        Set the objects training mode.
        
        :param mode: (Optional) If the training mode is enabled or disabled. (default: True)
        """
        self._training = mode
        for member_name in self.__dict__:
            obj = self.__dict__[member_name]
            if isinstance(obj, Sequence):
                for x in obj:
                    train_fn = getattr(x, "train", None)
                    if callable(train_fn):
                        train_fn(mode)
            else:
                train_fn = getattr(obj, "train", None)
                if callable(train_fn):
                    train_fn(mode)

    def load(self, checkpoint_file_path: str) -> None:
        """
        Load the state of the object from a checkpoint.
        
        :param checkpoint_file_path: The path to the checkpoint storing the state dict.
        """
        checkpoint = load_state(checkpoint_file_path)
        if "model" in checkpoint:
            self.load_state_dict(checkpoint["model"])
        else:
            babilim.error("Could not find state in checkpoint.")

    def save(self, checkpoint_file_path: str) -> None:
        """
        Save the state of the object to a checkpoint.
        
        :param checkpoint_file_path: The path to the checkpoint storing the state dict.
        """
        save_state({"model": self.state_dict()}, checkpoint_file_path)
