# AUTOGENERATED FROM: babilim/core/module_native.ipynb

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

> A module that is implemented by native function calls.
"""

# Cell: 1
from typing import Any

import babilim
from babilim import PYTORCH_BACKEND, TF_BACKEND, is_backend, get_backend
from babilim.core.annotations import RunOnlyOnce
from babilim.core.module import Module


# Cell: 2
class ModuleNative(Module):
    def __init__(self):
        """
        A module with a native implementation.
        
        This module is like a normal module, except that call and build call a "call_pytorch", "call_tf", "build_pytorch" and "build_tf" depending on what backend is set.
        """
        super().__init__()
        
    @RunOnlyOnce
    def build(self, *args, **kwargs) -> None:
        """
        Build the model, this function automatically calls the native build with the tensors unwrapped.

        This function gets called by `__call__` and itself passes all calls to `_build_pytorch` and `_build_tf`.
        Furthermore, it takes care of unwrapping the tensors into native tensors before calling and wrapping them again after calling.
        This allows the native functions `_build_pytorch` and `_build_tf` to be pure pytorch or tensorflow code.
        All subclasses must implement `_build_pytorch` and `_build_tf`.

        You should never call the build function directly. Call this module in the following style (this ensures the module is build on first run):
        ```
        module = MyModule()
        result = module(*args, **kwargs)  # <- Build gets called internally here.
        ```

        Parameters:
        :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.
        """
        args = self._wrapper.unwrap(args)
        kwargs = self._wrapper.unwrap(kwargs)
        if is_backend(PYTORCH_BACKEND):
            self._build_pytorch(*args, **kwargs)
        elif is_backend(TF_BACKEND):
            self._build_tf(*args, **kwargs)
        else:
            raise RuntimeError("Unknown Backend: {}".format(get_backend()))
            
    def _build_pytorch(self, *args, **kwargs) -> None:
        """
        A native build function in pytorch.
        
        Even though babilim never calls this function directly multiple times, it is recommended to add the RunOnlyOnce guard in case a user calls it multiple times.
        
        :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 _build_tf(self, *args, **kwargs) -> None:
        """
        A native build function in tensorflow.
        
        Even though babilim never calls this function directly multiple times, it is recommended to add the RunOnlyOnce guard in case a user calls it multiple times.
        
        :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 itself passes all calls to `_call_pytorch` and `_call_tf`.
        Furthermore, it takes care of unwrapping the tensors into native tensors before calling and wrapping them again after calling.
        This allows the native functions `_call_pytorch` and `_call_tf` to be pure pytorch or tensorflow code.
        All subclasses must implement `_call_pytorch` and `_call_tf`.

        You should call this module in the following style (this ensures the module is build on first run):
        ```
        module = MyModule()
        result = module(*args, **kwargs)
        ```

        Parameters:
        :param *args: You can specify any parameters you want.
        :param **kwargs: You can specify any named parameters you want.
        """
        args = self._wrapper.unwrap(args)
        kwargs = self._wrapper.unwrap(kwargs)
        if is_backend(PYTORCH_BACKEND):
            results = self._call_pytorch(*args, **kwargs)
        elif is_backend(TF_BACKEND):
            results = self._call_tf(*args, **kwargs)
        else:
            raise RuntimeError("Unknown Backend: {}".format(get_backend()))
        
        results = self._wrapper.wrap(results)
        return results
    
    def _call_pytorch(self, *args, **kwargs) -> Any:
        """
        A native call function in pytorch (like the forward).
        
        :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.
        """
        raise NotImplemented()
    
    def _call_tf(self, *args, **kwargs) -> Any:
        """
        A native call function in tensorflow.
        
        :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.
        """
        raise NotImplemented()
