import torch
from pina.model import AveragingNeuralOperator
from pina import LabelTensor
import pytest


batch_size = 15
n_layers = 4
embedding_dim = 24
func = torch.nn.Tanh
coordinates_indices = ["p"]
field_indices = ["v"]


def test_constructor():
    # working constructor
    lifting_net = torch.nn.Linear(
        len(coordinates_indices) + len(field_indices), embedding_dim
    )
    projecting_net = torch.nn.Linear(
        embedding_dim + len(field_indices), len(field_indices)
    )
    AveragingNeuralOperator(
        lifting_net=lifting_net,
        projecting_net=projecting_net,
        coordinates_indices=coordinates_indices,
        field_indices=field_indices,
        n_layers=n_layers,
        func=func,
    )

    # not working constructor
    with pytest.raises(ValueError):
        AveragingNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=projecting_net,
            coordinates_indices=coordinates_indices,
            field_indices=field_indices,
            n_layers=3.2,  # wrong
            func=func,
        )

        AveragingNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=projecting_net,
            coordinates_indices=coordinates_indices,
            field_indices=field_indices,
            n_layers=n_layers,
            func=1,
        )  # wrong

        AveragingNeuralOperator(
            lifting_net=[0],  # wrong
            projecting_net=projecting_net,
            coordinates_indices=coordinates_indices,
            field_indices=field_indices,
            n_layers=n_layers,
            func=func,
        )

        AveragingNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=[0],  # wront
            coordinates_indices=coordinates_indices,
            field_indices=field_indices,
            n_layers=n_layers,
            func=func,
        )

        AveragingNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=projecting_net,
            coordinates_indices=[0],  # wrong
            field_indices=field_indices,
            n_layers=n_layers,
            func=func,
        )

        AveragingNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=projecting_net,
            coordinates_indices=coordinates_indices,
            field_indices=[0],  # wrong
            n_layers=n_layers,
            func=func,
        )

        lifting_net = torch.nn.Linear(len(coordinates_indices), embedding_dim)
        AveragingNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=projecting_net,
            coordinates_indices=coordinates_indices,
            field_indices=field_indices,
            n_layers=n_layers,
            func=func,
        )

        lifting_net = torch.nn.Linear(
            len(coordinates_indices) + len(field_indices), embedding_dim
        )
        projecting_net = torch.nn.Linear(embedding_dim, len(field_indices))
        AveragingNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=projecting_net,
            coordinates_indices=coordinates_indices,
            field_indices=field_indices,
            n_layers=n_layers,
            func=func,
        )


def test_forward():
    lifting_net = torch.nn.Linear(
        len(coordinates_indices) + len(field_indices), embedding_dim
    )
    projecting_net = torch.nn.Linear(
        embedding_dim + len(field_indices), len(field_indices)
    )
    avno = AveragingNeuralOperator(
        lifting_net=lifting_net,
        projecting_net=projecting_net,
        coordinates_indices=coordinates_indices,
        field_indices=field_indices,
        n_layers=n_layers,
        func=func,
    )

    input_ = LabelTensor(
        torch.rand(
            batch_size, 100, len(coordinates_indices) + len(field_indices)
        ),
        ["p", "v"],
    )

    out = avno(input_)
    assert out.shape == torch.Size(
        [batch_size, input_.shape[1], len(field_indices)]
    )


def test_backward():
    lifting_net = torch.nn.Linear(
        len(coordinates_indices) + len(field_indices), embedding_dim
    )
    projecting_net = torch.nn.Linear(
        embedding_dim + len(field_indices), len(field_indices)
    )
    avno = AveragingNeuralOperator(
        lifting_net=lifting_net,
        projecting_net=projecting_net,
        coordinates_indices=coordinates_indices,
        field_indices=field_indices,
        n_layers=n_layers,
        func=func,
    )
    input_ = LabelTensor(
        torch.rand(
            batch_size, 100, len(coordinates_indices) + len(field_indices)
        ),
        ["p", "v"],
    )
    input_ = input_.requires_grad_()
    out = avno(input_)
    tmp = torch.linalg.norm(out)
    tmp.backward()
    grad = input_.grad
    assert grad.shape == torch.Size(
        [
            batch_size,
            input_.shape[1],
            len(coordinates_indices) + len(field_indices),
        ]
    )