import torch
from torch import nn as nn
import numpy as np


def hyper_weight_init(m, in_features_main_net, activation):
    if hasattr(m, 'weight'):
        nn.init.kaiming_normal_(m.weight, a=0.0, nonlinearity='relu', mode='fan_in')
        m.weight.data = m.weight.data / 1.e2

    if hasattr(m, 'bias'):
        with torch.no_grad():
            if activation == 'sine':
                m.bias.uniform_(-np.sqrt(6 / in_features_main_net) / 30, np.sqrt(6 / in_features_main_net) / 30)
            elif activation == 'leakyrelu_pe':
                m.bias.uniform_(-np.sqrt(6 / in_features_main_net), np.sqrt(6 / in_features_main_net))
            else:
                raise NotImplementedError


class ConvBlock(nn.Module):
    def __init__(
            self,
            in_channels, out_channels,
            kernel_size=4, stride=2, padding=1,
            norm_layer=nn.BatchNorm2d, activation=nn.ELU,
            bias=True,
    ):
        super(ConvBlock, self).__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias),
            norm_layer(out_channels) if norm_layer is not None else nn.Identity(),
            activation(),
        )

    def forward(self, x):
        return self.block(x)


class MaxPoolDownSize(nn.Module):
    def __init__(self, in_channels, mid_channels, out_channels, depth):
        super(MaxPoolDownSize, self).__init__()
        self.depth = depth
        self.reduce_conv = ConvBlock(in_channels, mid_channels, kernel_size=1, stride=1, padding=0)
        self.convs = nn.ModuleList([
            ConvBlock(mid_channels, out_channels, kernel_size=3, stride=1, padding=1)
            for conv_i in range(depth)
        ])
        self.pool2d = nn.MaxPool2d(kernel_size=2)

    def forward(self, x):
        outputs = []

        output = self.reduce_conv(x)

        for conv_i, conv in enumerate(self.convs):
            output = output if conv_i == 0 else self.pool2d(output)
            outputs.append(conv(output))

        return outputs


class convParams(nn.Module):
    def __init__(self, input_dim, INR_in_out, opt, hidden_mlp_num, hidden_dim=512, toRGB=False):
        super(convParams, self).__init__()
        self.INR_in_out = INR_in_out
        self.cont_split_weight = []
        self.cont_split_bias = []
        self.hidden_mlp_num = hidden_mlp_num
        self.param_factorize_dim = opt.param_factorize_dim
        output_dim = self.cal_params_num(INR_in_out, hidden_mlp_num, toRGB)
        self.output_dim = output_dim
        self.toRGB = toRGB
        self.cont_extraction_net = nn.Sequential(
            nn.Conv2d(input_dim, hidden_dim, kernel_size=3, stride=2, padding=1, bias=False),
            # nn.BatchNorm2d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, stride=1, padding=1, bias=False),
            # nn.BatchNorm2d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Conv2d(hidden_dim, output_dim, kernel_size=1, stride=1, padding=0, bias=True),
        )

        self.cont_extraction_net[-1].apply(lambda m: hyper_weight_init(m, INR_in_out[0], opt.activation))

        self.basic_params = nn.ParameterList()
        if opt.param_factorize_dim > 0:
            for id in range(self.hidden_mlp_num + 1):
                if id == 0:
                    inp, outp = self.INR_in_out[0], self.INR_in_out[1]
                else:
                    inp, outp = self.INR_in_out[1], self.INR_in_out[1]
                self.basic_params.append(nn.Parameter(torch.randn(1, 1, 1, inp, outp)))

            if toRGB:
                self.basic_params.append(nn.Parameter(torch.randn(1, 1, 1, self.INR_in_out[1], 3)))

    def forward(self, feat, outMore=False):
        cont_params = self.cont_extraction_net(feat)
        out_mlp = self.to_mlp(cont_params)
        if outMore:
            return out_mlp, cont_params
        return out_mlp

    def cal_params_num(self, INR_in_out, hidden_mlp_num, toRGB=False):
        cont_params = 0
        start = 0
        if self.param_factorize_dim == -1:
            cont_params += INR_in_out[0] * INR_in_out[1] + INR_in_out[1]
            self.cont_split_weight.append([start, cont_params - INR_in_out[1]])
            self.cont_split_bias.append([cont_params - INR_in_out[1], cont_params])
            start = cont_params

            for id in range(hidden_mlp_num):
                cont_params += INR_in_out[1] * INR_in_out[1] + INR_in_out[1]
                self.cont_split_weight.append([start, cont_params - INR_in_out[1]])
                self.cont_split_bias.append([cont_params - INR_in_out[1], cont_params])
                start = cont_params

            if toRGB:
                cont_params += INR_in_out[1] * 3 + 3
                self.cont_split_weight.append([start, cont_params - 3])
                self.cont_split_bias.append([cont_params - 3, cont_params])

        elif self.param_factorize_dim > 0:
            cont_params += INR_in_out[0] * self.param_factorize_dim + self.param_factorize_dim * INR_in_out[1] + \
                           INR_in_out[1]
            self.cont_split_weight.append(
                [start, start + INR_in_out[0] * self.param_factorize_dim, cont_params - INR_in_out[1]])
            self.cont_split_bias.append([cont_params - INR_in_out[1], cont_params])
            start = cont_params

            for id in range(hidden_mlp_num):
                cont_params += INR_in_out[1] * self.param_factorize_dim + self.param_factorize_dim * INR_in_out[1] + \
                               INR_in_out[1]
                self.cont_split_weight.append(
                    [start, start + INR_in_out[1] * self.param_factorize_dim, cont_params - INR_in_out[1]])
                self.cont_split_bias.append([cont_params - INR_in_out[1], cont_params])
                start = cont_params

            if toRGB:
                cont_params += INR_in_out[1] * self.param_factorize_dim + self.param_factorize_dim * 3 + 3
                self.cont_split_weight.append(
                    [start, start + INR_in_out[1] * self.param_factorize_dim, cont_params - 3])
                self.cont_split_bias.append([cont_params - 3, cont_params])

        return cont_params

    def to_mlp(self, params):
        all_weight_bias = []
        if self.param_factorize_dim == -1:
            for id in range(self.hidden_mlp_num + 1):
                if id == 0:
                    inp, outp = self.INR_in_out[0], self.INR_in_out[1]
                else:
                    inp, outp = self.INR_in_out[1], self.INR_in_out[1]
                weight = params[:, self.cont_split_weight[id][0]:self.cont_split_weight[id][1], :, :]
                weight = weight.permute(0, 2, 3, 1).contiguous().view(weight.shape[0], *weight.shape[2:],
                                                                      inp, outp)

                bias = params[:, self.cont_split_bias[id][0]:self.cont_split_bias[id][1], :, :]
                bias = bias.permute(0, 2, 3, 1).contiguous().view(bias.shape[0], *bias.shape[2:], 1, outp)
                all_weight_bias.append([weight, bias])

            if self.toRGB:
                inp, outp = self.INR_in_out[1], 3
                weight = params[:, self.cont_split_weight[-1][0]:self.cont_split_weight[-1][1], :, :]
                weight = weight.permute(0, 2, 3, 1).contiguous().view(weight.shape[0], *weight.shape[2:],
                                                                      inp, outp)

                bias = params[:, self.cont_split_bias[-1][0]:self.cont_split_bias[-1][1], :, :]
                bias = bias.permute(0, 2, 3, 1).contiguous().view(bias.shape[0], *bias.shape[2:], 1, outp)
                all_weight_bias.append([weight, bias])

            return all_weight_bias

        else:
            for id in range(self.hidden_mlp_num + 1):
                if id == 0:
                    inp, outp = self.INR_in_out[0], self.INR_in_out[1]
                else:
                    inp, outp = self.INR_in_out[1], self.INR_in_out[1]
                weight1 = params[:, self.cont_split_weight[id][0]:self.cont_split_weight[id][1], :, :]
                weight1 = weight1.permute(0, 2, 3, 1).contiguous().view(weight1.shape[0], *weight1.shape[2:],
                                                                        inp, self.param_factorize_dim)

                weight2 = params[:, self.cont_split_weight[id][1]:self.cont_split_weight[id][2], :, :]
                weight2 = weight2.permute(0, 2, 3, 1).contiguous().view(weight2.shape[0], *weight2.shape[2:],
                                                                        self.param_factorize_dim, outp)

                bias = params[:, self.cont_split_bias[id][0]:self.cont_split_bias[id][1], :, :]
                bias = bias.permute(0, 2, 3, 1).contiguous().view(bias.shape[0], *bias.shape[2:], 1, outp)

                all_weight_bias.append([torch.tanh(torch.matmul(weight1, weight2)) * self.basic_params[id], bias])

            if self.toRGB:
                inp, outp = self.INR_in_out[1], 3
                weight1 = params[:, self.cont_split_weight[-1][0]:self.cont_split_weight[-1][1], :, :]
                weight1 = weight1.permute(0, 2, 3, 1).contiguous().view(weight1.shape[0], *weight1.shape[2:],
                                                                        inp, self.param_factorize_dim)

                weight2 = params[:, self.cont_split_weight[-1][1]:self.cont_split_weight[-1][2], :, :]
                weight2 = weight2.permute(0, 2, 3, 1).contiguous().view(weight2.shape[0], *weight2.shape[2:],
                                                                        self.param_factorize_dim, outp)

                bias = params[:, self.cont_split_bias[-1][0]:self.cont_split_bias[-1][1], :, :]
                bias = bias.permute(0, 2, 3, 1).contiguous().view(bias.shape[0], *bias.shape[2:], 1, outp)

                all_weight_bias.append([torch.tanh(torch.matmul(weight1, weight2)) * self.basic_params[-1], bias])

            return all_weight_bias


class lineParams(nn.Module):
    def __init__(self, input_dim, INR_in_out, input_resolution, opt, hidden_mlp_num, toRGB=False,
                 hidden_dim=512):
        super(lineParams, self).__init__()
        self.INR_in_out = INR_in_out
        self.app_split_weight = []
        self.app_split_bias = []
        self.toRGB = toRGB
        self.hidden_mlp_num = hidden_mlp_num
        self.param_factorize_dim = opt.param_factorize_dim
        output_dim = self.cal_params_num(INR_in_out, hidden_mlp_num)
        self.output_dim = output_dim

        self.compress_layer = nn.Sequential(
            nn.Linear(input_resolution, 64, bias=False),
            nn.BatchNorm1d(input_dim),
            nn.ReLU(inplace=True),
            nn.Linear(64, 1, bias=True)
        )

        self.app_extraction_net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim, bias=False),
            # nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, hidden_dim, bias=False),
            # nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, output_dim, bias=True)
        )

        self.app_extraction_net[-1].apply(lambda m: hyper_weight_init(m, INR_in_out[0], opt.activation))

        self.basic_params = nn.ParameterList()
        if opt.param_factorize_dim > 0:
            for id in range(self.hidden_mlp_num + 1):
                if id == 0:
                    inp, outp = self.INR_in_out[0], self.INR_in_out[1]
                else:
                    inp, outp = self.INR_in_out[1], self.INR_in_out[1]
                self.basic_params.append(nn.Parameter(torch.randn(1, inp, outp)))
            if toRGB:
                self.basic_params.append(nn.Parameter(torch.randn(1, self.INR_in_out[1], 3)))

    def forward(self, feat):
        app_params = self.app_extraction_net(self.compress_layer(torch.flatten(feat, 2)).squeeze(-1))
        out_mlp = self.to_mlp(app_params)
        return out_mlp, app_params

    def cal_params_num(self, INR_in_out, hidden_mlp_num):
        app_params = 0
        start = 0
        if self.param_factorize_dim == -1:
            app_params += INR_in_out[0] * INR_in_out[1] + INR_in_out[1]
            self.app_split_weight.append([start, app_params - INR_in_out[1]])
            self.app_split_bias.append([app_params - INR_in_out[1], app_params])
            start = app_params

            for id in range(hidden_mlp_num):
                app_params += INR_in_out[1] * INR_in_out[1] + INR_in_out[1]
                self.app_split_weight.append([start, app_params - INR_in_out[1]])
                self.app_split_bias.append([app_params - INR_in_out[1], app_params])
                start = app_params

            if self.toRGB:
                app_params += INR_in_out[1] * 3 + 3
                self.app_split_weight.append([start, app_params - 3])
                self.app_split_bias.append([app_params - 3, app_params])

        elif self.param_factorize_dim > 0:
            app_params += INR_in_out[0] * self.param_factorize_dim + self.param_factorize_dim * INR_in_out[1] + \
                          INR_in_out[1]
            self.app_split_weight.append([start, start + INR_in_out[0] * self.param_factorize_dim,
                                          app_params - INR_in_out[1]])
            self.app_split_bias.append([app_params - INR_in_out[1], app_params])
            start = app_params

            for id in range(hidden_mlp_num):
                app_params += INR_in_out[1] * self.param_factorize_dim + self.param_factorize_dim * INR_in_out[1] + \
                              INR_in_out[1]
                self.app_split_weight.append(
                    [start, start + INR_in_out[1] * self.param_factorize_dim, app_params - INR_in_out[1]])
                self.app_split_bias.append([app_params - INR_in_out[1], app_params])
                start = app_params

            if self.toRGB:
                app_params += INR_in_out[1] * self.param_factorize_dim + self.param_factorize_dim * 3 + 3
                self.app_split_weight.append([start, start + INR_in_out[1] * self.param_factorize_dim,
                                              app_params - 3])
                self.app_split_bias.append([app_params - 3, app_params])

        return app_params

    def to_mlp(self, params):
        all_weight_bias = []
        if self.param_factorize_dim == -1:
            for id in range(self.hidden_mlp_num + 1):
                if id == 0:
                    inp, outp = self.INR_in_out[0], self.INR_in_out[1]
                else:
                    inp, outp = self.INR_in_out[1], self.INR_in_out[1]
                weight = params[:, self.app_split_weight[id][0]:self.app_split_weight[id][1]]
                weight = weight.view(weight.shape[0], inp, outp)

                bias = params[:, self.app_split_bias[id][0]:self.app_split_bias[id][1]]
                bias = bias.view(bias.shape[0], 1, outp)

                all_weight_bias.append([weight, bias])

            if self.toRGB:
                id = -1
                inp, outp = self.INR_in_out[1], 3
                weight = params[:, self.app_split_weight[id][0]:self.app_split_weight[id][1]]
                weight = weight.view(weight.shape[0], inp, outp)

                bias = params[:, self.app_split_bias[id][0]:self.app_split_bias[id][1]]
                bias = bias.view(bias.shape[0], 1, outp)

                all_weight_bias.append([weight, bias])

            return all_weight_bias

        else:
            for id in range(self.hidden_mlp_num + 1):
                if id == 0:
                    inp, outp = self.INR_in_out[0], self.INR_in_out[1]
                else:
                    inp, outp = self.INR_in_out[1], self.INR_in_out[1]
                weight1 = params[:, self.app_split_weight[id][0]:self.app_split_weight[id][1]]
                weight1 = weight1.view(weight1.shape[0], inp, self.param_factorize_dim)

                weight2 = params[:, self.app_split_weight[id][1]:self.app_split_weight[id][2]]
                weight2 = weight2.view(weight2.shape[0], self.param_factorize_dim, outp)

                bias = params[:, self.app_split_bias[id][0]:self.app_split_bias[id][1]]
                bias = bias.view(bias.shape[0], 1, outp)

                all_weight_bias.append([torch.tanh(torch.matmul(weight1, weight2)) * self.basic_params[id], bias])

            if self.toRGB:
                id = -1
                inp, outp = self.INR_in_out[1], 3
                weight1 = params[:, self.app_split_weight[id][0]:self.app_split_weight[id][1]]
                weight1 = weight1.view(weight1.shape[0], inp, self.param_factorize_dim)

                weight2 = params[:, self.app_split_weight[id][1]:self.app_split_weight[id][2]]
                weight2 = weight2.view(weight2.shape[0], self.param_factorize_dim, outp)

                bias = params[:, self.app_split_bias[id][0]:self.app_split_bias[id][1]]
                bias = bias.view(bias.shape[0], 1, outp)

                all_weight_bias.append([torch.tanh(torch.matmul(weight1, weight2)) * self.basic_params[id], bias])

            return all_weight_bias