
import scipy.integrate as integrate
from scipy.optimize import minimize
from scipy.misc import derivative
import numpy as np
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import linear_operator
from linear_operator.operators.diag_linear_operator import DiagLinearOperator
from linear_operator.operators.dense_linear_operator import DenseLinearOperator
from utils import (
    extendtensor, memoize, power_iteration_eig, compute_maxeig,
    BatchDenseLinearOperator, activation_derivative, timethis,
    extremal_sigs_noq, calculate_parameters, calculate_expected_eigenvalue
)


class ActivLinear(nn.Module):
    """
    Class for a linear layer followed by an activation, which is automatically rescaled.
    Takes care of measuring all important layer-wise measures we are interested in.
    In particular returns the Lipschitz constant of the layer, its conditioning with respect
    to weights and inputs, and the distance of the current weights from the initialization.
    """
    def __init__(self, in_features, out_features, activation=nn.Identity(), bias=True, renormalize=False, **kwargs):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.fc = nn.Linear(in_features, out_features, bias=bias)
        if isinstance(activation, str):
            activation = getattr(nn, activation)()
        self.activation_name = activation.__class__.__name__
        actmean, actvar, actlip, _ = calculate_parameters(activation)
        # Then we rescale the activation to have mean zero and variance one
        # And we remember the rescaled lipschitz constant
        self.activation = lambda x: (activation(x) - actmean) / math.sqrt(actvar)
        self.actvar = actvar
        self.actlipschitz = actlip / math.sqrt(actvar)
        self.renormalize = renormalize

    def init_weights_(self):
        """
        Initializes the weights of the Linear function and saves them.
        """
        # Provides a one-to-one scaling for the forward variance, since we are already scaling the activation
        nn.init.normal_(self.fc.weight, mean=0.0, std=1.0 / math.sqrt(self.in_features))
        nn.init.constant_(self.fc.bias, 0.0)
        self.initial_weight = self.fc.weight.clone().detach()
        self.initial_bias = self.fc.bias.clone().detach()

    def cache_initial_inputs(self):
        # TODO: Change when measuring SGD
        self.initial_inputs = self._cache_inputs.clone().detach()
        
    def forward(self, x):
        # We expect x to follow a standard gaussian.
        # And we cache results for later calculations.
        self._cache_inputs = x
        self._cache_activations = self.activation(x)
        if self.renormalize:
            # We renormalize such that the total norm is as we should expect in independent standard gaussians
            self._cache_activations = math.sqrt(self.in_features) * F.normalize(self._cache_activations, p=2, dim=1)
        return self.fc(self._cache_activations)

    def measure_input_matrix(self) -> float:
        """
        Returns the smallest on-diagonal entry, highest off-diagonal entry, the max of the sum of the off-diagonal entries, the mean of the sum of the off-diagonal entries, the max eigenvalue and the min eigenvalue of the kernel matrix.
        The on-diagonal entries are always one since we normalize inputs at each layer.
        """
        entry_matrix = (self._cache_activations @ self._cache_activations.t()).clone().detach() / self.in_features
        min_ondiag = torch.min(torch.diag(entry_matrix)).item()
        max_offdiagonal = torch.max(torch.abs(entry_matrix - torch.diag(torch.diag(entry_matrix)))).item()
        sumoff = torch.sum(torch.abs(entry_matrix - torch.diag(torch.diag(entry_matrix))), dim=1)
        max_sumoff = sumoff.max().item()
        mean_sumoff = sumoff.mean().item()
        op = BatchDenseLinearOperator(self._cache_activations, 1)
        maxeig, mineig = op.extremal_sig()
        return min_ondiag, max_offdiagonal, max_sumoff, mean_sumoff, maxeig, mineig
    
    def measure_approximate_input_conditioning(self) -> (float, float):
        """
        Returns the lowest and highest eigenvalues of the conditioning wrt the input.
        To be called before updating alpha parameters!

        Our computed function is $f(a, x) = a \phi(x)$, so that the gradient has to be
        $\nabla_x f(a, x)[x'] = a (\phi'(x) \odot x')$.
        """
        batch_size, _ = self._cache_inputs.size()
        # The diagonal operator corresponding to the activation gradients
        grads = activation_derivative(self.activation_name, self._cache_inputs.detach())
        # But we need gradients on the scaled activation:
        grads /= math.sqrt(self.actvar)
        if self.out_features >= self.in_features:
            # We create the linear operator from the linear layer
            linearop = BatchDenseLinearOperator(self.fc.weight, batch_size)
            diagop = DiagLinearOperator(grads.view(-1))
            maxeig, mineig = (linearop @ diagop).extremal_sig()
        else:
            # In the other case we have to recur to a more expensive computation
            matrix = self.fc.weight.unsqueeze(0).expand(batch_size, -1, -1)
            # and we have to postmultiply by the diagonal matrix
            newmatrix = matrix * grads.unsqueeze(1).expand(-1, self.out_features, -1)
            maxeig, mineig = extremal_sigs_noq(newmatrix)
        return mineig, maxeig

    def measure_norms(self) -> (float, float, float, float):
        """
        Returns the 2-norm and infty-norm of the weight matrix, both as matrix and as vector.
        """
        # We have to collect norms of matrices to test the theory about worse conditioning when moving
        # away from the initialization, which are bounded by the distance in matrix norm.
        maxeig = compute_maxeig(self.fc.weight)
        infnorm = torch.linalg.norm(self.fc.weight, ord=float("inf")).item()
        norm2 = torch.norm(self.fc.weight.view(-1)).item()
        norminf = torch.norm(self.fc.weight.view(-1), p=float("inf")).item()
        return maxeig, infnorm, norm2, norminf
        
    def measure_distances(self) -> dict:
        """
        Returns the distances of the current weights from their initialization,
        and from the originally cached inputs in a dictionary.
        Distances are taken wrt l2 and linf norm.
        """
        # We return here the distance of weights in their own space!!
        res = {}
        res["wdistance"] = torch.norm((self.fc.weight - self.initial_weight).view(-1)).item()
        res["bdistance"] = torch.norm((self.fc.bias - self.initial_bias).view(-1)).item()
        res["total"] = math.sqrt(res["wdistance"]**2 + res["bdistance"]**2)
        res["wdistanceinf"] = torch.norm((self.fc.weight - self.initial_weight).view(-1), float("inf")).item()
        res["bdistanceinf"] = torch.norm((self.fc.bias - self.initial_bias).view(-1), float("inf")).item()
        res["totalinf"] = max(res["wdistanceinf"], res["bdistanceinf"])
        res["idistance"] = torch.norm((self._cache_inputs - self.initial_inputs).view(-1)).item()
        res["idistanceinf"] = torch.norm((self._cache_inputs - self.initial_inputs).view(-1), float("inf")).item()
        res["imaxsingledistance"] = torch.norm(self._cache_inputs - self.initial_inputs, dim=1).max().item()
        res["imaxsingledistanceinf"] = torch.norm(self._cache_inputs - self.initial_inputs, p=float("inf"), dim=1).max().item()
        # Measure also the distance as matrices between the initial weights and the current ones
        diffweight = self.fc.weight - self.initial_weight
        x0 = torch.randn(self.out_features, dtype=self.fc.weight.dtype, device=self.fc.weight.device)
        res["wmdistance"] = power_iteration_eig(lambda x: diffweight @ (diffweight.t() @ x), x0)
        res["wmdistanceinf"] = torch.norm(diffweight, float("inf")).item()
        return res
