from __future__ import annotations

import numpy as np
import pytest

import pyvista as pv


@pytest.fixture
def tiny_rectilinear():
    xrng = np.linspace(0, 3, 4)
    yrng = np.linspace(0, 3, 5)
    zrng = np.linspace(0, 3, 6)
    return pv.RectilinearGrid(xrng, yrng, zrng)


@pytest.mark.parametrize('tetra_per_cell', [5, 6, 12])
def test_to_tetrahedral(tiny_rectilinear, tetra_per_cell):
    tet_grid = tiny_rectilinear.to_tetrahedra(tetra_per_cell=tetra_per_cell)
    assert tet_grid.n_cells == tiny_rectilinear.n_cells * tetra_per_cell


def test_to_tetrahedral_raise(tiny_rectilinear):
    with pytest.raises(ValueError, match='either 5, 6, or 12'):
        tiny_rectilinear.to_tetrahedra(tetra_per_cell=9)


def test_to_tetrahedral_mixed(tiny_rectilinear):
    data = np.empty(tiny_rectilinear.n_cells, dtype=int)
    half = tiny_rectilinear.n_cells // 2
    data[:half] = 5
    data[half:] = 12
    tiny_rectilinear['data'] = data
    tet_grid = tiny_rectilinear.to_tetrahedra(mixed=True)
    assert tet_grid.n_cells == (half * 5 + half * 12)

    tet_grid = tiny_rectilinear.to_tetrahedra(mixed='data')
    assert tet_grid.n_cells == (half * 5 + half * 12)

    tet_grid = tiny_rectilinear.to_tetrahedra(mixed=data)
    assert tet_grid.n_cells == (half * 5 + half * 12)

    with pytest.raises(TypeError, match='mixed'):
        tet_grid = tiny_rectilinear.to_tetrahedra(mixed=123)


def test_to_tetrahedral_edge_case():
    with pytest.raises(RuntimeError, match='is 1'):
        pv.ImageData(dimensions=(1, 2, 2)).to_tetrahedra(tetra_per_cell=12)


def test_to_tetrahedral_pass_cell_ids(tiny_rectilinear):
    tet_grid = tiny_rectilinear.to_tetrahedra(pass_cell_ids=False, pass_data=False)
    assert not tet_grid.cell_data
    tet_grid = tiny_rectilinear.to_tetrahedra(pass_cell_ids=True, pass_data=False)
    assert 'vtkOriginalCellIds' in tet_grid.cell_data
    assert np.issubdtype(tet_grid.cell_data['vtkOriginalCellIds'].dtype, np.integer)


def test_to_tetrahedral_pass_cell_data(tiny_rectilinear):
    # test that data isn't passed
    tiny_rectilinear['cell_data'] = np.ones(tiny_rectilinear.n_cells)
    tiny_rectilinear['point_data'] = np.arange(tiny_rectilinear.n_points)
    tiny_rectilinear.set_active_scalars('cell_data')
    tet_grid = tiny_rectilinear.to_tetrahedra(pass_cell_ids=False, pass_data=False)
    assert not tet_grid.cell_data
    assert not tet_grid.point_data

    # test with cell data
    tet_grid = tiny_rectilinear.to_tetrahedra(pass_cell_ids=False, pass_data=True)
    assert tet_grid.cell_data
    assert 'cell_data' in tet_grid.cell_data
    assert 'point_data' in tet_grid.point_data
    assert tet_grid.active_scalars_name == 'cell_data'

    # automatically added
    assert 'vtkOriginalCellIds' in tet_grid.cell_data

    # Test point data active
    tiny_rectilinear.set_active_scalars('point_data')
    assert tiny_rectilinear.active_scalars_name == 'point_data'
    tet_grid = tiny_rectilinear.to_tetrahedra(pass_cell_ids=False, pass_data=True)
    assert tet_grid.active_scalars_name == 'point_data'
