"""FiLM Siren MLP as per https://marcoamonteiro.github.io/pi-GAN-website/."""

from typing import Optional

import numpy as np
import torch
from torch import nn


def kaiming_leaky_init(m):
    classname = m.__class__.__name__
    if classname.find("Linear") != -1:
        torch.nn.init.kaiming_normal_(
            m.weight, a=0.2, mode="fan_in", nonlinearity="leaky_relu"
        )


def frequency_init(freq):
    def init(m):
        with torch.no_grad():
            if isinstance(m, nn.Linear):
                num_input = m.weight.size(-1)
                m.weight.uniform_(
                    -np.sqrt(6 / num_input) / freq, np.sqrt(6 / num_input) / freq
                )

    return init


def first_layer_film_sine_init(m):
    with torch.no_grad():
        if isinstance(m, nn.Linear):
            num_input = m.weight.size(-1)
            m.weight.uniform_(-1 / num_input, 1 / num_input)


class CustomMappingNetwork(nn.Module):
    def __init__(self, in_features, map_hidden_layers, map_hidden_dim, map_output_dim):
        super().__init__()

        self.network = []

        for _ in range(map_hidden_layers):
            self.network.append(nn.Linear(in_features, map_hidden_dim))
            self.network.append(nn.LeakyReLU(0.2, inplace=True))
            in_features = map_hidden_dim

        self.network.append(nn.Linear(map_hidden_dim, map_output_dim))

        self.network = nn.Sequential(*self.network)

        self.network.apply(kaiming_leaky_init)
        with torch.no_grad():
            self.network[-1].weight *= 0.25

    def forward(self, z):
        frequencies_offsets = self.network(z)
        frequencies = frequencies_offsets[
            ..., : torch.div(frequencies_offsets.shape[-1], 2, rounding_mode="floor")
        ]
        phase_shifts = frequencies_offsets[
            ..., torch.div(frequencies_offsets.shape[-1], 2, rounding_mode="floor") :
        ]

        return frequencies, phase_shifts


class FiLMLayer(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.layer = nn.Linear(input_dim, hidden_dim)

    def forward(self, x, freq, phase_shift):
        x = self.layer(x)
        freq = freq.expand_as(x)
        phase_shift = phase_shift.expand_as(x)
        return torch.sin(freq * x + phase_shift)


class FiLMSiren(nn.Module):
    """FiLM Conditioned Siren network."""

    def __init__(
        self,
        in_dim: int,
        hidden_layers: int,
        hidden_features: int,
        mapping_network_in_dim: int,
        mapping_network_layers: int,
        mapping_network_features: int,
        out_dim: int,
        outermost_linear: bool = False,
        out_activation: Optional[nn.Module] = None,
    ) -> None:
        super().__init__()
        self.in_dim = in_dim
        assert self.in_dim > 0
        self.out_dim = out_dim if out_dim is not None else hidden_features
        self.hidden_layers = hidden_layers
        self.hidden_features = hidden_features
        self.mapping_network_in_dim = mapping_network_in_dim
        self.mapping_network_layers = mapping_network_layers
        self.mapping_network_features = mapping_network_features
        self.outermost_linear = outermost_linear
        self.out_activation = out_activation

        self.net = nn.ModuleList()

        self.net.append(FiLMLayer(self.in_dim, self.hidden_features))

        for _ in range(self.hidden_layers - 1):
            self.net.append(FiLMLayer(self.hidden_features, self.hidden_features))

        self.final_layer = None
        if self.outermost_linear:
            self.final_layer = nn.Linear(self.hidden_features, self.out_dim)
            self.final_layer.apply(frequency_init(25))
        else:
            final_layer = FiLMLayer(self.hidden_features, self.out_dim)
            self.net.append(final_layer)

        self.mapping_network = CustomMappingNetwork(
            in_features=self.mapping_network_in_dim,
            map_hidden_layers=self.mapping_network_layers,
            map_hidden_dim=self.mapping_network_features,
            map_output_dim=(len(self.net)) * self.hidden_features * 2,
        )

        self.net.apply(frequency_init(25))
        self.net[0].apply(first_layer_film_sine_init)

    def forward_with_frequencies_phase_shifts(self, x, frequencies, phase_shifts):
        """Get conditiional frequencies and phase shifts from mapping network."""
        frequencies = frequencies * 15 + 30

        for index, layer in enumerate(self.net):
            start = index * self.hidden_features
            end = (index + 1) * self.hidden_features
            x = layer(x, frequencies[..., start:end], phase_shifts[..., start:end])

        x = self.final_layer(x) if self.final_layer is not None else x
        output = self.out_activation(x) if self.out_activation is not None else x
        return output

    def forward(self, x, conditioning_input):
        """Forward pass."""
        frequencies, phase_shifts = self.mapping_network(conditioning_input)
        return self.forward_with_frequencies_phase_shifts(x, frequencies, phase_shifts)