import torch

from pina import LabelTensor
from pina.domain import Union, EllipsoidDomain, CartesianDomain


def test_constructor_two_CartesianDomains():
    Union(
        [
            CartesianDomain({"x": [0, 1], "y": [0, 1]}),
            CartesianDomain({"x": [0.5, 2], "y": [-1, 0.1]}),
        ]
    )


def test_constructor_two_EllipsoidDomains():
    Union(
        [
            EllipsoidDomain({"x": [-1, 1], "y": [-1, 1], "z": [-1, 1]}),
            EllipsoidDomain(
                {"x": [-0.5, 0.5], "y": [-0.5, 0.5], "z": [-0.5, 0.5]}
            ),
        ]
    )


def test_constructor_EllipsoidDomain_CartesianDomain():
    Union(
        [
            EllipsoidDomain({"x": [-1, 1], "y": [-1, 1]}),
            CartesianDomain({"x": [-0.5, 0.5], "y": [-0.5, 0.5]}),
        ]
    )


def test_is_inside_two_CartesianDomains():
    pt_1 = LabelTensor(torch.tensor([[0.5, 0.5]]), ["x", "y"])
    pt_2 = LabelTensor(torch.tensor([[-1, -1]]), ["x", "y"])
    domain = Union(
        [
            CartesianDomain({"x": [0, 1], "y": [0, 1]}),
            CartesianDomain({"x": [0.5, 2], "y": [-1, 0.1]}),
        ]
    )
    assert domain.is_inside(pt_1) == True
    assert domain.is_inside(pt_2) == False


def test_is_inside_two_EllipsoidDomains():
    pt_1 = LabelTensor(torch.tensor([[0.5, 0.5, 0.5]]), ["x", "y", "z"])
    pt_2 = LabelTensor(torch.tensor([[-1, -1, -1]]), ["x", "y", "z"])
    domain = Union(
        [
            EllipsoidDomain({"x": [-1, 1], "y": [-1, 1], "z": [-1, 1]}),
            EllipsoidDomain(
                {"x": [-0.5, 0.5], "y": [-0.5, 0.5], "z": [-0.5, 0.5]}
            ),
        ]
    )
    assert domain.is_inside(pt_1) == True
    assert domain.is_inside(pt_2) == False


def test_is_inside_EllipsoidDomain_CartesianDomain():
    pt_1 = LabelTensor(torch.tensor([[0.5, 0.5]]), ["x", "y"])
    pt_2 = LabelTensor(torch.tensor([[-1, -1]]), ["x", "y"])
    domain = Union(
        [
            EllipsoidDomain(
                {
                    "x": [-1, 1],
                    "y": [-1, 1],
                }
            ),
            CartesianDomain({"x": [0.6, 1.5], "y": [-2, 0]}),
        ]
    )
    assert domain.is_inside(pt_1) == True
    assert domain.is_inside(pt_2) == False


def test_sample():
    n = 100
    domain = Union(
        [
            EllipsoidDomain({"x": [-1, 1], "y": [-1, 1]}),
            CartesianDomain({"x": [-0.5, 0.5], "y": [-0.5, 0.5]}),
        ]
    )
    pts = domain.sample(n)
    assert isinstance(pts, LabelTensor)
    assert pts.shape[0] == n

    n = 105
    pts = domain.sample(n)
    assert isinstance(pts, LabelTensor)
    assert pts.shape[0] == n