from itertools import product

import numpy as np
from pathlib import Path
import pytest

vtk = pytest.importorskip("vtk")
from vtk.util import numpy_support as nps

import openmc

@pytest.fixture
def model():
    openmc.reset_auto_ids()

    surf1 = openmc.Sphere(r=10, boundary_type='vacuum')
    surf2 = openmc.XPlane(x0=-0.001, boundary_type='vacuum')

    cell = openmc.Cell(region=-surf1 & -surf2)

    geometry = openmc.Geometry([cell])

    settings = openmc.Settings()
    settings.batches = 2
    settings.particles = 100
    settings.run_mode = 'fixed source'

    source = openmc.IndependentSource()
    source.angle = openmc.stats.Isotropic()
    source.energy = openmc.stats.Discrete([1.0e6], [1.0])
    source.space = openmc.stats.Point((-0.01, -0.01, -0.01))

    settings.source = source

    model = openmc.Model(geometry=geometry, settings=settings)

    return model


regular_mesh = openmc.RegularMesh()
regular_mesh.lower_left = (-10, -10, -10)
regular_mesh.upper_right = (10, 10, 10)
regular_mesh.dimension = [30, 20, 10]

rectilinear_mesh = openmc.RectilinearMesh()
rectilinear_mesh.x_grid = np.linspace(-10, 10, 6)
rectilinear_mesh.y_grid = np.logspace(0, 1, 7)
rectilinear_mesh.y_grid = \
    np.concatenate((-rectilinear_mesh.y_grid[::-1], rectilinear_mesh.y_grid))
rectilinear_mesh.z_grid = np.linspace(-10, 10, 11)

cylinder_mesh = openmc.CylindricalMesh(
    r_grid=np.linspace(0, 10, 23),
    z_grid=np.linspace(0, 1, 15)
)
cylinder_mesh.phi_grid = np.linspace(0, np.pi, 21)

spherical_mesh = openmc.SphericalMesh(
    r_grid=np.linspace(1, 10, 30),
    phi_grid=np.linspace(0, 0.8*np.pi, 25),
    theta_grid=np.linspace(0, np.pi / 2, 15),
)

MESHES = [cylinder_mesh, regular_mesh, rectilinear_mesh, spherical_mesh]

x_plane = openmc.XPlane(x0=-0.001, boundary_type='vacuum')
y_plane = openmc.YPlane(y0=-0.001, boundary_type='vacuum')
z_plane = openmc.ZPlane(z0=-0.001, boundary_type='vacuum')

SURFS = [x_plane, y_plane, z_plane]


def ids(mesh):
    if isinstance(mesh, openmc.CylindricalMesh):
        return 'cylindrical_mesh'
    elif isinstance(mesh, openmc.RegularMesh):
        return 'regular_mesh'
    elif isinstance(mesh, openmc.RectilinearMesh):
        return 'rectilinear_mesh'
    elif isinstance(mesh, openmc.SphericalMesh):
        return 'spherical_mesh'


@pytest.mark.parametrize("mesh", MESHES, ids=ids)
def test_write_data_to_vtk(mesh, tmpdir):
    # BUILD
    filename = Path(tmpdir) / "out.vtk"

    # use mesh element volumes as data to check volume-normalization ordering
    # kji (i changing fastest) orering is expected for input data
    # by using the volumes transposed as the data here, we can ensure the
    # normalization is happening correctly
    data = mesh.volumes

    # RUN
    mesh.write_data_to_vtk(filename=filename, datasets={"label1": data, "label2": data})

    # TEST
    assert filename.is_file()

    # read file
    reader = vtk.vtkStructuredGridReader()
    reader.SetFileName(str(filename))
    reader.Update()

    # check name of datasets
    vtk_grid = reader.GetOutput()
    array1 = vtk_grid.GetCellData().GetArray(0)
    array2 = vtk_grid.GetCellData().GetArray(1)

    assert array1.GetName() == "label1"
    assert array2.GetName() == "label2"

    # check size of datasets
    data1 = nps.vtk_to_numpy(array1)
    data2 = nps.vtk_to_numpy(array2)
    assert data1.size == data.size
    assert data2.size == data.size

    assert all(data1 == data2)
    assert all(data1 == 1.0)


@pytest.mark.parametrize("mesh", MESHES, ids=ids)
def test_write_data_to_vtk_size_mismatch(mesh):
    """Checks that an error is raised when the size of the dataset
    doesn't match the mesh number of cells

    Parameters
    ----------
    mesh : openmc.StructuredMesh
        The mesh to test
    """
    right_size = mesh.n_elements
    data = np.random.random(right_size + 1)

    # Error message has \ in to escape characters that are otherwise recognized
    # by regex. These are needed to make the test string match the error message
    # string when using the match argument as that uses regular expression
    expected_error_msg = (
        fr"The size of the dataset 'label' \({len(data)}\) should be equal to "
        fr"the number of mesh cells \({mesh.n_elements}\)"
    )
    with pytest.raises(ValueError, match=expected_error_msg):
        mesh.write_data_to_vtk(filename="out.vtk", datasets={"label": data})

def test_write_data_to_vtk_round_trip(run_in_tmpdir):
    cmesh = openmc.CylindricalMesh(
        r_grid=(0.0, 1.0, 2.0),
        z_grid=(0.0, 2.0, 4.0, 5.0),
        phi_grid=(0.0, 3.0, 6.0),
    )

    smesh = openmc.SphericalMesh(
        r_grid=(0.0, 1.0, 2.0),
        theta_grid=(0.0, 0.5, 1.0, 2.0),
        phi_grid=(0.0, 3.0, 6.0),
    )
    rmesh = openmc.RegularMesh()
    rmesh.lower_left = (0.0, 0.0, 0.0)
    rmesh.upper_right = (1.0, 3.0, 5.0)
    rmesh.dimension = (2, 1, 6)

    for mesh in [smesh, cmesh, rmesh]:

        filename = "mesh.vtk"
        data = np.array([1.0] * 12)  # there are 12 voxels in each mesh
        mesh.write_data_to_vtk(
            filename=filename,
            datasets={"normalized": data},
            volume_normalization=True
        )

        reader = vtk.vtkStructuredGridReader()
        reader.SetFileName(filename)
        reader.ReadAllFieldsOn()
        reader.Update()

        cell_data = reader.GetOutput().GetCellData()
        uniform_array = cell_data.GetArray("normalized")
        num_tuples = uniform_array.GetNumberOfTuples()
        vtk_values = [uniform_array.GetValue(i) for i in range(num_tuples)]

        # checks that the vtk cell values are equal to the data / mesh volumes
        assert np.allclose(vtk_values, data / mesh.volumes.T.flatten())

        mesh.write_data_to_vtk(
            filename=filename,
            datasets={"not_normalized": data},
            volume_normalization=False,
        )

        reader = vtk.vtkStructuredGridReader()
        reader.SetFileName(filename)
        reader.ReadAllFieldsOn()
        reader.Update()

        cell_data = reader.GetOutput().GetCellData()
        uniform_array = cell_data.GetArray("not_normalized")
        num_tuples = uniform_array.GetNumberOfTuples()
        vtk_values = [uniform_array.GetValue(i) for i in range(num_tuples)]

        # checks that the vtk cell values are equal to the data
        assert np.array_equal(vtk_values, data)

def mesh_surf_id(param):
    if isinstance(param, openmc.MeshBase):
        return ids(param)
    elif isinstance(param, openmc.XPlane):
        return 'XPlane'
    elif isinstance(param, openmc.YPlane):
        return 'YPlane'
    elif isinstance(param, openmc.ZPlane):
        return 'ZPlane'


@pytest.mark.parametrize("mesh,surface", product(MESHES, SURFS), ids=mesh_surf_id)
def test_vtk_write_ordering(run_in_tmpdir, model, mesh, surface):

    tally = openmc.Tally()
    tally.scores = ['flux']
    # use the mesh on the specified tally
    mesh_filter = openmc.MeshFilter(mesh)
    tally.filters = [mesh_filter]

    model.tallies = openmc.Tallies([tally])

    # run the problem
    sp_filename = model.run()

    with openmc.StatePoint(sp_filename) as sp:
        mean = sp.tallies[tally.id].mean

    # write the data to a VTK file
    vtk_filename = 'test.vtk'
    mesh.write_data_to_vtk(vtk_filename, datasets={'mean': mean})

     # read file
    reader = vtk.vtkStructuredGridReader()
    reader.SetFileName(str(vtk_filename))
    reader.Update()

    # check name of datasets
    vtk_grid = reader.GetOutput()
    array = vtk_grid.GetCellData().GetArray(0)
    vtk_data = nps.vtk_to_numpy(array)

    # convenience function for determining if a mesh
    # element has vertices in the geometry. This
    # particular geometry allows us to assume that tally results
    # in the element should be zero if none of its vertices lie in the geometry
    def in_geom(cell):
        point_ids = cell.GetPointIds()

        for i in range(point_ids.GetNumberOfIds()):
            p = vtk_grid.GetPoint(point_ids.GetId(i))
            if model.geometry.find(p):
                return True

        return False

    # reshape mean according to mesh dimensions
    mean = mean.reshape(mesh.dimension[::-1]).T
    centroid = [0.0, 0.0, 0.0]

    # check that tally and vtk array results are zero where expected
    for ijk in mesh.indices:
        ijk = tuple(n - 1 for n in ijk)
        # get the cell from the stuctured mesh object
        cell = vtk_grid.GetCell(*ijk)
        if not in_geom(cell):
            cell.GetCentroid(centroid)
            err_msg = f'IJK: {ijk} should be zero but is not. Centroid: {centroid}'
            assert mean[ijk] == 0.0, err_msg

            # need to get flat index with axes reversed due to ordering passed into the VTK file
            flat_idx = np.ravel_multi_index(tuple(ijk[::-1]), mesh.dimension[::-1])
            assert vtk_data[flat_idx] == 0.0, err_msg


def test_sphere_mesh_coordinates(run_in_tmpdir):
    mesh = openmc.SphericalMesh(
        r_grid=np.linspace(0.1, 10, 30),
        phi_grid=np.linspace(0, 1.5*np.pi, 25),
        theta_grid=np.linspace(0, np.pi / 2, 15),
    )
    # write the data to a VTK file (no data)
    vtk_filename = 'test.vtk'
    mesh.write_data_to_vtk(vtk_filename, {})

    # read file
    reader = vtk.vtkStructuredGridReader()
    reader.SetFileName(str(vtk_filename))
    reader.Update()

    vtk_grid = reader.GetOutput()

    # create a region that matches the spherical mesh description
    x = openmc.XPlane()
    z = openmc.ZPlane()
    y = openmc.YPlane()
    s = openmc.Sphere(r=10.0)

    region = +z & +y & -s | -x & -y & +z & -s

    # the VTK interface will update this list when GetCentroid is called
    centroid = np.zeros(3)

    # ensure all centroids of the sphere mesh are inside the cell region
    for i in range(vtk_grid.GetNumberOfCells()):
        # get the cell from the stuctured mesh object
        cell = vtk_grid.GetCell(i)
        cell.GetCentroid(centroid)

        # if the coordinate conversion is happening correctly,
        # every one of the cell centroids should be in the CSG region
        assert centroid in region, \
            f'Cell centroid {centroid} not in equivalent ' \
            f'CSG region for spherical mesh {mesh}'