import torch
import pytest

from pina.model.block import LowRankBlock
from pina import LabelTensor


input_dimensions=2
embedding_dimenion=1
rank=4
inner_size=20
n_layers=2
func=torch.nn.Tanh
bias=True

def test_constructor():
    LowRankBlock(input_dimensions=input_dimensions,
                 embedding_dimenion=embedding_dimenion,
                 rank=rank,
                 inner_size=inner_size,
                 n_layers=n_layers,
                 func=func,
                 bias=bias)
    
def test_constructor_wrong():
    with pytest.raises(ValueError):
        LowRankBlock(input_dimensions=input_dimensions,
                 embedding_dimenion=embedding_dimenion,
                 rank=0.5,
                 inner_size=inner_size,
                 n_layers=n_layers,
                 func=func,
                 bias=bias)
           
def test_forward():
    block = LowRankBlock(input_dimensions=input_dimensions,
                 embedding_dimenion=embedding_dimenion,
                 rank=rank,
                 inner_size=inner_size,
                 n_layers=n_layers,
                 func=func,
                 bias=bias)
    data = LabelTensor(torch.rand(10, 30, 3), labels=['x', 'y', 'u'])
    block(data.extract('u'), data.extract(['x', 'y']))

def test_backward():
    block = LowRankBlock(input_dimensions=input_dimensions,
                 embedding_dimenion=embedding_dimenion,
                 rank=rank,
                 inner_size=inner_size,
                 n_layers=n_layers,
                 func=func,
                 bias=bias)
    data = LabelTensor(torch.rand(10, 30, 3), labels=['x', 'y', 'u'])
    data.requires_grad_(True)
    out = block(data.extract('u'), data.extract(['x', 'y']))
    loss = out.mean()
    loss.backward()