import pytest
import torch

import sys, os

# expecting to find lietorch package in parent directory
sys.path.append("..")
import lietorch

RNG_SEED = 0

if RNG_SEED:
    torch.manual_seed(RNG_SEED)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False


def print_graph(g, level=0):
    if g == None:
        return
    print(" |--" * level, g)
    for subg in g.next_functions:
        print_graph(subg[0], level + 1)


def test_adp_cpu():
    from lietorch.nn.m2 import anisotropic_dilated_project_m2

    input = torch.randn(2, 2, 8, 10, 10)
    output = anisotropic_dilated_project_m2(input, 3, 2.5, 2.0 / 3.0)

    assert list(output.shape) == [2, 2, 10, 10]


def test_adp_cuda():
    from lietorch.nn.m2 import anisotropic_dilated_project_m2

    input = torch.randn(2, 2, 8, 10, 10).cuda()
    output = anisotropic_dilated_project_m2(input, 3, 2.5, 2.0 / 3.0)

    assert list(output.shape) == [2, 2, 10, 10]


def test_adp_compare():
    from lietorch.nn.m2 import anisotropic_dilated_project_m2

    input = torch.randn(2, 2, 8, 10, 10)
    output = anisotropic_dilated_project_m2(input, 3, 2.5, 2.0 / 3.0)

    input_cuda = input.cuda()
    output_cuda = anisotropic_dilated_project_m2(input, 3, 2.5, 2.0 / 3.0)

    assert output.allclose(output_cuda.cpu())


def test_adp_autograd_cpu():
    from lietorch.nn import AnisotropicDilatedProjectM2
    from torch.autograd.gradcheck import gradcheck

    input = (torch.randn(2, 2, 8, 10, 10, dtype=torch.float64, requires_grad=True),)

    assert gradcheck(AnisotropicDilatedProjectM2(), input)


def test_adp_autograd_cuda():
    from lietorch.nn import AnisotropicDilatedProjectM2
    from torch.autograd.gradcheck import gradcheck

    input = (
        torch.randn(2, 2, 8, 10, 10, dtype=torch.float64, requires_grad=True).cuda(),
    )

    assert gradcheck(AnisotropicDilatedProjectM2(), input)


def test_lift_m2_cartesian_autograd():
    from lietorch.nn import LiftM2Cartesian
    from lietorch.nn.functional import lift_m2_cartesian
    from torch.autograd.gradcheck import gradcheck

    input = (torch.randn(1, 2, 10, 10, dtype=torch.float64, requires_grad=True),)
    module = LiftM2Cartesian(
        in_channels=2, out_channels=3, orientations=8, kernel_size=5
    ).double()

    assert gradcheck(module, input)

    input_cuda = (input[0].cuda(),)
    module_cuda = module.cuda()

    assert gradcheck(module_cuda, input_cuda, atol=1e-3)


def test_convection_m2_autograd_cpu():
    from lietorch.nn import ConvectionM2
    from lietorch.nn.functional import convection_m2
    from torch.autograd.gradcheck import gradcheck

    input = (torch.randn(1, 2, 8, 5, 5, dtype=torch.float64, requires_grad=True),)
    module = ConvectionM2(channels=2, parameter_dtype=torch.float64)

    g1 = torch.tensor(
        [[0, 0, 1.0], [0.5, 0, 1.0]], dtype=torch.float64, requires_grad=True
    )
    module.g0 = torch.nn.Parameter(g1)
    assert gradcheck(module, input, atol=1e-3)

    g2 = torch.tensor(
        [[0, 1.0, 0], [0.5, 1.0, 0.3]], dtype=torch.float64, requires_grad=True
    )
    module.g0 = torch.nn.Parameter(g2)
    assert gradcheck(module, input, atol=1e-3)

    g3 = torch.tensor(
        [[1.0, 0, 0], [1.1, -0.3, 2.1]], dtype=torch.float64, requires_grad=True
    )
    module.g0 = torch.nn.Parameter(g3)
    assert gradcheck(module, input, atol=1e-3)

    g4 = torch.tensor(
        [[0, 1.0, 1.0], [0.5, 0.00001, 1.0]], dtype=torch.float64, requires_grad=True
    )
    module.g0 = torch.nn.Parameter(g4)
    assert gradcheck(module, input, atol=1e-3)

    g5 = torch.tensor(
        [[10.0, 0, 1.0], [-1.4, 0, -5.0]], dtype=torch.float64, requires_grad=True
    )
    module.g0 = torch.nn.Parameter(g5)
    assert gradcheck(module, input, atol=1e-3)

    input2 = (
        torch.randn(1, 1, 4, 2, 2, dtype=torch.float64, requires_grad=True),
        torch.tensor([[0.1, 0.1, 0.1]], dtype=torch.float64, requires_grad=True),
    )
    assert gradcheck(convection_m2, input2, atol=1e-3)

    input3 = (
        torch.randn(1, 3, 8, 5, 5, dtype=torch.float64, requires_grad=True),
        torch.tensor(
            [[0.2, 0.5, 0.5], [-1.9, 1.2, 1.2], [10.1, -1.1, -1.2]],
            dtype=torch.float64,
            requires_grad=True,
        ),
    )
    assert gradcheck(convection_m2, input3, atol=1e-3)


def test_convection_m2_autograd_cuda():
    from lietorch.nn import ConvectionM2
    from lietorch.nn.functional import convection_m2
    from torch.autograd.gradcheck import gradcheck

    input = (
        torch.randn(1, 2, 8, 5, 5, dtype=torch.float64, requires_grad=True).cuda(),
    )
    module = ConvectionM2(channels=2, parameter_dtype=torch.float64)

    g1 = torch.tensor(
        [[0, 0, 1.0], [0.5, 0, 1.0]], dtype=torch.float64, requires_grad=True
    ).cuda()
    module.g0 = torch.nn.Parameter(g1)
    assert gradcheck(module, input, atol=1e-3)

    g2 = torch.tensor(
        [[0, 1.0, 0], [0.5, 1.0, 0.3]], dtype=torch.float64, requires_grad=True
    ).cuda()
    module.g0 = torch.nn.Parameter(g2)
    assert gradcheck(module, input, atol=1e-3)

    g3 = torch.tensor(
        [[1.0, 0, 0], [1.1, -0.3, 2.1]], dtype=torch.float64, requires_grad=True
    ).cuda()
    module.g0 = torch.nn.Parameter(g3)
    assert gradcheck(module, input, atol=1e-3)

    g4 = torch.tensor(
        [[0, 1.0, 1.0], [0.5, 0.00001, 1.0]], dtype=torch.float64, requires_grad=True
    ).cuda()
    module.g0 = torch.nn.Parameter(g4)
    assert gradcheck(module, input, atol=1e-3)

    g5 = torch.tensor(
        [[10.0, 0, 1.0], [-1.4, 0, -5.0]], dtype=torch.float64, requires_grad=True
    ).cuda()
    module.g0 = torch.nn.Parameter(g5)
    assert gradcheck(module, input, atol=1e-3)

    input2 = (
        torch.randn(1, 1, 4, 2, 2, dtype=torch.float64, requires_grad=True).cuda(),
        torch.tensor([[0.1, 0.1, 0.1]], dtype=torch.float64, requires_grad=True).cuda(),
    )
    assert gradcheck(convection_m2, input2, atol=1e-3)

    input3 = (
        torch.randn(1, 3, 8, 5, 5, dtype=torch.float64, requires_grad=True).cuda(),
        torch.tensor(
            [[0.2, 0.5, 0.5], [-1.9, 1.2, 1.2], [10.1, -1.1, -1.2]],
            dtype=torch.float64,
            requires_grad=True,
        ).cuda(),
    )
    assert gradcheck(convection_m2, input3, atol=1e-3)


def test_m2_lift_shape():
    x = torch.randn(2, 3, 100, 200)
    module = lietorch.nn.m2.LiftM2Cartesian(
        3, 4, orientations=8, kernel_size=5, spline_order=2
    )
    os1 = module(x)
    os2 = lietorch.nn.m2.lift_m2_cartesian(
        x, module.weights, orientations=8, spline_order=2
    )

    assert os1.shape == os2.shape == (2, 4, 8, 100 - 4, 200 - 4)
    assert (os1 - os2).abs().max() < 1e-8


def test_maxproject():
    # test forward
    a = torch.tensor([[[1, 2], [3, 4]]]).float().requires_grad_()
    expected = torch.tensor([[2, 4]]).float()

    compare = lietorch.nn.m2.max_project_m2(a) == expected
    assert compare.all().item()

    # test backward
    torch.autograd.gradcheck(lietorch.nn.m2.max_project_m2, a.double())

    # test modular interface
    module = lietorch.nn.m2.MaxProjectM2()
    compare = module(a) == expected
    assert compare.all().item()

    # test backward
    torch.autograd.gradcheck(module, a.double())


def test_reflectionpadding():
    a = torch.tensor([[[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]]])
    b = torch.tensor(
        [
            [
                [
                    [
                        [4.0, 3.0, 3.0, 4.0, 4.0, 3.0],
                        [2.0, 1.0, 1.0, 2.0, 2.0, 1.0],
                        [2.0, 1.0, 1.0, 2.0, 2.0, 1.0],
                        [4.0, 3.0, 3.0, 4.0, 4.0, 3.0],
                        [4.0, 3.0, 3.0, 4.0, 4.0, 3.0],
                        [2.0, 1.0, 1.0, 2.0, 2.0, 1.0],
                    ],
                    [
                        [8.0, 7.0, 7.0, 8.0, 8.0, 7.0],
                        [6.0, 5.0, 5.0, 6.0, 6.0, 5.0],
                        [6.0, 5.0, 5.0, 6.0, 6.0, 5.0],
                        [8.0, 7.0, 7.0, 8.0, 8.0, 7.0],
                        [8.0, 7.0, 7.0, 8.0, 8.0, 7.0],
                        [6.0, 5.0, 5.0, 6.0, 6.0, 5.0],
                    ],
                ]
            ]
        ]
    )

    assert (lietorch.nn.m2.reflection_pad_m2(a, 2) == b).all().item()


def test_morphological_convolution_m2_cpu_autograd():
    from lietorch.nn import MorphologicalConvolutionM2
    from lietorch.nn.functional import morphological_convolution_m2
    from torch.autograd.gradcheck import gradcheck

    input = (
        torch.randn(2, 2, 8, 7, 7, requires_grad=True, dtype=torch.float64),
        torch.randn(2, 3, 3, 3, requires_grad=True, dtype=torch.float64),
    )

    # input = (
    #     torch.randperm(8 * 7 * 7, dtype=torch.float64, requires_grad=True).view(
    #         1, 1, 8, 7, 7
    #     ),
    #     torch.randn(3 * 3 * 3, dtype=torch.float64, requires_grad=True).view(
    #         1, 3, 3, 3
    #     ),
    # )
    assert gradcheck(morphological_convolution_m2, input, atol=1e-4)

    module = MorphologicalConvolutionM2(channels=2, kernel_size=(2, 2, 2))

    assert gradcheck(
        module,
        torch.randn(1, 2, 8, 9, 9, requires_grad=True, dtype=torch.float64),
        atol=1e-4,
    )


def test_morphological_convolution_m2_cuda_autograd():
    from lietorch.nn import MorphologicalConvolutionM2
    from lietorch.nn.functional import morphological_convolution_m2
    from torch.autograd.gradcheck import gradcheck

    input = (
        torch.randn(2, 2, 8, 7, 7, requires_grad=True, dtype=torch.float64).cuda(),
        torch.randn(2, 3, 3, 3, requires_grad=True, dtype=torch.float64).cuda(),
    )

    assert gradcheck(morphological_convolution_m2, input, atol=1e-4)

    module = MorphologicalConvolutionM2(channels=2, kernel_size=(2, 2, 2))

    assert gradcheck(
        module.cuda(),
        torch.randn(1, 2, 8, 9, 9, requires_grad=True, dtype=torch.float64).cuda(),
        atol=1e-4,
    )


def test_logarithmic_metric_estimate():
    import torch
    from torch.autograd.gradcheck import gradcheck

    rho = torch.ops.lietorch.m2_logarithmic_metric_estimate

    params = torch.rand([2, 4, 3], dtype=torch.float64, requires_grad=True)

    assert gradcheck(rho, (params, [5, 5, 5], 8), atol=1e-4)
    assert gradcheck(rho, (params.cuda(), [5, 5, 5], 8), atol=1e-4)

    assert gradcheck(rho, (params, [4, 4, 4], 8), atol=1e-4)
    assert gradcheck(rho, (params.cuda(), [4, 4, 4], 8), atol=1e-4)


def test_morphological_kernel():
    import torch
    from torch.autograd.gradcheck import gradcheck

    kf = torch.ops.lietorch.m2_morphological_kernel

    params = torch.rand([2, 4, 3], dtype=torch.float64, requires_grad=True)

    k = kf(params, [5, 5, 5], 8, 0.55)
    k = kf(params, [5, 5, 5], 8, 1.0)

    assert gradcheck(kf, (params, [5, 5, 5], 8, 0.55), atol=1e-4)
    assert gradcheck(kf, (params, [5, 5, 5], 8, 1.0), atol=1e-4)


def test_diffusion_kernel():
    import torch
    from torch.autograd.gradcheck import gradcheck

    kf = torch.ops.lietorch.m2_diffusion_kernel

    params = torch.rand([2, 4, 3], dtype=torch.float64, requires_grad=True)

    k = kf(
        torch.tensor(
            [[1, 1, 0.5], [0.5, 2.0, 0.1]], dtype=torch.float64, requires_grad=True
        ),
        [5, 5, 5],
        8,
    )

    assert (k.sum([-3, -2, -1]) - 1).abs().max() < 1e-8

    assert gradcheck(kf, (params, [5, 5, 5], 8))
    assert gradcheck(kf, (params, [6, 6, 6], 7))


def test_fractional_dilation():
    import torch
    from torch.autograd.gradcheck import gradcheck
    from lietorch.nn.functional import fractional_dilation_m2
    from lietorch.nn import FractionalDilationM2

    x = torch.randn([2, 2, 8, 5, 5], dtype=torch.float64, requires_grad=True)
    p = torch.rand([2, 3], dtype=torch.float64, requires_grad=True)

    assert gradcheck(fractional_dilation_m2, (x, p, [3, 3, 3], 0.55), atol=1e-4)
    assert gradcheck(fractional_dilation_m2, (x, p, [4, 4, 4], 1.0), atol=1e-4)

    xc = x.cuda()
    pc = p.cuda()

    assert gradcheck(fractional_dilation_m2, (xc, pc, [3, 3, 3], 0.55), atol=1e-4)
    assert gradcheck(fractional_dilation_m2, (xc, pc, [4, 4, 4], 1.0), atol=1e-4)

    module = FractionalDilationM2(2, [2, 2, 2], alpha=0.7)

    assert gradcheck(module, (x,), atol=1e-4)
    assert gradcheck(module.cuda(), (x.cuda(),), atol=1e-4)


def test_fractional_erosion():
    import torch
    from torch.autograd.gradcheck import gradcheck
    from lietorch.nn.functional import fractional_erosion_m2
    from lietorch.nn import FractionalErosionM2

    x = torch.randn([2, 2, 8, 5, 5], dtype=torch.float64, requires_grad=True)
    p = torch.rand([2, 3], dtype=torch.float64, requires_grad=True)

    assert gradcheck(fractional_erosion_m2, (x, p, [3, 3, 3], 0.55), atol=1e-4)
    assert gradcheck(fractional_erosion_m2, (x, p, [4, 4, 4], 1.0), atol=1e-4)

    xc = x.cuda()
    pc = p.cuda()

    assert gradcheck(fractional_erosion_m2, (xc, pc, [3, 3, 3], 0.55), atol=1e-4)
    assert gradcheck(fractional_erosion_m2, (xc, pc, [4, 4, 4], 1.0), atol=1e-4)

    module = FractionalErosionM2(2, [2, 2, 2], alpha=0.7)

    assert gradcheck(module, (x,), atol=1e-4)
    assert gradcheck(module.cuda(), (x.cuda(),), atol=1e-4)


def test_linear():
    import torch
    from torch.autograd.gradcheck import gradcheck
    from lietorch.nn.functional import linear_m2
    from lietorch.nn import LinearM2

    a = torch.tensor([[[[[1, 1], [1, 1]]]]], dtype=torch.float64)
    v = torch.tensor([[2, 3]], dtype=torch.float64)
    b = linear_m2(a, v)
    b_exp = torch.tensor(
        [[[[[2, 2], [2, 2]]], [[[3, 3], [3, 3]]]]], dtype=torch.float64
    )
    assert b.allclose(b_exp)
    assert linear_m2(a.cuda(), v.cuda()).allclose(b_exp.cuda())

    x = torch.randn([2, 3, 8, 5, 5], dtype=torch.float64, requires_grad=True)
    w = torch.randn([3, 6], dtype=torch.float64, requires_grad=True)

    assert gradcheck(linear_m2, (x, w), atol=1e-4)
    assert gradcheck(linear_m2, (x.cuda(), w.cuda()), atol=1e-4)

    module = LinearM2(3, 5).double()

    assert gradcheck(module, (x,), atol=1e-4)


def test_convection_dilation_pde():
    from torch.autograd.gradcheck import gradcheck
    from lietorch.nn import ConvectionDilationPdeM2

    x = torch.randn([1, 2, 8, 7, 7], dtype=torch.float64, requires_grad=True)
    module = ConvectionDilationPdeM2(channels=2, kernel_size=[5, 5, 5]).double()
    module2 = ConvectionDilationPdeM2(
        channels=2, kernel_size=[4, 4, 4], alpha_dilation=1.0, iterations=2,
    ).double()

    assert gradcheck(module, (x,), atol=1e-4)
    assert gradcheck(module2, (x,), atol=1e-4)
    assert gradcheck(module.cuda(), (x.cuda(),), atol=1e-4)
    assert gradcheck(module2.cuda(), (x.cuda(),), atol=1e-4)


def test_convection_erosion_pde():
    from torch.autograd.gradcheck import gradcheck
    from lietorch.nn import ConvectionErosionPdeM2

    x = torch.randn([1, 2, 8, 7, 7], dtype=torch.float64, requires_grad=True)
    module = ConvectionErosionPdeM2(channels=2, kernel_size=[5, 5, 5]).double()
    module2 = ConvectionErosionPdeM2(
        channels=2, kernel_size=[4, 4, 4], alpha_erosion=1.0, iterations=2,
    ).double()

    assert gradcheck(module, (x,), atol=1e-4)
    assert gradcheck(module2, (x,), atol=1e-4)
    assert gradcheck(module.cuda(), (x.cuda(),), atol=1e-4)
    assert gradcheck(module2.cuda(), (x.cuda(),), atol=1e-4)


def test_cde_pde_layer():
    from torch.autograd.gradcheck import gradcheck
    from lietorch.nn import CDEPdeLayerM2

    x = torch.randn([1, 2, 4, 5, 5], dtype=torch.float64, requires_grad=True)
    module = CDEPdeLayerM2(
        in_channels=2, out_channels=3, kernel_size=[3, 3, 3]
    ).double()

    assert gradcheck(module, (x,), atol=1e-4)
    assert gradcheck(module.cuda(), (x.cuda(),), atol=1e-4, nondet_tol=1e-4)

