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


batch_size = 15
n_layers = 4
embedding_dim = 24
func = torch.nn.Tanh
rank = 4
n_kernel_layers = 3
field_indices = ["u"]
coordinates_indices = ["x", "y"]


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(coordinates_indices), len(field_indices)
    )
    LowRankNeuralOperator(
        lifting_net=lifting_net,
        projecting_net=projecting_net,
        coordinates_indices=coordinates_indices,
        field_indices=field_indices,
        n_kernel_layers=n_kernel_layers,
        rank=rank,
    )

    # not working constructor
    with pytest.raises(ValueError):
        LowRankNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=projecting_net,
            coordinates_indices=coordinates_indices,
            field_indices=field_indices,
            n_kernel_layers=3.2,  # wrong
            rank=rank,
        )

        LowRankNeuralOperator(
            lifting_net=[0],  # wrong
            projecting_net=projecting_net,
            coordinates_indices=coordinates_indices,
            field_indices=field_indices,
            n_kernel_layers=n_kernel_layers,
            rank=rank,
        )

        LowRankNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=[0],  # wront
            coordinates_indices=coordinates_indices,
            field_indices=field_indices,
            n_kernel_layers=n_kernel_layers,
            rank=rank,
        )

        LowRankNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=projecting_net,
            coordinates_indices=[0],  # wrong
            field_indices=field_indices,
            n_kernel_layers=n_kernel_layers,
            rank=rank,
        )

        LowRankNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=projecting_net,
            coordinates_indices=coordinates_indices,
            field_indices=[0],  # wrong
            n_kernel_layers=n_kernel_layers,
            rank=rank,
        )

        lifting_net = torch.nn.Linear(len(coordinates_indices), embedding_dim)
        LowRankNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=projecting_net,
            coordinates_indices=coordinates_indices,
            field_indices=field_indices,
            n_kernel_layers=n_kernel_layers,
            rank=rank,
        )

        lifting_net = torch.nn.Linear(
            len(coordinates_indices) + len(field_indices), embedding_dim
        )
        projecting_net = torch.nn.Linear(embedding_dim, len(field_indices))
        LowRankNeuralOperator(
            lifting_net=lifting_net,
            projecting_net=projecting_net,
            coordinates_indices=coordinates_indices,
            field_indices=field_indices,
            n_kernel_layers=n_kernel_layers,
            rank=rank,
        )


def test_forward():
    lifting_net = torch.nn.Linear(
        len(coordinates_indices) + len(field_indices), embedding_dim
    )
    projecting_net = torch.nn.Linear(
        embedding_dim + len(coordinates_indices), len(field_indices)
    )
    lno = LowRankNeuralOperator(
        lifting_net=lifting_net,
        projecting_net=projecting_net,
        coordinates_indices=coordinates_indices,
        field_indices=field_indices,
        n_kernel_layers=n_kernel_layers,
        rank=rank,
    )

    input_ = LabelTensor(
        torch.rand(
            batch_size, 100, len(coordinates_indices) + len(field_indices)
        ),
        coordinates_indices + field_indices,
    )

    out = lno(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(coordinates_indices), len(field_indices)
    )
    lno = LowRankNeuralOperator(
        lifting_net=lifting_net,
        projecting_net=projecting_net,
        coordinates_indices=coordinates_indices,
        field_indices=field_indices,
        n_kernel_layers=n_kernel_layers,
        rank=rank,
    )
    input_ = LabelTensor(
        torch.rand(
            batch_size, 100, len(coordinates_indices) + len(field_indices)
        ),
        coordinates_indices + field_indices,
    )
    input_ = input_.requires_grad_()
    out = lno(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),
        ]
    )