import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


device = "cuda" if torch.cuda.is_available() else "cpu"

class PositionalEncoding(nn.Module):
    def __init__(self, n_functions):
        super(PositionalEncoding, self).__init__()

        self.register_buffer('frequencies', 2.0 ** torch.arange(n_functions))

    def forward(self, x):
        """
        Args:
            x: tensor of shape [..., dim]

        Returns:
            embedding: a temporal embedding of `x` of shape [..., n_functions * dim * 2]
        """
        freq = (x[..., None] * self.frequencies).view(*x.shape[:-1], -1)

        embedding = torch.zeros(*freq.shape[:-1], freq.shape[-1] * 2).to(device)
        embedding[..., 0::2] = freq.sin()
        embedding[..., 1::2] = freq.cos()

        return embedding


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

        self.w0 = w0

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


class SirenBlock(nn.Module):
    def __init__(self, in_features, out_features, w0=30, c=6, is_first=False, use_bias=True, activation=None):
        super(SirenBlock, self).__init__()

        self.in_features = in_features
        self.is_first = is_first

        weight = torch.zeros(out_features, in_features)
        bias = torch.zeros(out_features) 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, bias, c, w0):
        n = self.in_features

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

        if bias is not None:
            bias.uniform_(-w_std, w_std)

    def forward(self, x):
        out = F.linear(x, self.weight, self.bias)

        return self.activation(out)


class FCBlock(nn.Module):
    def __init__(self, in_features, out_features, norm_layer=False, activation=None):
        super(FCBlock, self).__init__()

        self.fc = nn.Linear(in_features, out_features)
        self.residual = (in_features == out_features)  # when the input and output have the same dimensions, build a residual block
        self.norm_layer = nn.LayerNorm(out_features) if norm_layer else None
        self.activation = nn.ReLU(inplace=True) if activation is None else activation

    def forward(self, x):
        """
        Args:
            x: (B, T, D), features are in the last dimension.
        """
        out = self.fc(x)

        if self.norm_layer is not None:
            out = self.norm_layer(out)

        if self.residual:
            return self.activation(out) + x

        return self.activation(out)


class NeuralMotionField(nn.Module):
    def __init__(self, args):
        super(NeuralMotionField, self).__init__()

        self.args = args

        if args.bandwidth != 0:
            self.positional_encoding = PositionalEncoding(args.bandwidth)
            embedding_dim = args.bandwidth * 2
        else:
            embedding_dim = 1

        hidden_neuron = args.hidden_neuron
        in_features = 1 + args.local_z if args.siren else embedding_dim + args.local_z
        local_in_features = hidden_neuron + in_features if args.skip_connection else hidden_neuron
        global_in_features = local_in_features + args.global_z if args.skip_connection else hidden_neuron

        layers = [
            SirenBlock(in_features, hidden_neuron, is_first=True) if args.siren else FCBlock(in_features, hidden_neuron, norm_layer=args.norm_layer),
            SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
            SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
            SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
            SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
            SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
            SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
            SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
            SirenBlock(global_in_features, hidden_neuron) if args.siren else FCBlock(global_in_features, hidden_neuron, norm_layer=args.norm_layer),
            SirenBlock(global_in_features, hidden_neuron) if args.siren else FCBlock(global_in_features, hidden_neuron, norm_layer=args.norm_layer),
            SirenBlock(global_in_features, hidden_neuron) if args.siren else FCBlock(global_in_features, hidden_neuron, norm_layer=args.norm_layer)
        ]
        self.mlp = nn.ModuleList(layers)
        self.skip_layers = [] if not args.skip_connection else list(range(1, len(self.mlp)))
        self.local_layers = list(range(8))
        self.local_linear = nn.Sequential(nn.Linear(hidden_neuron, args.local_output))
        self.global_layers = list(range(8, len(self.mlp)))
        self.global_linear = nn.Sequential(nn.Linear(hidden_neuron, args.global_output))

    def forward(self, t, z_l=None, z_g=None):
        if self.args.bandwidth != 0 and self.args.siren is False:
            t = self.positional_encoding(t)

        if z_l is not None:
            z_l = z_l.unsqueeze(1).expand(-1, t.shape[1], -1)
            x = torch.cat([t, z_l], dim=-1)
        else:
            x = t

        skip_in = x
        for i in self.local_layers:
            if i in self.skip_layers:
                x = torch.cat((x, skip_in), dim=-1)
            x = self.mlp[i](x)
        local_output = self.local_linear(x)

        if z_g is not None:
            z_g = z_g.unsqueeze(1).expand(-1, t.shape[1], -1)
            for i in self.global_layers:
                if i in self.skip_layers:
                    x = torch.cat((x, skip_in, z_g), dim=-1)
                x = self.mlp[i](x)
            global_output = self.global_linear(x)
        else:
            global_output = None

        return local_output, global_output