import torch
from pina.model import KernelNeuralOperator, FeedForward

input_dim = 2
output_dim = 4
embedding_dim = 24
batch_size = 10
numb = 256
data = torch.rand(size=(batch_size, numb, input_dim), requires_grad=True)
output_shape = torch.Size([batch_size, numb, output_dim])


lifting_operator = FeedForward(
    input_dimensions=input_dim, output_dimensions=embedding_dim
)
projection_operator = FeedForward(
    input_dimensions=embedding_dim, output_dimensions=output_dim
)
integral_kernels = torch.nn.Sequential(
    FeedForward(
        input_dimensions=embedding_dim, output_dimensions=embedding_dim
    ),
    FeedForward(
        input_dimensions=embedding_dim, output_dimensions=embedding_dim
    ),
)


def test_constructor():
    KernelNeuralOperator(
        lifting_operator=lifting_operator,
        integral_kernels=integral_kernels,
        projection_operator=projection_operator,
    )


def test_forward():
    operator = KernelNeuralOperator(
        lifting_operator=lifting_operator,
        integral_kernels=integral_kernels,
        projection_operator=projection_operator,
    )
    out = operator(data)
    assert out.shape == output_shape


def test_backward():
    operator = KernelNeuralOperator(
        lifting_operator=lifting_operator,
        integral_kernels=integral_kernels,
        projection_operator=projection_operator,
    )
    out = operator(data)
    loss = torch.nn.functional.mse_loss(out, torch.zeros_like(out))
    loss.backward()
    grad = data.grad
    assert grad.shape == data.shape