"""
SIREN NETWORK
"""
import math
import torch
import torch.nn as nn
import numpy as np


class SineLayer(nn.Module):
    def __init__(self, in_features, out_features, bias=True,
                 is_first=False, omega_0=30):
        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, bias=bias)
        self.init_weights()

    def init_weights(self):
        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_(-math.sqrt(6 / self.in_features) / self.omega_0,
                                            math.sqrt(6 / self.in_features) / self.omega_0)

    def forward(self, x):
        y = torch.sin(self.omega_0 * self.linear(x))
        return y


class Siren(nn.Module):
    def __init__(self, in_features, hidden_features, hidden_layers, out_features, outermost_linear=False,
                 first_omega_0=30, hidden_omega_0=30., last_omega_0=1):
        super().__init__()

        self.net = []
        self.net.append(SineLayer(in_features, hidden_features,
                                  is_first=True, omega_0=first_omega_0))

        for i in range(hidden_layers):
            self.net.append(SineLayer(hidden_features, hidden_features,
                                      is_first=False, omega_0=hidden_omega_0))

        if outermost_linear:
            final_linear = nn.Linear(hidden_features, out_features)

            with torch.no_grad():
                final_linear.weight.uniform_(-math.sqrt(6 / hidden_features) / last_omega_0,
                                             math.sqrt(6 / hidden_features) / last_omega_0)

            self.net.append(final_linear)
        else:
            self.net.append(SineLayer(hidden_features, out_features,
                                      is_first=False, omega_0=hidden_omega_0))

        self.net = nn.Sequential(*self.net)
        self.hidden_features = hidden_features
        self.last_omega_0 = last_omega_0

    def forward(self, coords):
        output = self.net(coords)
        return output

    def reset(self):
        for layer in self.net:
            if isinstance(layer, SineLayer):
                layer.init_weights()

    def reset_last_layer(self):
        self.net[-1].weight.data.uniform_(-math.sqrt(6 / self.hidden_features) / self.last_omega_0,
                                          math.sqrt(6 / self.hidden_features) / self.last_omega_0)

    def freeze_sine_layers(self):
        for layer in self.net:
            if isinstance(layer, SineLayer):
                layer.requires_grad_ = False



def make_grid2d(height, width):
    lin_y = torch.linspace(-1., 1., height)
    lin_x = torch.linspace(-1., 1., width)
    grid_h, grid_w = torch.meshgrid([lin_y, lin_x])
    grid = torch.cat((grid_w[None, :, :, None], grid_h[None, :, :, None]), 3)
    return grid.view(1,-1,2)


class BottleNeck(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        #x is bnc
        avg = x.mean(dim=-1).unsqueeze(2)
        return torch.cat((x, avg), dim=2)


class SirenVision(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.sirens = nn.Sequential(
                Siren(2+32, 256, 10, 64),
                BottleNeck(),
                Siren(65, 256, 10, num_classes)
                # Siren(17, 128, 2, 64),
                # BottleNeck(),
                # Siren(65, 128, 2, 64),
                # BottleNeck(),
                # Siren(65, 64, 2, num_classes)
                )
        self.grid = None
        self.conv = nn.Conv2d(3, 32, 5, 1, 2)

    def forward(self, x):
        b,c,h,w = x.shape
        if self.grid is None:
            self.grid = make_grid2d(h,w).to(x)
        grid = self.grid.expand(b,h*w,2)
        x = self.conv(x)
        x = x.permute(0,2,3,1).reshape(b,-1,32)
        x = torch.cat((x,grid), dim=2)
        y = self.sirens(x).mean(dim=1)
        return y

# if __name__ == '__main__':
#     x = torch.randn(32,3,32,32)
#     net = SirenVision()
#     y = net(x)
#     print(y.shape)







"""
SIREN NETWORK
"""
import math
import torch
import torch.nn as nn
import numpy as np


class SineLayer(nn.Module):
    def __init__(self, in_features, out_features, bias=True,
                 is_first=False, omega_0=30):
        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, bias=bias)
        self.init_weights()

    def init_weights(self):
        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_(-math.sqrt(6 / self.in_features) / self.omega_0,
                                            math.sqrt(6 / self.in_features) / self.omega_0)

    def forward(self, x):
        y = torch.sin(self.omega_0 * self.linear(x))
        return y


class Siren(nn.Module):
    def __init__(self, in_features, hidden_features, hidden_layers, out_features, outermost_linear=False,
                 first_omega_0=30, hidden_omega_0=30., last_omega_0=1):
        super().__init__()

        self.net = []
        self.net.append(SineLayer(in_features, hidden_features,
                                  is_first=True, omega_0=first_omega_0))

        for i in range(hidden_layers):
            self.net.append(SineLayer(hidden_features, hidden_features,
                                      is_first=False, omega_0=hidden_omega_0))

        if outermost_linear:
            final_linear = nn.Linear(hidden_features, out_features)

            with torch.no_grad():
                final_linear.weight.uniform_(-math.sqrt(6 / hidden_features) / last_omega_0,
                                             math.sqrt(6 / hidden_features) / last_omega_0)

            self.net.append(final_linear)
        else:
            self.net.append(SineLayer(hidden_features, out_features,
                                      is_first=False, omega_0=hidden_omega_0))

        self.net = nn.Sequential(*self.net)
        self.hidden_features = hidden_features
        self.last_omega_0 = last_omega_0

    def forward(self, coords):
        output = self.net(coords)
        return output

    def reset(self):
        for layer in self.net:
            if isinstance(layer, SineLayer):
                layer.init_weights()

    def reset_last_layer(self):
        self.net[-1].weight.data.uniform_(-math.sqrt(6 / self.hidden_features) / self.last_omega_0,
                                          math.sqrt(6 / self.hidden_features) / self.last_omega_0)

    def freeze_sine_layers(self):
        for layer in self.net:
            if isinstance(layer, SineLayer):
                layer.requires_grad_ = False



def make_grid2d(height, width):
    lin_y = torch.linspace(-1., 1., height)
    lin_x = torch.linspace(-1., 1., width)
    grid_h, grid_w = torch.meshgrid([lin_y, lin_x])
    grid = torch.cat((grid_w[None, :, :, None], grid_h[None, :, :, None]), 3)
    return grid.view(1,-1,2)


class BottleNeck(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        #x is bnc
        avg = x.mean(dim=-1).unsqueeze(2)
        return torch.cat((x, avg), dim=2)


class SirenVision(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.sirens = nn.Sequential(
                Siren(2+32, 256, 10, 64),
                BottleNeck(),
                Siren(65, 256, 10, num_classes)
                # Siren(17, 128, 2, 64),
                # BottleNeck(),
                # Siren(65, 128, 2, 64),
                # BottleNeck(),
                # Siren(65, 64, 2, num_classes)
                )
        self.grid = None
        self.conv = nn.Conv2d(3, 32, 5, 1, 2)

    def forward(self, x):
        b,c,h,w = x.shape
        if self.grid is None:
            self.grid = make_grid2d(h,w).to(x)
        grid = self.grid.expand(b,h*w,2)
        x = self.conv(x)
        x = x.permute(0,2,3,1).reshape(b,-1,32)
        x = torch.cat((x,grid), dim=2)
        y = self.sirens(x).mean(dim=1)
        return y

# if __name__ == '__main__':
#     x = torch.randn(32,3,32,32)
#     net = SirenVision()
#     y = net(x)
#     print(y.shape)







SIREN NETWORK #x is bnc # Siren(17, 128, 2, 64), # BottleNeck(), # Siren(65, 128, 2, 64), # BottleNeck(), # Siren(65, 64, 2, num_classes) # if __name__ == '__main__': #     x = torch.randn(32,3,32,32) #     net = SirenVision() #     y = net(x) #     print(y.shape)