
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from linear_operator.operators.diag_linear_operator import DiagLinearOperator
from linear_module import ActivLinear, calculate_parameters, calculate_expected_eigenvalue
from utils import extendtensor, memoize, power_iteration_eig, ConvolutionOperator, activation_derivative


# NOTE: We remove padding because zeros on the borders lower the variance of the entries!
class ActivConvolutional(ActivLinear):
    """
    Class for a convolutional layer followed by an activation, which is automatically rescaled.
    Takes care of measuring all important layer-wise measures we are interested in.
    """
    def __init__(self, in_features, out_features, kernel_size, image_size, activation=nn.Identity(), bias=True, renormalize=False):
        super(ActivLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.image_size = image_size
        in_pixels = image_size[0] * image_size[1]
        self.in_pixels = in_pixels
        self.kernel_size = kernel_size
        if self.kernel_size % 2 == 0:
            raise Exception(f"Even kernel_size {self.kernel_size}!")
        self.conv = nn.Conv2d(in_features, out_features, kernel_size, bias=bias)
        if isinstance(activation, str):
            activation = getattr(nn, activation)()
        self.activation_name = activation.__class__.__name__
        actmean, actvar, actlip, _ = calculate_parameters(activation)
        # Rescale activation to have zero mean and unitary variance
        self.activation = lambda x: (activation(x) - actmean) / math.sqrt(actvar)
        self.actlipschitz = actlip / math.sqrt(actvar)
        self.renormalize = renormalize

    def init_weights_(self):
        """
        Initializes the weights of the Convolutional layer.
        """
        # We need to initialize kxk weights to preserve input variance, thus 1/sqrt(k^2) = 1/k.
        nn.init.normal_(self.conv.weight, mean=0.0, std=1.0 / self.kernel_size)
        nn.init.constant_(self.conv.bias, 0.0)
        self.initial_weight = self.conv.weight.clone().detach()
        self.initial_bias = self.conv.bias.clone().detach()

    # To reuse many of the other functions it is enough to redefine
    # self._cache_inputs and self._cache_activations to a flat view
    def forward(self, x):
        # inputs are in the form [batch, in_channels, height, width]
        batch = x.size()[0]
        self._cache_inputs = x
        x = self.activation(x)
        self._cache_activations = x.view(batch, -1)
        if self.renormalize:
            self._cache_activations = math.sqrt(self.in_features * self.in_pixels) * F.normalize(self._cache_activations, p=2, dim=1)
            out = self._cache_activations.view(*x.size())
            x = out
        return self.conv(x)

    def measure_approximate_input_conditioning(self) -> (float, float):
        # TODO: Not possible to use this since we always have output_dim <= input_dim
        batch_size, _, in_height, in_width = self._cache_inputs.size()
        # To measure approximate conditioning, we create the given convolution operator
        convop = ConvolutionOperator(self.conv.weight, batch_size, (in_height, in_width))
        # And the diagonal operator correponding to the gradients of the activations
        grads = activation_derivative(self.activation_name, self._cache_inputs.detach())
        # grads are batch x (in_channels * height * width), as expected by ConvolutionOperator
        # We create a diagonal linear operator with the gradients on the diagonal
        diagop = DiagLinearOperator(grads.view(-1))
        # We then create Conv @ Diag and compute its extremal singular values.
        maxeig, mineig = (convop @ diagop).extremal_sig()
        return mineig, maxeig
    
    def measure_norms(self) -> (float, float):
        """
        Measure the norms of the weight matrices as convolutions.
        The inf-norm is calculated using the maximum entry, while 2-norm by an iterative eigenvalue problem.
        """
        infnorm = self.conv.weight.max().abs().item()
        # The 2-norm is calculated by iterating convolution and convolution with transposed filter
        x0 = torch.randn(1, self.in_features, *self.image_size, dtype=self.conv.weight.dtype, device=self.conv.weight.device)
        def Amul(x0):
            x0 = F.conv2d(x0, self.conv.weight)
            return F.conv_transpose2d(x0, self.conv.weight)
        maxeig = math.sqrt(power_iteration_eig(Amul, x0))
        return maxeig, infnorm

    def measure_distances(self) -> dict:
        """
        Returns the distances of the current weights from their initialization in a dictionary.
        This returns distance of weights as weights, not as matrices!
        """
        # Obviously those weights act on multiple pixels at the same time, so that the distance
        # between the "expanded" weights has to be multiplied by the number of pixels in the
        # output image (these are the number of times that the weights are repeated).
        res = {}
        res["wdistance"] = torch.norm((self.conv.weight - self.initial_weight).view(-1)).item()
        res["bdistance"] = torch.norm((self.conv.bias - self.initial_bias).view(-1)).item()
        res["total"] = math.sqrt(res["wdistance"]**2 + res["bdistance"]**2)
        res["wdistanceinf"] = torch.norm((self.conv.weight - self.initial_weight).view(-1), float("inf")).item()
        res["bdistanceinf"] = torch.norm((self.conv.bias - self.initial_bias).view(-1), float("inf")).item()
        res["totalinf"] = max(res["wdistanceinf"], res["bdistanceinf"])
        return res


