
import torch
import math
import linear_operator as lop
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
from time import time
from contextlib import contextmanager
from typing import Optional, Tuple
import scipy.integrate as integrate
from scipy.optimize import minimize
from scipy.misc import derivative
from linear_operator.operators.dense_linear_operator import DenseLinearOperator
from linear_operator.settings import cg_tolerance
from warnings import warn


if torch.__version__ == "1.7.1":
    # Substitute with method that is instead found only in version 1.8.0
    torch.linalg.svd = torch.svd


def sign(x: torch.Tensor) -> torch.Tensor:
    "Returns -1 if the entries are negative and +1 otherwise"
    return torch.where(x >= 0.0, 1.0, -1.0)


@torch.jit.script
def vecmmul(v: torch.Tensor, A: torch.Tensor) -> torch.Tensor:
    """
    Performs multiplication v v^T A, where

    A: batch x dim1 x dim2
    v: batch x dim1
    returns: batch x dim1 x dim2
    """
    batch, dim1, dim2 = A.size()
    vta = (v[:, :, None] * A[:, :, :]).sum(1) # batch x dim2
    res = v[:, :, None] * vta[:, None, :]
    return res


@torch.jit.script
def qr_decomposition(A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Returns Q, R of the QR decomposition of A computed via Householder reflections.
    """
    single = True if len(A.size()) == 2 else False
    if single:
        A = A.unsqueeze(0)
    
    batches, outdim, indim = A.size()
    R = A.clone()
    Qt = torch.zeros((batches, outdim, outdim), dtype=A.dtype, device=A.device)
    for k in range(outdim):
        Qt[:, k, k] = 1.0

    # Allocate needed space outside of for loop to cut malloc calls
    ebase = torch.zeros(outdim, dtype=A.dtype, device=A.device)
    ebase[0] = 1.0
    accumsigns = torch.zeros(batches, dtype=A.dtype, device=A.device)
    accumvk = torch.zeros((batches, outdim), dtype=A.dtype, device=A.device)
    accumnorm = torch.zeros(batches, dtype=A.dtype, device=A.device)
    accummult = torch.zeros((batches, outdim), dtype=A.dtype, device=A.device)
    for k in range(indim):
        x = R[:, k:, k] # batches x (outdim - k)
        accumsigns = torch.where(x[:, 0] >= 0.0, 1.0, -1.0)
        accumvk[:, :(outdim-k)] = accumsigns[:, None]
        accumnorm = torch.norm(x, dim=-1)
        accumvk[:, :(outdim-k)] *= accumnorm[:, None]
        accumvk[:, :(outdim-k)] *= ebase[None, :(outdim-k)]
        accumvk[:, :(outdim-k)] += x
        accumnorm = torch.norm(accumvk[:, :(outdim-k)], dim=-1)
        accumvk[:, :(outdim-k)] /= accumnorm[:, None]
        # Now multiply R[:, k:, k:] by accumvk[:, :(outdim-k)] as vk (vk^T R) 
        # R[:, k:, k:] = R[:, k:, k:] - 2 * vecmmul(vk, R[:, k:, k:])
        accummult[:, :(indim-k)] = -2.0 * (accumvk[:, :(outdim-k), None] * R[:, k:, k:]).sum(1)
        R[:, k:, k:] += accumvk[:, :(outdim-k), None] * accummult[:, None, :(indim-k)]
        # And we also compute Qt in the same way as R
        # Qt[:, k:, :] = Qt[:, k:, :] - 2 * vecmmul(accumvk[:, :(outdim-k)], Qt[:, k:, :])
        accummult[:, :] = -2.0 * (accumvk[:, :(outdim-k), None] * Qt[:, k:, :]).sum(1)
        Qt[:, k:, :] += accumvk[:, :(outdim-k), None] * accummult[:, None, :]

    Qt = Qt.permute(0, 2, 1)
    if single:
        Qt = Qt.squeeze(0)
        R = R.squeeze(0)
    return Qt, R


@torch.jit.script
def r_of_qr_decomposition(A: torch.Tensor) -> torch.Tensor:
    """
    Returns only R of the QR decomposition of A computed via Householder reflections.
    """
    single = True if len(A.size()) == 2 else False
    if single:
        A = A.unsqueeze(0)
    
    batches, outdim, indim = A.size()
    R = A.clone()

    # Allocate needed space outside of for loop to cut malloc calls
    ebase = torch.zeros(outdim, dtype=A.dtype, device=A.device)
    ebase[0] = 1.0
    accumsigns = torch.zeros(batches, dtype=A.dtype, device=A.device)
    accumvk = torch.zeros((batches, outdim), dtype=A.dtype, device=A.device)
    accumnorm = torch.zeros(batches, dtype=A.dtype, device=A.device)
    accummult = torch.zeros((batches, outdim), dtype=A.dtype, device=A.device)
    for k in range(indim):
        x = R[:, k:, k] # batches x (outdim - k)
        accumsigns = torch.where(x[:, 0] >= 0.0, 1.0, -1.0)
        accumvk[:, :(outdim-k)] = accumsigns[:, None]
        accumnorm = torch.norm(x, dim=-1)
        accumvk[:, :(outdim-k)] *= accumnorm[:, None]
        accumvk[:, :(outdim-k)] *= ebase[None, :(outdim-k)]
        accumvk[:, :(outdim-k)] += x
        accumnorm = torch.norm(accumvk[:, :(outdim-k)], dim=-1)
        accumvk[:, :(outdim-k)] /= accumnorm[:, None]
        # Now multiply R[:, k:, k:] by accumvk[:, :(outdim-k)] as vk (vk^T R) 
        # R[:, k:, k:] = R[:, k:, k:] - 2 * vecmmul(vk, R[:, k:, k:])
        accummult[:, :(indim-k)] = -2.0 * (accumvk[:, :(outdim-k), None] * R[:, k:, k:]).sum(1)
        R[:, k:, k:] += accumvk[:, :(outdim-k), None] * accummult[:, None, :(indim-k)]

    if single:
        R = R.squeeze(0)
    return R


@torch.jit.script
def triangular_solve(b: torch.Tensor, A: torch.Tensor, upper: bool = True) -> torch.Tensor:
    """
    We solve the triangular system AX = B.
    """
    single = True if len(A.size()) == 2 else False
    if single:
        A = A.unsqueeze(0)
        b = b.unsqueeze(0)
    
    # b: batch x size x postbatch, A: batch x size x size, out: batch x size x postbatch
    batch, size, postbatch = b.size()
    assert A.size() == (batch, size, size)

    x = torch.zeros((batch, size, postbatch), device=b.device, dtype=b.dtype)
    if upper:
        for i in range(size-1, -1, -1):
            dotprod = (A[:, i, i+1:, None] * x[:, i+1:, :]).sum(-2) # dotprod: batch x postbatch
            x[:, i, :] = (b[:, i, :] - dotprod) / A[:, i, i, None]
    else:
        for i in range(size):
            dotprod = (A[:, i, :i, None] * x[:, :i, :]).sum(-2) # dotprod: batch x postbatch
            x[:, i, :] = (b[:, i, :] - dotprod) / A[:, i, i, None]

    if single:
        x = x.squeeze(0)
    return x


def extendtensor(tensor, size):
    """
    Extends a matrix with zeros to be of the given sizes.
    """
    msize = tensor.size()
    newm = torch.zeros(size).to(tensor.device)
    it = [slice(0, minsize) for minsize in msize]
    newm[it] = tensor
    return newm


def power_iteration_eig(Amul, x0, prec=1e-10, maxstep=100):
    """
    Method that computes the highest eigenvalue of A via power iteration.
    
    @param Amul: A function accepting x and returning Ax
    @param x0: The starting point for the power iteration
    @param prec: The required precision on the distance from the real eigenvalue
    @param maxstep: The maximum amount of steps of optimization to perform

    @ret The calculated approximation of the eigenvalue
    """
    eig = float("inf")
    x2 = x0
    x2x2 = 1.0
    x0x0 = 1.0
    x2x0 = 0.0
    step = 0
    while x2x2 / x0x0 - (x2x0 / x0x0)**2 > prec**2:
        x0 = x2 / torch.norm(x2).item()
        x2 = Amul(x0)
        x2x2 = torch.dot(x2.view(-1), x2.view(-1)).item()
        x2x0 = torch.dot(x2.view(-1), x0.view(-1)).item()
        x0x0 = max(1e-15, torch.dot(x0.view(-1), x0.view(-1)).item())
        eig = x2x0 / x0x0
        step += 1
        if step > maxstep:
            warn(f"Reached maximum amount of {maxstep} permitted steps in eigenvalue research.")
            break
    return eig


def power_iteration_eig_matrix(A: torch.Tensor, x0: torch.Tensor, prec : float =1e-8, maxstep : int=100):
    """
    Method that computes the highest eigenvalue of A via power iteration.
    
    @param Amul: A function accepting x and returning Ax
    @param x0: The starting point for the power iteration
    @param prec: The required precision on the distance from the real eigenvalue
    @param maxstep: The maximum amount of steps of optimization to perform

    @ret The calculated approximation of the eigenvalue
    """
    eig = float("inf")
    x2 = x0
    x2x2 = 1.0
    x0x0 = 1.0
    x2x0 = 0.0
    step = 0
    while x2x2 / x0x0 - (x2x0 / x0x0)**2 > prec**2 and step <= maxstep:
        x0 = x2 / torch.norm(x2).item()
        x2 = A @ x0
        x2x2 = torch.dot(x2.view(-1), x2.view(-1)).item()
        x2x0 = torch.dot(x2.view(-1), x0.view(-1)).item()
        x0x0 = max(1e-15, torch.dot(x0.view(-1), x0.view(-1)).item())
        eig = x2x0 / x0x0
        step += 1
    return eig


def compute_maxeig(matrix: torch.Tensor) -> float:
    height, width = matrix.size()
    if width == height:
        matrix = matrix @ matrix.t()
        x0 = torch.randn(width, dtype=matrix.dtype, device=matrix.device)
        return math.sqrt(power_iteration_eig_matrix(matrix, x0))
    # Otherwise we consider the minimum length and compute the eigenvalues of the crossproduct
    if height > width:
        matrix = matrix.t()
    # Now height < width so that we can compute the smallest
    matrix = matrix @ matrix.t()
    x0 = torch.randn(matrix.size()[0], dtype=matrix.dtype, device=matrix.device)
    return math.sqrt(power_iteration_eig_matrix(matrix, x0))


def memoize(fn):
    """
    Memoize decorator for a function of a single argument.
    """
    cache = {}
    def wrapped(p):
        pc = p
        if isinstance(p, type):
            pc = p.__name__
        elif isinstance(p, str):
            pc = p
        elif isinstance(p, object):
            pc = p.__class__.__name__
        if pc not in cache:
            cache[pc] = fn(p)
        return cache[pc]
    return wrapped


# Measures that are saved in files that need to be added to DVC
saved_measures = ["distance", "general", "iconditioning", "input_matrix", "norms"]


def experiment_slug(args):
    "Generates the name of the file under which the experiment metrics are saved."
    run_slug = []
    for argument in ["dataset", "examples", "network", "loss_function", "n_layers", "width",
                     "batch_size", "lr", "activation", "epochs", "kernel_size", "seed", "renormalize"]:
        if argument == "kernel_size" and args.network == "FCN":
            continue
        elif argument == "batch_size" and not hasattr(args, argument):
            argval = args.examples
        else:
            argval = getattr(args, argument)
        run_slug.append(f"{argument}_{argval}")
    run_slug = "@".join(run_slug)
    return run_slug


def one_hot_encode(target: torch.Tensor, classes: int) -> torch.Tensor:
    """
    Encodes the original targets of size [batch]::int into a matrix of size [batch x classes]::float
    which has ones exactly where the targets should be.
    """
    batch = target.size()[0]
    expanded_target = target.unsqueeze(-1).expand(-1, classes)
    ones = torch.ones(batch, classes, device=target.device)
    return torch.zeros(batch, classes, device=target.device).scatter(dim=1, index=expanded_target, src=ones)


def extremal_sig(self, maxstep=100) -> (float, float):
    """
    Uses the definitions of LinearOperators to compute the extremal singular values of the operator,
    excluding the ones that would be zero because of the dimension).
    
    These correspond to the square roots of the eigenvalues of self @ self.t()
    """
    outdim, indim = self.size()
    if self.__class__ != BatchDenseLinearOperator:
        # This is the only case in which it works for the product of dense with diagonal
        assert outdim >= indim
    # Compute the smallest resulting matrix for smaller computation and discarding zero eigenvalues.
    A = self @ self.t() if outdim < indim else self.t() @ self
    s, t = A.size()
    assert s == t
    x0 = torch.randn(t, device=A.device)
    maxeig = power_iteration_eig(lambda x: A @ x, x0, maxstep=maxstep)
    invmineig = power_iteration_eig(lambda x: A.inv_matmul(x.unsqueeze(-1)).squeeze(-1), x0, maxstep=maxstep)
    mineig = 1.0 / invmineig if abs(invmineig) >= 1e-20 else float("NaN")
    if maxeig < -0.1 or mineig < -0.1:
        raise Exception(f"maxeig {maxeig} or mineig {mineig} are very negative")
    return math.sqrt(max(maxeig, 0.0)), math.sqrt(max(0.0, mineig))

lop.LinearOperator.extremal_sig = extremal_sig


def inv_matmul(self, right_tensor, left_tensor=None):
    # We have (LR)^{-1} = R^{-1}L^{-1}
    partial = self.left_linear_operator.inv_matmul(right_tensor)
    return self.right_linear_operator.inv_matmul(partial, left_tensor)

lop.operators.MatmulLinearOperator.inv_matmul = inv_matmul


class ConvolutionOperator(lop.LinearOperator):
    """
    A linear operator that represents the action of convolution on images.
    Does not support padding nor strides.
    """
    def __init__(self, weight, batch_size, input_size, transposed=False):
        self._args = (weight,)
        self._kwargs = {}
        self.weight = weight.clone().detach().requires_grad_(False)
        self.transposed = transposed
        self.input_size = input_size
        self.batch = batch_size
        if not self.transposed:
            out_channels, in_channels, kH, kW = weight.size()
            self.output_size = (input_size[0] - (kH - 1), input_size[1] - (kW - 1))
        else:
            in_channels, out_channels, kH, kW = weight.size()
            self.output_size = (input_size[0] + (kH - 1), input_size[1] + (kW - 1))
        self.mshape = (
            batch_size * out_channels * self.output_size[0] * self.output_size[1],
            batch_size * in_channels * input_size[0] * input_size[1],
        )

    def _size(self):
        return torch.Size(self.mshape)

    def _transpose_nonbatch(self):
        return ConvolutionOperator(self.weight, self.batch, self.output_size, transposed=True)
    
    def _matmul(self, rhs):
        """
        Multiplying it by another matrix means applying convolution to a batch
        """
        # rhs [batch * in_channel * input[0] * input[1], rhsbatch]
        _, rhsbatch = rhs.size()
        if not self.transposed:
            _, in_channels, _, _ = self.weight.size()
        else:
            in_channels, _, _, _ = self.weight.size()
        iH, iW = self.input_size
        rhs = rhs.view(self.batch, in_channels, iH, iW, rhsbatch)
        rhs = rhs.permute(4, 0, 1, 2, 3).reshape(rhsbatch * self.batch, in_channels, iH, iW)
        if not self.transposed:
            out = F.conv2d(rhs, self.weight)
        else:
            out = F.conv_transpose2d(rhs, self.weight)
        # Return out to its originally expected dimensions
        out = out.view(rhsbatch, -1).permute(1, 0)
        return out


class BatchDenseLinearOperator(lop.LinearOperator):
    """
    A linear operator that represents the action of a linear layer.
    """
    def __init__(self, weight, batch_size):
        self._args = (weight,)
        self._kwargs = dict(batch_size=batch_size)
        self.weight = weight.clone().detach().requires_grad_(False)
        self.batch = batch_size
        output_size, input_size = self.weight.size()
        self.output_size = output_size
        self.input_size = input_size
        self.mshape = (
            batch_size * output_size,
            batch_size * input_size,
        )

    def to_dense(self):
        dense = self.weight.unsqueeze(-1).expand(-1, -1, self.batch)
        dense = torch.diag_embed(dense, dim1=0, dim2=2).view(*self.mshape)
        return dense

    def _size(self):
        return torch.Size(self.mshape)

    def _base_qr(self):
        if hasattr(self, "_computed_qr"):
            return self._computed_qr
        self._computed_qr = qr_decomposition(self.weight)
        return self._computed_qr
    
    def _base_t_qr(self):
        if hasattr(self, "_computed_t_qr"):
            return self._computed_t_qr
        self._computed_t_qr = qr_decomposition(self.weight.t())
        return self._computed_t_qr

    def _transpose_nonbatch(self):
        return BatchDenseLinearOperator(self.weight.t(), self.batch)

    def inv_matmul(self, right_tensor, left_tensor=None):
        if left_tensor is not None:
            raise NotImplementedError("inv_matmul with left_tensor is not implemented.")
        # Resize input tensor
        *prebatches, total_inner, postbatches = right_tensor.size()
        if total_inner != self.batch * self.output_size:
            raise Exception(f"Unmatched dimensions {total_inner} and {self.batch_size} * {self.output_size}")
        new_right_tensor = right_tensor.view(-1, self.output_size, postbatches)
        curbatches = new_right_tensor.size()[0]
        # Depending whether our output size or input size is bigger we need different decompositions
        if self.input_size < self.output_size:
            # Overdetermined: least squares
            Q, R = self._base_qr()
            result = least_squares_solution(new_right_tensor,
                                            R.unsqueeze(0).expand(curbatches, -1, -1),
                                            Q.unsqueeze(0).expand(curbatches, -1, -1))

        else:
            # Underdetermined: minimum norm
            Q, R = self._base_t_qr()
            result = minimum_norm_solution(new_right_tensor,
                                           R.unsqueeze(0).expand(curbatches, -1, -1),
                                           Q.unsqueeze(0).expand(curbatches, -1, -1))
        return result.view(*prebatches, self.batch * self.input_size, postbatches)
    
    def _matmul(self, rhs):
        """
        Multiplying it by another matrix means applying the layer to a batch
        """
        # rhs [batch * input_size, rhsbatch]
        _, rhsbatch = rhs.size()
        output_size, input_size = self.weight.size()
        rhs = rhs.view(self.batch, input_size, rhsbatch).permute(2, 0, 1).reshape(rhsbatch * self.batch, input_size)
        out = F.linear(rhs, self.weight)
        # Return out to its originally expected dimensions
        out = out.view(rhsbatch, -1).permute(1, 0)
        return out


def activation_derivative(activation, tensor_input):
    """
    Define derivatives for the activation functions that we will test.
    This is done for memory issues in using torch.autograd.functional.jvp
    which requires to pass v=ones, but we don't have enough memory to do it.
    """
    if not isinstance(activation, str):
        activation = activation.__name__
    # We pass activation functions only from torch.nn so names are cased
    if activation == "Identity":
        return torch.ones_like(tensor_input)
    elif activation == "ReLU":
        return torch.where(tensor_input >= 0.0, 1.0, 0.0)
    elif activation == "Tanh":
        return 1.0 / torch.cosh(tensor_input)**2
    else:
        raise Exception(f"Unknown derivative of activation {activation}")


@contextmanager
def timethis(name="<Unnamed>", report=True):
    """
    Context manager for timing a computation.
    """
    start = time()
    yield
    end = time()
    if report:
        print(f"Computation {name} took {end - start} seconds")


def minimum_norm_solution(b, R, Q):
    """
    Solves the minimum norm solution of Ax = b given the QR factorization of A^T.
    The system has to be underdetermined, i.e. A.out_dim <= A.in_dim.
    """
    # Q: prebatch x in x in, R: prebatch x in x out, b: prebatch x out x postbatch
    # We have to compute QR^{-T}b, and to do it we can restrict R to the first square,
    # and the solution will have zeros on the other parts.
    prebatch, in_dim, out_dim = R.size()
    assert out_dim <= in_dim
    ccc, binner, postbatch = b.size()
    assert binner == out_dim
    assert ccc == prebatch
    if Q is not None:
        assert Q.size() == (prebatch, in_dim, in_dim)
    solution = torch.zeros((prebatch, in_dim, postbatch),
                           dtype=b.dtype, device=b.device)
    multR = R[:, :out_dim, :].permute(0, 2, 1)
    solution[:, :out_dim, :] = triangular_solve(b, multR, upper=False)
    if Q is not None:
        solution = torch.bmm(Q, solution)
    assert solution.size() == (prebatch, in_dim, postbatch)
    return solution


def least_squares_solution(b, R, Q):
    """
    Solves the least square solution of |Ax - b| given the QR factorization of A.
    The system has to be overdetermined, i.e. A.out_dim >= A.in_dim.
    """
    # Q: prebatch x out x out, R: prebatch x out x in, b: prebatch x out x postbatch
    # We have first to compute c = Q^Tb, and then solve |Rx - c| as least square,
    # which consists in solving the restricted system with R.
    _, out_dim, in_dim = R.size()
    assert out_dim >= in_dim
    prebatch, binner, postbatch = b.size()
    assert binner == out_dim
    if Q is not None:
        assert Q.size() == (prebatch, out_dim, out_dim)
        c = torch.bmm(Q.permute(0, 2, 1), b)[:, :in_dim, :]
    else:
        c = b[:, :in_dim, :]
    solution = triangular_solve(c, R[:, :in_dim, :], upper=True)
    assert solution.size() == (prebatch, in_dim, postbatch)
    return solution


def extremal_sigs(matrix):
    "Calculates extremal singular values of matrix: batch x d1 x d2"
    # We have to use iterative methods for max and min singular value,
    # using the two functions above. We will need QR decomposition of
    # matrix and of its transpose.
    batch, d1, d2 = matrix.size()
    # Ensure d1 <= d2 to not have to deal with multiple cases
    if d1 > d2:
        matrix = matrix.permute(0, 2, 1)
        batch, d1, d2 = matrix.size()
    assert d1 <= d2
    # The smallest matrix product is then A = matrix @ matrix.t()
    # First we compute the maximal eigenvalue which is easy
    x0 = torch.randn((batch, d1, 1), dtype=matrix.dtype, device=matrix.device)
    def Amul(x0):
        x0 = torch.bmm(matrix.permute(0, 2, 1), x0)
        x0 = torch.bmm(matrix, x0)
        return x0
    maxeig = power_iteration_eig(Amul, x0)
    # Then we compute the QR decomposition and use it to find alternatively
    # the minimum norm solution and the least square solution to converge
    # to the minimum eigenvalue.
    Q_trs, R_trs = qr_decomposition(matrix.permute(0, 2, 1))
    def Ainv(x0):
        # We have to evaluate inversion of matrix @ matrix.t()
        # We first invert matrix, which is underdetermined, so minimum norm (with A^T = QR)
        x0 = minimum_norm_solution(x0, R_trs, Q_trs)
        # Then to invert matrix.t(), which is overdetermined, we need least squares (A = QR)
        x0 = least_squares_solution(x0, R_trs, Q_trs)
        return x0
    mineig = 1.0 / max(1e-30, power_iteration_eig(Ainv, x0))
    if mineig < -0.1 or maxeig < -0.1:
        raise Exception(f"Either maxsig {maxeig} or minsig {mineig} are very negative.")
    mineig = max(1e-18, mineig)
    maxeig = max(1e-18, maxeig)
    return math.sqrt(maxeig), math.sqrt(mineig)


def extremal_sigs_noq(matrix):
    "Calculates extremal singular values of matrix: batch x d1 x d2 using R-only QR decomposition"
    batch, d1, d2 = matrix.size()
    # Ensure d1 <= d2 to not have to deal with multiple cases
    if d1 > d2:
        matrix = matrix.permute(0, 2, 1)
        batch, d1, d2 = matrix.size()
    assert d1 <= d2
    # The smallest matrix product is then A = matrix @ matrix.t()
    # First we compute the maximal eigenvalue which is easy
    x0 = torch.randn((batch, d1, 1), dtype=matrix.dtype, device=matrix.device)
    def Amul(x0):
        x0 = torch.bmm(matrix.permute(0, 2, 1), x0)
        x0 = torch.bmm(matrix, x0)
        return x0
    maxeig = power_iteration_eig(Amul, x0)
    # Then we compute the R-only QR decomposition and use it to find alternatively
    # the minimum norm solution and the least square solution to converge to minsig.
    R_trs = r_of_qr_decomposition(matrix.permute(0, 2, 1))
    def Ainv(x0):
        # We have to evaluate inversion of matrix @ matrix.t() = R^T @ R: batch x d1 -> batch x d1
        # We first invert matrix, which is underdetermined
        x0 = minimum_norm_solution(x0, R_trs, Q=None)
        # Then to invert matrix.t(), which is overdetermined
        x0 = least_squares_solution(x0, R_trs, Q=None)
        return x0
    mineig = 1.0 / max(1e-30, power_iteration_eig(Ainv, x0))
    if mineig < -0.1 or maxeig < -0.1:
        raise Exception(f"Either maxsig {maxeig} or minsig {mineig} are very negative.")
    mineig = max(1e-18, mineig)
    maxeig = max(1e-18, maxeig)
    return math.sqrt(maxeig), math.sqrt(mineig)


@memoize
def calculate_parameters(activation) -> (float, float, float, float):
    """
    Calculates mean and variance over a standard gaussian,
    as well as the lipschitz constant of the activation.
    """
    if isinstance(activation, str):
        activation = getattr(nn, activation)()
    gaussian = lambda x: np.exp(- x**2 / 2.0) / np.sqrt(2 * np.pi)
    npactiv = lambda x: np.array(activation(torch.tensor(x)))
    actmean = integrate.quad(lambda x: gaussian(x) * npactiv(x), -np.inf, np.inf)[0]
    actvar = integrate.quad(lambda x: gaussian(x) * (npactiv(x) - actmean)**2, -np.inf, np.inf)[0]
    # Find the maximum of the differential
    actlipsol = minimize(lambda x: -derivative(npactiv, x), np.array(0.0))
    actlip = derivative(npactiv, actlipsol.x).item()
    # Finds the smoothness constant in mean by computing sqrt(E_{x ~ N(0, 1)} phi'(x)^2)
    actmeansmoothness = math.sqrt(integrate.quad(lambda x: gaussian(x) * derivative(npactiv, x)**2, -np.inf, np.inf)[0])
    # print(f"Activation {activation} has mean {actmean}, variance {actvar}, lipschitz {actlip}, mean smoothness {actmeansmoothness}")
    # print(f"Expected smoothness: {actmeansmoothness / math.sqrt(actvar)}")
    return actmean, actvar, actlip, actmeansmoothness


def calculate_hat_sigma(activation, amount: float) -> float:
    if isinstance(activation, str):
        activation = getattr(nn, activation)()
    actmean, actvar, _, _ = calculate_parameters(activation)
    npactiv = lambda x: (np.array(activation(torch.tensor(x))) - actmean) / np.sqrt(actvar)
    def gaussian_pdf(variables, mu, invsigma):
        dimensions = variables.size
        normalization = np.sqrt((1.0 / np.linalg.det(invsigma)) * np.power(2.0 * np.pi, dimensions))
        exponential = np.exp(- 0.5 * (variables - mu).transpose() @ invsigma @ (variables - mu))
        return exponential / normalization
    # Define the function hat(sigma):
    # hat(sigma)(rho) = \int_{x, y} gaussian_2dim(x, y) * sigma(x) * sigma(y)
    def hatsigma(rho):
        invsigma = np.linalg.inv([[1.0, rho], [rho, 1.0]])
        mean = np.array([0, 0])
        return integrate.nquad(
            lambda x, y: gaussian_pdf(np.array([x, y]), mean, invsigma) * npactiv(x) * npactiv(y),
            [[-np.inf, np.inf], [-np.inf, np.inf]],
        )[0]
    return hatsigma(amount)

def calculate_expected_eigenvalue(activation, delta: float) -> float:
    """
    Calculates hat(sigma)(1.0) - hat(sigma)(1.0 - delta),
    which is the amount that should be gained in a single layer.
    """
    return 1.0 - calculate_hat_sigma(activation, 1.0 - delta)
