""" This files contains the function approximators that are used by DeepMoD. 
Note that the output of the function approximator has to be a Torch module. 
"""


import torch
import torch.nn as nn
from typing import List
import numpy as np


class NN(nn.Module):
    def __init__(self, n_in: int, n_hidden: List[int], n_out: int) -> None:
        """ This class sets up a simple feed forward neural network as function approximator.

        Args:
            n_in (int): Number of input features.
            n_hidden (List[int]): Number of neurons in each layer.
            n_out (int): Number of output features.
        """
        super().__init__()
        self.network = self.build_network(n_in, n_hidden, n_out)

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """Performs a forward pass through the network. 

        Args:
            input (torch.Tensor): Input tensor. 

        Returns:
            torch.Tensor: Output tensor. 
        """
        coordinates = input.clone().detach().requires_grad_(True)
        return self.network(coordinates), coordinates

    def build_network(self, n_in: int, n_hidden: List[int], n_out: int) -> torch.nn.Sequential:
        """ Constructs a feed-forward neural network.

        Args:
            n_in (int): Number of input features.
            n_hidden (list[int]): Number of neurons in each layer.
            n_out (int): Number of output features.

        Returns:
            torch.Sequential: Pytorch module
        """

        network = []
        architecture = [n_in] + n_hidden + [n_out]
        for layer_i, layer_j in zip(architecture, architecture[1:]):
            network.append(nn.Linear(layer_i, layer_j))
            network.append(nn.Tanh())
        network.pop()  # get rid of last activation function
        return nn.Sequential(*network)


class SineLayer(nn.Module):
    def __init__(self, in_features: int, out_features: int, omega_0: float = 30, is_first: bool = False):
        """ Sine activation function with omega_0 scaling. 

        Args:
            in_features (int): Number of input features.
            out_features (int): Number of output features.
            omega_0 (float, optional): Scaling factor of the Sine function. Defaults to 30.
            is_first (bool, optional): Defaults to False.
        """
        super().__init__()
        self.omega_0 = omega_0
        self.is_first = is_first

        self.in_features = in_features
        self.linear = nn.Linear(in_features, out_features)

        self.init_weights()

    def init_weights(self):
        """Initialization of the weigths. 
        """
        with torch.no_grad():
            if self.is_first:
                self.linear.weight.uniform_(-1 / self.in_features, 1 / self.in_features)
            else:
                self.linear.weight.uniform_(-np.sqrt(6 / self.in_features) / self.omega_0,
                                            np.sqrt(6 / self.in_features) / self.omega_0)

    def forward(self, input):
        """Performs the forward pass through the network. 

        Args:
            input (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor.
        """
        return torch.sin(self.omega_0 * self.linear(input))


class Siren(nn.Module):
    def __init__(self, n_in: int, n_hidden: List[int], n_out: int, first_omega_0: float = 30., hidden_omega_0: float = 30.):
        """ SIREN model from the paper [Implicit Neural Representations with
        Periodic Activation Functions](https://arxiv.org/abs/2006.09661).

        Args:
            n_in (int): Number of input features.
            n_hidden (list[int]): Number of neurons in each layer.
            n_out (int): Number of output features.
            first_omega_0 (float, optional): Scaling factor of the Sine function of the first layer. Defaults to 30.
            hidden_omega_0 (float, optional): Scaling factor of the Sine function of the hidden layers. Defaults to 30.
        """
        super().__init__()
        self.network = self.build_network(n_in, n_hidden, n_out, first_omega_0, hidden_omega_0)

    def forward(self, input: torch.Tensor):
        """Performs the forward pass through the network. 

        Args:
            input (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor.
        """
        coordinates = input.clone().detach().requires_grad_(True)
        return self.network(coordinates), coordinates

    def build_network(self, n_in: int, n_hidden: List[int], n_out: int, first_omega_0: float, hidden_omega_0: float):
        """Constructs the Siren neural network.

        Args:
            n_in (int): Number of input features.
            n_hidden (list[int]): Number of neurons in each layer.
            n_out (int): Number of output features.
            first_omega_0 (float, optional): Scaling factor of the Sine function of the first layer. Defaults to 30.
            hidden_omega_0 (float, optional): Scaling factor of the Sine function of the hidden layers. Defaults to 30.
        Returns:
            torch.Sequential: Pytorch module
        """
        network = []
        # Input layer
        network.append(SineLayer(n_in, n_hidden[0], is_first=True, omega_0=first_omega_0))

        # Hidden layers
        for layer_i, layer_j in zip(n_hidden, n_hidden[1:]):
            network.append(SineLayer(layer_i, layer_j, is_first=False, omega_0=hidden_omega_0))

        # Output layer
        final_linear = nn.Linear(n_hidden[-1], n_out)
        with torch.no_grad():
            final_linear.weight.uniform_(-np.sqrt(6 / n_hidden[-1]) / hidden_omega_0,
                                         np.sqrt(6 / n_hidden[-1]) / hidden_omega_0)
            network.append(final_linear)

        return nn.Sequential(*network)
