
import unittest
from hypothesis import given, example, settings, assume
from hypothesis.strategies import integers, one_of, just, booleans
import numpy as np
import torch
import math
import torch.nn as nn
from linear_operator.operators.diag_linear_operator import DiagLinearOperator
from linear_operator.settings import max_cg_iterations, cg_tolerance
from utils import (
    BatchDenseLinearOperator, activation_derivative,
    minimum_norm_solution, least_squares_solution,
    extremal_sigs, extremal_sigs_noq, triangular_solve, qr_decomposition,
)


class TensorError(Exception):
    def __init__(self, msg, **extra):
        self.msg = msg
        self.extra = extra

    def __str__(self):
        out = self.msg + "\n"
        for key, value in self.extra.items():
            out += f"\t{key}: {value}\n\n"
        return out


def assertTensorsAlmostEqual(self, t1, t2, allowed_diff=1e-6):
    if isinstance(t1, float) and isinstance(t2, float):
        maxdiff = abs(t1 - t2)
    else:
        maxdiff = (t1 - t2).max().item()
    if maxdiff >= allowed_diff:
        raise TensorError(f"Tensors are {maxdiff} unequal.", t1=t1, t2=t2)

unittest.TestCase.assertTensorsAlmostEqual = assertTensorsAlmostEqual


def assertTensorsRelativelyEqual(self, t1, t2, allowed_rel=1e-2):
    if isinstance(t1, float) and isinstance(t2, float):
        rel = max(t1 / t2, t2 / t1)
    else:
        rel = max((t1 / t2).max().item(), (t2 / t1).max().item())
    if rel >= 1.0 + allowed_rel:
        raise TensorError(f"Tensors are {rel} relatively unequal.", t1=t1, t2=t2)

unittest.TestCase.assertTensorsRelativelyEqual = assertTensorsRelativelyEqual


class TestActivationDerivatives(unittest.TestCase):
    @settings(max_examples=15)
    @given(actname=one_of(just("Identity"), just("ReLU"), just("Tanh")),
           seed=integers(min_value=0, max_value=1000))
    def test_activation_derivative(self, actname: str, seed: int):
        torch.manual_seed(seed)
        N = 1000
        inputs = torch.randn(N)
        computed_grads = activation_derivative(actname, inputs)
        activation = getattr(nn, actname)()
        real_grads = torch.autograd.functional.jvp(activation, inputs, v=torch.ones(inputs.size()))[1]
        self.assertTensorsAlmostEqual(real_grads, computed_grads, allowed_diff=1e-3)


class TestLinearSystemsSolution(unittest.TestCase):
    @settings(deadline=None)
    @given(in_dim=integers(min_value=1, max_value=10),
           out_dim=integers(min_value=1, max_value=10),
           batches=integers(min_value=1, max_value=3),
           seed=integers(min_value=0, max_value=1000))
    def test_minimum_norm_solution(self, batches: int, in_dim: int, out_dim: int, seed:int):
        assume(out_dim <= in_dim)
        torch.manual_seed(seed)
        A = torch.randn((batches, out_dim, in_dim))
        b = torch.randn((batches, out_dim, 1))
        Q, R = qr_decomposition(A.permute(0, 2, 1))
        oursol = minimum_norm_solution(b, R, Q)
        torchsol = torch.zeros((batches, in_dim, 1))
        for i in range(batches):
            torchsol[i, :, :], _ = torch.lstsq(b[i, :, :], A[i, :, :])
        assert torchsol.size() == oursol.size()
        self.assertTensorsRelativelyEqual(torchsol, oursol, allowed_rel=1e-1)

    @settings(deadline=None)
    @given(in_dim=integers(min_value=1, max_value=10),
           out_dim=integers(min_value=1, max_value=10),
           batches=integers(min_value=1, max_value=3),
           seed=integers(min_value=0, max_value=1000))
    def test_least_squares_solution(self, batches: int, in_dim: int, out_dim: int, seed:int):
        assume(out_dim >= in_dim)
        torch.manual_seed(seed)
        A = torch.randn((batches, out_dim, in_dim))
        b = torch.randn((batches, out_dim, 1))
        Q, R = qr_decomposition(A)
        oursol = least_squares_solution(b, R, Q)
        torchsol = torch.zeros((batches, in_dim, 1))
        for i in range(batches):
            torchsol[i, :, :] = torch.lstsq(b[i, :, :], A[i, :, :])[0][:in_dim, :]
        self.assertTensorsRelativelyEqual(torchsol, oursol, allowed_rel=1e-1)

    @settings(deadline=None)
    @given(size=integers(min_value=1, max_value=30),
           batches=integers(min_value=1, max_value=10),
           postbatches=integers(min_value=1, max_value=10),
           upper=booleans(),
           seed=integers(min_value=0, max_value=1000))
    def test_triangular_solve(self, batches: int, size: int, postbatches: int, seed: int, upper: bool):
        torch.manual_seed(seed)
        A = torch.randn((batches, size, size))
        A = A.triu() if upper == True else A.tril()
        b = torch.randn((batches, size, postbatches))
        oursol = triangular_solve(b, A, upper=upper)
        torchsol, _ = torch.triangular_solve(b, A, upper=upper)
        assert torchsol.size() == oursol.size()
        self.assertTensorsRelativelyEqual(torchsol, oursol, allowed_rel=1e-1)

    @given(outdim=integers(min_value=1, max_value=30),
           indim=integers(min_value=1, max_value=30),
           batches=integers(min_value=1, max_value=10),
           seed=integers(min_value=0, max_value=1000))
    def test_qr_decomposition(self, batches: int, outdim: int, indim: int, seed: int):
        torch.manual_seed(seed)
        assume(outdim >= indim)
        A = torch.randn((batches, outdim, indim))
        ourQ, ourR = qr_decomposition(A)
        Id = torch.eye(outdim)[None, :]
        self.assertTensorsAlmostEqual(ourQ @ ourR, A, allowed_diff=1e-4)
        self.assertTensorsAlmostEqual(ourQ @ ourQ.permute(0, 2, 1), Id, allowed_diff=1e-4)
        self.assertTensorsAlmostEqual(ourR, ourR.triu(), allowed_diff=1e-4)

    
class TestBatchDenseLop(unittest.TestCase):
    @given(
        N=integers(min_value=1, max_value=100),
        M=integers(min_value=1, max_value=100),
        batches=integers(min_value=1, max_value=8),
        seed=integers(min_value=0, max_value=1000),
    )
    def test_multiplication_to_dense(self, N, M, batches, seed):
        torch.manual_seed(seed)
        layer = nn.Linear(N, M)
        nn.init.normal_(layer.weight)
        lop = BatchDenseLinearOperator(layer.weight, batches)
        x = torch.randn(N * batches)
        asdense = lop.to_dense() @ x
        asoper = lop @ x
        self.assertTensorsAlmostEqual(asdense, asoper, allowed_diff=1e-4)

    @settings(deadline=None, max_examples=30)
    @given(
        N=integers(min_value=1, max_value=100),
        M=integers(min_value=1, max_value=100),
        seed=integers(min_value=0, max_value=1000),
    )
    def test_extremal_sigs(self, N: int, M: int, seed: int):
        torch.manual_seed(seed)
        weight = torch.zeros((N, M))
        grads = torch.diag_embed(torch.randn(M))
        nn.init.normal_(weight)
        complete = weight @ grads
        _, eigs, _ = torch.linalg.svd(complete.unsqueeze(0), compute_uv=False)
        maxsig, minsig = extremal_sigs_noq(complete.unsqueeze(0))
        self.assertTensorsRelativelyEqual(eigs.max().item(), maxsig, allowed_rel=0.1)
        self.assertTensorsRelativelyEqual(eigs.min().item(), minsig, allowed_rel=0.1)

    @settings(deadline=None, max_examples=30)
    @given(
        N=integers(min_value=1, max_value=100),
        M=integers(min_value=1, max_value=100),
        batches=integers(min_value=1, max_value=1),
        seed=integers(min_value=0, max_value=1000),
    )
    def test_eigenvalues_batch_dense_lop_with_diagonal(self, N: int, M: int, batches: int, seed: int):
        torch.manual_seed(seed)
        assume(N >= M)
        layer = nn.Linear(M, N)
        grads = torch.randn(batches, M)
        nn.init.normal_(layer.weight)
        linearop = BatchDenseLinearOperator(layer.weight, batches)
        diagop = DiagLinearOperator(grads.view(-1))
        computed_maxsig, computed_minsig = (linearop @ diagop).extremal_sig(maxstep=100)
        _, real_sigs, _ = torch.linalg.svd((linearop @ diagop).to_dense(), compute_uv=False)
        real_maxsig = real_sigs.max().item()
        real_minsig = real_sigs.min().item()
        self.assertTensorsRelativelyEqual(real_maxsig, computed_maxsig, allowed_rel=1e-1)
        self.assertTensorsRelativelyEqual(real_minsig, computed_minsig, allowed_rel=1e-1)


if __name__ == "__main__":
    unittest.run()
