import math
import torch
import numpy as np
from torch import nn
from typing import Optional, Tuple
import torch.nn.functional as F

# Adapted from https://github.com/lucidrains/siren-pytorch
# Licensed under MIT by Phil Wang

# helpers


def exists(val) -> bool:
    return val is not None


def cast_tuple(val, repeat: int = 1) -> Tuple:
    return val if isinstance(val, tuple) else ((val,) * repeat)


# sin activation


class Sine(nn.Module):
    def __init__(self, w0: float = 1.0):
        super().__init__()
        self.w0 = w0

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.sin(self.w0 * x)


# siren layer


class Siren(nn.Module):
    def __init__(
        self,
        dim_in: int,
        dim_out: int,
        w0: float = 1.0,
        c: float = 6.0,
        is_first: bool = False,
        use_bias: bool = True,
        activation: Optional[nn.Module] = None,
    ):
        super().__init__()
        self.dim_in = dim_in
        self.is_first = is_first

        weight = torch.zeros(dim_out, dim_in)
        bias = torch.zeros(dim_out) if use_bias else None
        self.init_(weight, bias, c=c, w0=w0)

        self.weight = nn.Parameter(weight)
        self.bias = nn.Parameter(bias) if use_bias else None
        self.activation = Sine(w0) if activation is None else activation

    def init_(self, weight: torch.Tensor, bias: Optional[torch.Tensor], c: float, w0: float) -> None:
        dim = self.dim_in

        w_std = (1 / dim) if self.is_first else (math.sqrt(c / dim) / w0)
        weight.uniform_(-w_std, w_std)

        if exists(bias):
            bias.uniform_(-w_std, w_std)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.linear(x, self.weight, self.bias)
        out = self.activation(out)
        return out


class SirenNet(nn.Module):
    def __init__(
        self,
        dim_in: int,
        dim_hidden: int,
        dim_out: int,
        num_layers: int,
        w0: float = 1.0,
        w0_initial: float = 30.0,
        use_bias: bool = True,
        final_activation: Optional[nn.Module] = None,
    ):
        super().__init__()
        self.num_layers = num_layers
        self.dim_hidden = dim_hidden

        self.layers = nn.ModuleList()
        for ind in range(num_layers):
            is_first = ind == 0
            layer_w0 = w0_initial if is_first else w0
            layer_dim_in = dim_in if is_first else dim_hidden

            self.layers.append(
                Siren(
                    dim_in=layer_dim_in,
                    dim_out=dim_hidden,
                    w0=layer_w0,
                    use_bias=use_bias,
                    is_first=is_first,
                )
            )

        final_activation = nn.Identity() if not exists(final_activation) else final_activation
        self.last_layer = Siren(
            dim_in=dim_hidden,
            dim_out=dim_out,
            w0=w0,
            use_bias=use_bias,
            activation=final_activation,
        )

    def forward(self, x: torch.Tensor, mods: Optional[Tuple[torch.Tensor, ...]] = None) -> torch.Tensor:
        bs, num, in_dim = x.size()

        if mods is not None:
            mods = cast_tuple(mods, self.num_layers)
            for layer, mod in zip(self.layers, mods):
                x = layer(x)
                x *= mod.view(bs, 1, -1)
        else:
            for layer in self.layers:
                x = layer(x)

        return self.last_layer(x)


# modulatory feed forward


class Modulator(nn.Module):
    def __init__(self, dim_in: int, dim_hidden: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList([])

        for ind in range(num_layers):
            is_first = ind == 0
            dim = dim_in if is_first else (dim_hidden + dim_in)

            self.layers.append(nn.Sequential(nn.Linear(dim, dim_hidden), nn.Sigmoid()))

    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, ...]:
        x = z
        hiddens = []

        for layer in self.layers:
            x = layer(x)
            hiddens.append(x)
            x = torch.cat((x, z), dim=-1)

        return tuple(hiddens)
