import torch
import pytest

from pina.model.block import PeriodicBoundaryEmbedding, FourierFeatureEmbedding

# test tolerance
tol = 1e-6

def check_same_columns(tensor):
    # Get the first column and compute residual
    residual = tensor - tensor[0]
    zeros = torch.zeros_like(residual)
    # Compare each column with the first column
    all_same = torch.allclose(input=residual,other=zeros,atol=tol)
    return all_same

def grad(u, x):
    """
    Compute the first derivative of u with respect to x.
    """
    return torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u),
                               create_graph=True, allow_unused=True,
                               retain_graph=True)[0]

def test_constructor_PeriodicBoundaryEmbedding():
    PeriodicBoundaryEmbedding(input_dimension=1, periods=2)
    PeriodicBoundaryEmbedding(input_dimension=1, periods={'x': 3, 'y' : 4})
    PeriodicBoundaryEmbedding(input_dimension=1, periods={0: 3, 1 : 4})
    PeriodicBoundaryEmbedding(input_dimension=1, periods=2, output_dimension=10)
    with pytest.raises(TypeError):
        PeriodicBoundaryEmbedding()
    with pytest.raises(ValueError):
        PeriodicBoundaryEmbedding(input_dimension=1., periods=1)
        PeriodicBoundaryEmbedding(input_dimension=1, periods=1,
                                  output_dimension=1.)
        PeriodicBoundaryEmbedding(input_dimension=1, periods={'x':'x'})
        PeriodicBoundaryEmbedding(input_dimension=1, periods={0:'x'})


@pytest.mark.parametrize("period", [1, 4, 10])
@pytest.mark.parametrize("input_dimension", [1, 2, 3])
def test_forward_backward_same_period_PeriodicBoundaryEmbedding(input_dimension,
                                                                period):
    func = torch.nn.Sequential(
        PeriodicBoundaryEmbedding(input_dimension=input_dimension,
                     output_dimension=60, periods=period),
        torch.nn.Tanh(),
        torch.nn.Linear(60, 60),
        torch.nn.Tanh(),
        torch.nn.Linear(60, 1)
    )
    # coordinates
    x = period * torch.tensor([[0.],[1.]])
    if input_dimension == 2:
        x = torch.cartesian_prod(x.flatten(),x.flatten())
    elif input_dimension == 3:
        x = torch.cartesian_prod(x.flatten(),x.flatten(),x.flatten())
    x.requires_grad = True
    # output
    f = func(x)
    assert check_same_columns(f)
    # compute backward
    loss = f.mean()
    loss.backward()

def test_constructor_FourierFeatureEmbedding():
    FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
                            sigma=1)
    with pytest.raises(TypeError): 
        FourierFeatureEmbedding()
    with pytest.raises(RuntimeError): 
        FourierFeatureEmbedding(input_dimension=1, output_dimension=3, sigma=1)
    with pytest.raises(ValueError):
        FourierFeatureEmbedding(input_dimension='x', output_dimension=20,
                                sigma=1)
        FourierFeatureEmbedding(input_dimension=1, output_dimension='x',
                                sigma=1)
        FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
                                sigma='x')

@pytest.mark.parametrize("output_dimension", [2, 4, 6])
@pytest.mark.parametrize("input_dimension", [1, 2, 3])
@pytest.mark.parametrize("sigma", [10, 1, 0.1])
def test_forward_backward_FourierFeatureEmbedding(input_dimension,
                                                  output_dimension,
                                                  sigma):
    func = FourierFeatureEmbedding(input_dimension, output_dimension,
                                   sigma)
    # coordinates
    x = torch.rand((10, input_dimension), requires_grad=True)
    # output
    f = func(x)
    assert f.shape[-1] == output_dimension
    # compute backward
    loss = f.mean()
    loss.backward()