from math import pi
from tempfile import TemporaryDirectory
from pathlib import Path

import h5py
import numpy as np
from scipy.stats import chi2
import pytest
import openmc
import openmc.lib
from openmc.utility_funcs import change_directory
from uncertainties.unumpy import uarray, nominal_values, std_devs


@pytest.mark.parametrize("val_left,val_right", [(0, 0), (-1., -1.), (2.0, 2)])
def test_raises_error_when_flat(val_left, val_right):
    """Checks that an error is raised when a mesh is flat"""
    mesh = openmc.RegularMesh()

    # Same X
    with pytest.raises(ValueError):
        mesh.lower_left = [val_left, -25, -25]
        mesh.upper_right = [val_right, 25, 25]

    with pytest.raises(ValueError):
        mesh.upper_right = [val_right, 25, 25]
        mesh.lower_left = [val_left, -25, -25]

    # Same Y
    with pytest.raises(ValueError):
        mesh.lower_left = [-25, val_left, -25]
        mesh.upper_right = [25, val_right, 25]

    with pytest.raises(ValueError):
        mesh.upper_right = [25, val_right, 25]
        mesh.lower_left = [-25, val_left, -25]

    # Same Z
    with pytest.raises(ValueError):
        mesh.lower_left = [-25, -25, val_left]
        mesh.upper_right = [25, 25, val_right]

    with pytest.raises(ValueError):
        mesh.upper_right = [25, 25, val_right]
        mesh.lower_left = [-25, -25, val_left]


def test_regular_mesh_bounding_box():
    mesh = openmc.RegularMesh()
    mesh.lower_left = [-2, -3, -5]
    mesh.upper_right = [2, 3, 5]
    bb = mesh.bounding_box
    assert isinstance(bb, openmc.BoundingBox)
    np.testing.assert_array_equal(bb.lower_left, (-2, -3 ,-5))
    np.testing.assert_array_equal(bb.upper_right, (2, 3, 5))


def test_rectilinear_mesh_bounding_box():
    mesh = openmc.RectilinearMesh()
    mesh.x_grid = [0., 1., 5., 10.]
    mesh.y_grid = [-10., -5., 0.]
    mesh.z_grid = [-100., 0., 100.]
    bb = mesh.bounding_box
    assert isinstance(bb, openmc.BoundingBox)
    np.testing.assert_array_equal(bb.lower_left, (0., -10. ,-100.))
    np.testing.assert_array_equal(bb.upper_right, (10., 0., 100.))


def test_cylindrical_mesh_bounding_box():
    # test with mesh at origin (0, 0, 0)
    mesh = openmc.CylindricalMesh(
        r_grid=[0.1, 0.2, 0.5, 1.],
        z_grid=[0.1, 0.2, 0.4, 0.6, 1.],
        origin=(0, 0, 0)
    )
    np.testing.assert_array_equal(mesh.upper_right, (1, 1, 1))
    np.testing.assert_array_equal(mesh.lower_left, (-1, -1, 0.1))
    bb = mesh.bounding_box
    assert isinstance(bb, openmc.BoundingBox)
    np.testing.assert_array_equal(bb.lower_left, (-1, -1, 0.1))
    np.testing.assert_array_equal(bb.upper_right, (1, 1, 1))

    # test with mesh at origin (3, 5, 7)
    mesh.origin = (3, 5, 7)
    np.testing.assert_array_equal(mesh.upper_right, (4, 6, 8))
    np.testing.assert_array_equal(mesh.lower_left, (2, 4, 7.1))
    bb = mesh.bounding_box
    assert isinstance(bb, openmc.BoundingBox)
    np.testing.assert_array_equal(bb.lower_left, (2, 4, 7.1))
    np.testing.assert_array_equal(bb.upper_right, (4, 6, 8))

    # changing z grid to contain negative numbers
    mesh.z_grid = [-10, 0, 10]
    np.testing.assert_array_equal(mesh.lower_left, (2, 4, -3))
    np.testing.assert_array_equal(mesh.upper_right, (4, 6, 17))


def test_spherical_mesh_bounding_box():
    # test with mesh at origin (0, 0, 0)
    mesh = openmc.SphericalMesh([0.1, 0.2, 0.5, 1.], origin=(0., 0., 0.))
    np.testing.assert_array_equal(mesh.upper_right, (1, 1, 1))
    np.testing.assert_array_equal(mesh.lower_left, (-1, -1, -1))
    bb = mesh.bounding_box
    assert isinstance(bb, openmc.BoundingBox)
    np.testing.assert_array_equal(bb.lower_left, (-1, -1, -1))
    np.testing.assert_array_equal(bb.upper_right, (1, 1, 1))

    # test with mesh at origin (3, 5, 7)
    mesh.origin = (3, 5, 7)
    np.testing.assert_array_equal(mesh.upper_right, (4, 6, 8))
    np.testing.assert_array_equal(mesh.lower_left, (2, 4, 6))
    bb = mesh.bounding_box
    assert isinstance(bb, openmc.BoundingBox)
    np.testing.assert_array_equal(bb.lower_left, (2, 4, 6))
    np.testing.assert_array_equal(bb.upper_right, (4, 6, 8))


def test_SphericalMesh_initiation():
    # test defaults
    mesh = openmc.SphericalMesh(r_grid=(0, 10))
    assert (mesh.origin == np.array([0, 0, 0])).all()
    assert (mesh.r_grid == np.array([0, 10])).all()
    assert (mesh.theta_grid == np.array([0, pi])).all()
    assert (mesh.phi_grid == np.array([0, 2*pi])).all()

    # test setting on creation
    mesh = openmc.SphericalMesh(
        origin=(1, 2, 3),
        r_grid=(0, 2),
        theta_grid=(1, 3),
        phi_grid=(2, 4)
    )
    assert (mesh.origin == np.array([1, 2, 3])).all()
    assert (mesh.r_grid == np.array([0., 2.])).all()
    assert (mesh.theta_grid == np.array([1, 3])).all()
    assert (mesh.phi_grid == np.array([2, 4])).all()

    # test attribute changing
    mesh.r_grid = (0, 11)
    assert (mesh.r_grid == np.array([0., 11.])).all()

    # test invalid r_grid values
    with pytest.raises(ValueError):
        openmc.SphericalMesh(r_grid=[1, 1])

    with pytest.raises(ValueError):
        openmc.SphericalMesh(r_grid=[0])

    # test invalid theta_grid values
    with pytest.raises(ValueError):
        openmc.SphericalMesh(r_grid=[1, 2], theta_grid=[1, 1])

    with pytest.raises(ValueError):
        openmc.SphericalMesh(r_grid=[1, 2], theta_grid=[0])

    # test invalid phi_grid values
    with pytest.raises(ValueError):
        openmc.SphericalMesh(r_grid=[1, 2], phi_grid=[1, 1])

    with pytest.raises(ValueError):
        openmc.SphericalMesh(r_grid=[1, 2], phi_grid=[0])

    # waffles and pancakes are unfortunately not valid radii
    with pytest.raises(TypeError):
        openmc.SphericalMesh(('🧇', '🥞'))


def test_CylindricalMesh_initiation():
    # test defaults
    mesh = openmc.CylindricalMesh(r_grid=(0, 10), z_grid=(0, 10))
    assert (mesh.origin == np.array([0, 0, 0])).all()
    assert (mesh.r_grid == np.array([0, 10])).all()
    assert (mesh.phi_grid == np.array([0, 2*pi])).all()
    assert (mesh.z_grid == np.array([0, 10])).all()

    # test setting on creation
    mesh = openmc.CylindricalMesh(
        origin=(1, 2, 3),
        r_grid=(0, 2),
        z_grid=(1, 3),
        phi_grid=(2, 4)
    )
    assert (mesh.origin == np.array([1, 2, 3])).all()
    assert (mesh.r_grid == np.array([0., 2.])).all()
    assert (mesh.z_grid == np.array([1, 3])).all()
    assert (mesh.phi_grid == np.array([2, 4])).all()

    # test attribute changing
    mesh.r_grid = (0., 10.)
    assert (mesh.r_grid == np.array([0, 10.])).all()
    mesh.z_grid = (0., 4.)
    assert (mesh.z_grid == np.array([0, 4.])).all()

    # waffles and pancakes are unfortunately not valid radii
    with pytest.raises(TypeError):
        openmc.SphericalMesh(('🧇', '🥞'))


def test_invalid_cylindrical_mesh_errors():
    # Test invalid r_grid values
    with pytest.raises(ValueError):
        openmc.CylindricalMesh(r_grid=[5, 1], phi_grid=[0, pi], z_grid=[0, 10])

    with pytest.raises(ValueError):
        openmc.CylindricalMesh(r_grid=[1, 2, 4, 3], phi_grid=[0, pi], z_grid=[0, 10])

    with pytest.raises(ValueError):
        openmc.CylindricalMesh(r_grid=[1], phi_grid=[0, pi], z_grid=[0, 10])

    # Test invalid phi_grid values
    with pytest.raises(ValueError):
        openmc.CylindricalMesh(r_grid=[0, 1, 2], phi_grid=[-1, 3], z_grid=[0, 10])

    with pytest.raises(ValueError):
        openmc.CylindricalMesh(
            r_grid=[0, 1, 2],
            phi_grid=[0, 2*pi + 0.1],
            z_grid=[0, 10]
        )

    with pytest.raises(ValueError):
        openmc.CylindricalMesh(r_grid=[0, 1, 2], phi_grid=[pi], z_grid=[0, 10])

    # Test invalid z_grid values
    with pytest.raises(ValueError):
        openmc.CylindricalMesh(r_grid=[0, 1, 2], phi_grid=[0, pi], z_grid=[5])

    with pytest.raises(ValueError):
        openmc.CylindricalMesh(r_grid=[0, 1, 2], phi_grid=[0, pi], z_grid=[5, 1])

    with pytest.raises(ValueError):
        openmc.CylindricalMesh(r_grid=[1, 2, 4, 3], phi_grid=[0, pi], z_grid=[0, 10, 5])


def test_centroids():
    # regular mesh
    mesh = openmc.RegularMesh()
    mesh.lower_left = (1., 2., 3.)
    mesh.upper_right = (11., 12., 13.)
    mesh.dimension = (1, 1, 1)
    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [6., 7., 8.])

    # rectilinear mesh
    mesh = openmc.RectilinearMesh()
    mesh.x_grid = [1., 11.]
    mesh.y_grid = [2., 12.]
    mesh.z_grid = [3., 13.]
    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [6., 7., 8.])

    # cylindrical mesh
    mesh = openmc.CylindricalMesh(r_grid=(0, 10), z_grid=(0, 10), phi_grid=(0, np.pi))
    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [0.0, 5.0, 5.0])
    # ensure that setting an origin is handled correctly
    mesh.origin = (5.0, 0, -10)
    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [5.0, 5.0, -5.0])

    # spherical mesh, single element xyz-positive octant
    mesh = openmc.SphericalMesh(r_grid=[0, 10], theta_grid=[0, 0.5*np.pi], phi_grid=[0, np.pi])
    x = 5.*np.cos(0.5*np.pi)*np.sin(0.25*np.pi)
    y = 5.*np.sin(0.5*np.pi)*np.sin(0.25*np.pi)
    z = 5.*np.sin(0.25*np.pi)
    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [x, y, z])

    mesh.origin = (-5.0, -5.0, 5.0)
    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [x-5.0, y-5.0, z+5.0])


@pytest.mark.parametrize('mesh_type', ('regular', 'rectilinear', 'cylindrical', 'spherical'))
def test_mesh_vertices(mesh_type):

    ijk = (2, 3, 2)

    # create a new mesh object
    if mesh_type == 'regular':
        mesh = openmc.RegularMesh()
        ll = np.asarray([0.]*3)
        width = np.asarray([0.5]*3)
        mesh.lower_left = ll
        mesh.width = width
        mesh.dimension = (5, 7, 9)

        # spot check that an element has the correct vertex coordinates asociated with it
        # (using zero-indexing here)
        exp_i_j_k = ll + np.asarray(ijk, dtype=float) * width
        np.testing.assert_equal(mesh.vertices[ijk], exp_i_j_k)

        # shift the mesh using the llc
        shift  = np.asarray((3.0, 6.0, 10.0))
        mesh.lower_left += shift
        np.testing.assert_equal(mesh.vertices[ijk], exp_i_j_k+shift)
    elif mesh_type == 'rectilinear':
        mesh = openmc.RectilinearMesh()
        w = np.asarray([0.5] * 3)
        ll = np.asarray([0.]*3)
        dims = (5, 7, 9)
        mesh.x_grid = np.linspace(ll[0], w[0]*dims[0], dims[0])
        mesh.y_grid = np.linspace(ll[1], w[1]*dims[1], dims[1])
        mesh.z_grid = np.linspace(ll[2], w[2]*dims[2], dims[2])
        exp_vert = np.asarray((mesh.x_grid[2], mesh.y_grid[3], mesh.z_grid[2]))
        np.testing.assert_equal(mesh.vertices[ijk], exp_vert)
    elif mesh_type == 'cylindrical':
        r_grid = np.linspace(0, 5, 10)
        z_grid = np.linspace(-10, 10, 20)
        phi_grid = np.linspace(0, 2*np.pi, 8)
        mesh = openmc.CylindricalMesh(r_grid=r_grid, z_grid=z_grid, phi_grid=phi_grid)
        exp_vert = np.asarray((mesh.r_grid[2], mesh.phi_grid[3], mesh.z_grid[2]))
        np.testing.assert_equal(mesh.vertices_cylindrical[ijk], exp_vert)
    elif mesh_type == 'spherical':
        r_grid = np.linspace(0, 13, 14)
        theta_grid = np.linspace(0, np.pi, 11)
        phi_grid = np.linspace(0, 2*np.pi, 7)
        mesh = openmc.SphericalMesh(r_grid=r_grid, theta_grid=theta_grid, phi_grid=phi_grid)
        exp_vert = np.asarray((mesh.r_grid[2], mesh.theta_grid[3], mesh.phi_grid[2]))
        np.testing.assert_equal(mesh.vertices_spherical[ijk], exp_vert)


def test_CylindricalMesh_get_indices_at_coords():
    # default origin (0, 0, 0) and default phi grid (0, 2*pi)
    mesh = openmc.CylindricalMesh(r_grid=(0, 5, 10), z_grid=(0, 5, 10))
    assert mesh.get_indices_at_coords([1, 0, 1]) == (0, 0, 0)
    assert mesh.get_indices_at_coords([6, 0, 1]) == (1, 0, 0)
    assert mesh.get_indices_at_coords([9, 0, 1]) == (1, 0, 0)
    assert mesh.get_indices_at_coords([0, 6, 0]) == (1, 0, 0)
    assert mesh.get_indices_at_coords([0, 9, 6]) == (1, 0, 1)
    assert mesh.get_indices_at_coords([-2, -2, 9]) == (0, 0, 1)

    with pytest.raises(ValueError):
        assert mesh.get_indices_at_coords([8, 8, 1])  # resulting r value to large
    with pytest.raises(ValueError):
        assert mesh.get_indices_at_coords([-8, -8, 1])  # resulting r value to large
    with pytest.raises(ValueError):
        assert mesh.get_indices_at_coords([1, 0, -1])  # z value below range
    with pytest.raises(ValueError):
        assert mesh.get_indices_at_coords([1, 0, 11])  # z value above range

    assert mesh.get_indices_at_coords([1, 1, 1]) == (0, 0, 0)

    # negative range on z grid
    mesh = openmc.CylindricalMesh(
        r_grid=(0, 5, 10),
        phi_grid=(0, 0.5 * pi, pi, 1.5 * pi, 1.9 * pi),
        z_grid=(-5, 0, 5, 10),
    )
    assert mesh.get_indices_at_coords([1, 1, 1]) == (0, 0, 1)  # first angle quadrant
    assert mesh.get_indices_at_coords([2, 2, 6]) == (0, 0, 2)  # first angle quadrant
    assert mesh.get_indices_at_coords([-2, 0.1, -1]) == (0, 1, 0)  # second angle quadrant
    assert mesh.get_indices_at_coords([-2, -0.1, -1]) == (0, 2, 0)  # third angle quadrant
    assert mesh.get_indices_at_coords([2, -0.9, -1]) == (0, 3, 0)  # forth angle quadrant

    with pytest.raises(ValueError):
        assert mesh.get_indices_at_coords([2, -0.1, 1])  # outside of phi range

    # origin of mesh not default
    mesh = openmc.CylindricalMesh(
        r_grid=(0, 5, 10),
        phi_grid=(0, 0.5 * pi, pi, 1.5 * pi, 1.9 * pi),
        z_grid=(-5, 0, 5, 10),
        origin=(100, 200, 300),
    )
    assert mesh.get_indices_at_coords([101, 201, 301]) == (0, 0, 1)  # first angle quadrant
    assert mesh.get_indices_at_coords([102, 202, 306]) == (0, 0, 2)  # first angle quadrant
    assert mesh.get_indices_at_coords([98, 200.1, 299]) == (0, 1, 0)  # second angle quadrant
    assert mesh.get_indices_at_coords([98, 199.9, 299]) == (0, 2, 0)  # third angle quadrant
    assert mesh.get_indices_at_coords([102, 199.1, 299]) == (0, 3, 0)  # forth angle quadrant


def test_mesh_name_roundtrip(run_in_tmpdir):

    mesh = openmc.RegularMesh()
    mesh.name = 'regular-mesh'
    mesh.lower_left = (-1, -1, -1)
    mesh.width = (1, 1, 1)
    mesh.dimension = (1, 1, 1)

    mesh_filter = openmc.MeshFilter(mesh)
    tally = openmc.Tally()
    tally.filters = [mesh_filter]
    tally.scores = ['flux']

    openmc.Tallies([tally]).export_to_xml()

    xml_tallies = openmc.Tallies.from_xml()

    mesh = xml_tallies[0].find_filter(openmc.MeshFilter).mesh
    assert mesh.name == 'regular-mesh'


def test_umesh_roundtrip(run_in_tmpdir, request):
    umesh = openmc.UnstructuredMesh(request.path.parent / 'test_mesh_tets.e', 'moab')
    umesh.output = True

    # create a tally using this mesh
    mf = openmc.MeshFilter(umesh)
    tally = openmc.Tally()
    tally.filters = [mf]
    tally.scores = ['flux']

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

    xml_tallies = openmc.Tallies.from_xml()
    xml_tally = xml_tallies[0]
    xml_mesh = xml_tally.filters[0].mesh

    assert umesh.id == xml_mesh.id


@pytest.fixture(scope='module')
def simple_umesh(request):
    """Fixture returning UnstructuredMesh with all attributes"""
    surf1 = openmc.Sphere(r=20.0, boundary_type="vacuum")
    material1 = openmc.Material()
    material1.add_element("H", 1.0)
    material1.set_density('g/cm3', 1.0)

    materials = openmc.Materials([material1])
    cell1 = openmc.Cell(region=-surf1, fill=material1)
    geometry = openmc.Geometry([cell1])

    umesh = openmc.UnstructuredMesh(
       filename=request.path.parent.parent
        / "regression_tests/external_moab/test_mesh_tets.h5m",
       library="moab",
       mesh_id=1
    )
    # setting ID to make it easier to get the mesh from the statepoint later
    mesh_filter = openmc.MeshFilter(umesh)

    # Create flux mesh tally to score alpha production
    mesh_tally = openmc.Tally(name="test_tally")
    mesh_tally.filters = [mesh_filter]
    mesh_tally.scores = ["total"]

    tallies = openmc.Tallies([mesh_tally])

    settings = openmc.Settings()
    settings.run_mode = "fixed source"
    settings.batches = 2
    settings.particles = 100
    settings.source = openmc.IndependentSource(
        space=openmc.stats.Point((0.1, 0.1, 0.1))
    )

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

    with change_directory(tmpdir=True):
        statepoint_file = model.run()
        with openmc.StatePoint(statepoint_file) as sp:
            return sp.meshes[1]


@pytest.mark.skipif(not openmc.lib._dagmc_enabled(), reason="DAGMC not enabled.")
@pytest.mark.parametrize('export_type', ('.vtk', '.vtu'))
def test_umesh(run_in_tmpdir, simple_umesh, export_type):
    """Performs a minimal UnstructuredMesh simulation, reads in the resulting
    statepoint file and writes the mesh data to vtk and vtkhdf files. It is
    necessary to read in the unstructured mesh from a statepoint file to ensure
    it has all the required attributes
    """
    # Get VTK modules
    vtkIOLegacy = pytest.importorskip("vtkmodules.vtkIOLegacy")
    vtkIOXML = pytest.importorskip("vtkmodules.vtkIOXML")

    # Sample some random data and write to VTK
    rng = np.random.default_rng()
    ref_data = rng.random(simple_umesh.dimension)
    filename = f"test_mesh{export_type}"
    simple_umesh.write_data_to_vtk(datasets={"mean": ref_data}, filename=filename)

    assert Path(filename).exists()

    if export_type == ".vtk":
        reader = vtkIOLegacy.vtkGenericDataObjectReader()
    elif export_type == ".vtu":
        reader = vtkIOXML.vtkXMLGenericDataObjectReader()
    reader.SetFileName(str(filename))
    reader.Update()

    # Get mean from file and make sure it matches original data
    arr = reader.GetOutput().GetCellData().GetArray("mean")
    mean = np.array([arr.GetTuple1(i) for i in range(ref_data.size)])
    np.testing.assert_almost_equal(mean, ref_data)

    # attempt to apply a dataset with an improper size to a VTK write
    with pytest.raises(ValueError, match='Cannot apply dataset "mean"'):
        simple_umesh.write_data_to_vtk(datasets={'mean': ref_data[:-2]}, filename=filename)


@pytest.mark.skipif(not openmc.lib._dagmc_enabled(), reason="DAGMC not enabled.")
def test_write_vtkhdf(request, run_in_tmpdir):
    """Performs a minimal UnstructuredMesh simulation, reads in the resulting
    statepoint file and writes the mesh data to vtk and vtkhdf files. It is
    necessary to read in the unstructured mesh from a statepoint file to ensure
    it has all the required attributes
    """
    model = openmc.Model()

    surf1 = openmc.Sphere(r=1000.0, boundary_type="vacuum")
    cell1 = openmc.Cell(region=-surf1)
    model.geometry = openmc.Geometry([cell1])

    umesh = openmc.UnstructuredMesh(
        request.path.parent / "test_mesh_dagmc_tets.vtk",
        "moab",
        mesh_id = 1
    )
    mesh_filter = openmc.MeshFilter(umesh)

    # Create flux mesh tally to score alpha production
    mesh_tally = openmc.Tally(name="test_tally")
    mesh_tally.filters = [mesh_filter]
    mesh_tally.scores = ["flux"]

    model.tallies = [mesh_tally]

    model.settings.run_mode = "fixed source"
    model.settings.batches = 2
    model.settings.particles = 10

    statepoint_file = model.run()

    with openmc.StatePoint(statepoint_file) as statepoint:
        my_tally = statepoint.get_tally(name="test_tally")

    umesh_from_sp = statepoint.meshes[umesh.id]

    datasets={
        "mean": my_tally.mean.flatten(),
        "std_dev": my_tally.std_dev.flatten()
    }

    umesh_from_sp.write_data_to_vtk(datasets=datasets, filename="test_mesh.vtkhdf")
    umesh_from_sp.write_data_to_vtk(datasets=datasets, filename="test_mesh.vtk")

    with pytest.raises(ValueError, match="Unsupported file extension"):
        # Supported file extensions are vtk or vtkhdf, not hdf5, so this should raise an error
        umesh_from_sp.write_data_to_vtk(
            datasets=datasets,
            filename="test_mesh.hdf5",
        )
    with pytest.raises(ValueError, match="Cannot apply dataset"):
        # The shape of the data should match the shape of the mesh, so this should raise an error
        umesh_from_sp.write_data_to_vtk(
            datasets={'incorrectly_shaped_data': np.array(([1,2,3]))},
            filename="test_mesh_incorrect_shape.vtkhdf",
        )

    assert Path("test_mesh.vtk").exists()
    assert Path("test_mesh.vtkhdf").exists()

    # just ensure we can open the file without error
    with h5py.File("test_mesh.vtkhdf", "r"):
        ...


def test_mesh_get_homogenized_materials():
    """Test the get_homogenized_materials method"""
    # Simple model with 1 cm of Fe56 next to 1 cm of H1
    fe = openmc.Material()
    fe.add_nuclide('Fe56', 1.0)
    fe.set_density('g/cm3', 5.0)
    h = openmc.Material()
    h.add_nuclide('H1', 1.0)
    h.set_density('g/cm3', 1.0)

    x0 = openmc.XPlane(-1.0, boundary_type='vacuum')
    x1 = openmc.XPlane(0.0)
    x2 = openmc.XPlane(1.0)
    x3 = openmc.XPlane(2.0, boundary_type='vacuum')
    cell1 = openmc.Cell(fill=fe, region=+x0 & -x1)
    cell2 = openmc.Cell(fill=h, region=+x1 & -x2)
    cell_empty = openmc.Cell(region=+x2 & -x3)
    model = openmc.Model(geometry=openmc.Geometry([cell1, cell2, cell_empty]))
    model.settings.particles = 1000
    model.settings.batches = 10

    mesh = openmc.RegularMesh()
    mesh.lower_left = (-1., -1., -1.)
    mesh.upper_right = (1., 1., 1.)
    mesh.dimension = (3, 1, 1)
    m1, m2, m3 = mesh.get_homogenized_materials(model, n_samples=10_000)

    # Left mesh element should be only Fe56
    assert m1.get_mass_density('Fe56') == pytest.approx(5.0)

    # Middle mesh element should be 50% Fe56 and 50% H1
    assert m2.get_mass_density('Fe56') == pytest.approx(2.5, rel=1e-2)
    assert m2.get_mass_density('H1') == pytest.approx(0.5, rel=1e-2)

    # Right mesh element should be only H1
    assert m3.get_mass_density('H1') == pytest.approx(1.0)

    mesh_void = openmc.RegularMesh()
    mesh_void.lower_left = (0.5, 0.5, -1.)
    mesh_void.upper_right = (1.5, 1.5, 1.)
    mesh_void.dimension = (1, 1, 1)
    m4, = mesh_void.get_homogenized_materials(model, n_samples=(100, 100, 0))

    # Mesh element that overlaps void should have half density
    assert m4.get_mass_density('H1') == pytest.approx(0.5, rel=1e-2)

    # If not including void, density of homogenized material should be same as
    # original material
    m5, = mesh_void.get_homogenized_materials(
        model, n_samples=1000, include_void=False)
    assert m5.get_mass_density('H1') == pytest.approx(1.0)


@pytest.fixture
def sphere_model():
    # Model with three materials separated by planes x=0 and z=0
    mats = []
    for i in range(3):
        mat = openmc.Material()
        mat.add_nuclide('H1', 1.0)
        mat.set_density('g/cm3', float(i + 1))
        mats.append(mat)

    sph = openmc.Sphere(r=25.0, boundary_type='vacuum')
    x0 = openmc.XPlane(0.0)
    z0 = openmc.ZPlane(0.0)
    cell1 = openmc.Cell(fill=mats[0], region=-sph & +x0 & +z0)
    cell2 = openmc.Cell(fill=mats[1], region=-sph & -x0 & +z0)
    cell3 = openmc.Cell(fill=mats[2], region=-sph & -z0)
    model = openmc.Model()
    model.geometry = openmc.Geometry([cell1, cell2, cell3])
    model.materials = openmc.Materials(mats)
    return model


@pytest.mark.parametrize("n_rays", [1000, (10, 10, 0), (10, 0, 10), (0, 10, 10)])
def test_material_volumes_regular_mesh(sphere_model, n_rays):
    """Test the material_volumes method on a regular mesh"""
    mesh = openmc.RegularMesh()
    mesh.lower_left = (-1., -1., -1.)
    mesh.upper_right = (1., 1., 1.)
    mesh.dimension = (2, 2, 2)
    volumes = mesh.material_volumes(sphere_model, n_rays)
    mats = sphere_model.materials
    np.testing.assert_almost_equal(volumes[mats[0].id], [0., 0., 0., 0., 0., 1., 0., 1.])
    np.testing.assert_almost_equal(volumes[mats[1].id], [0., 0., 0., 0., 1., 0., 1., 0.])
    np.testing.assert_almost_equal(volumes[mats[2].id], [1., 1., 1., 1., 0., 0., 0., 0.])
    assert volumes.by_element(4) == [(mats[1].id, 1.)]
    assert volumes.by_element(0) == [(mats[2].id, 1.)]


def test_material_volumes_cylindrical_mesh(sphere_model):
    """Test the material_volumes method on a cylindrical mesh"""
    cyl_mesh = openmc.CylindricalMesh(
        [0., 1.], [-1., 0., 1.,], [0.0, pi/4, 3*pi/4, 5*pi/4, 7*pi/4, 2*pi])
    volumes = cyl_mesh.material_volumes(sphere_model, (0, 100, 100))
    mats = sphere_model.materials
    np.testing.assert_almost_equal(volumes[mats[0].id], [
        0., 0., 0., 0., 0.,
        pi/8, pi/8, 0., pi/8, pi/8
    ])
    np.testing.assert_almost_equal(volumes[mats[1].id], [
        0., 0., 0., 0., 0.,
        0., pi/8, pi/4, pi/8, 0.
    ])
    np.testing.assert_almost_equal(volumes[mats[2].id], [
        pi/8, pi/4, pi/4, pi/4, pi/8,
        0., 0., 0., 0., 0.
    ])


def test_mesh_material_volumes_serialize():
    materials = np.array([
        [1, -1, -2],
        [-1, -2, -2],
        [2, 1, -2],
        [2, -2, -2]
    ])
    volumes = np.array([
        [0.5, 0.5, 0.0],
        [1.0, 0.0, 0.0],
        [0.5, 0.5, 0.0],
        [1.0, 0.0, 0.0]
    ])
    volumes = openmc.MeshMaterialVolumes(materials, volumes)
    with TemporaryDirectory() as tmpdir:
        path = f'{tmpdir}/volumes.npz'
        volumes.save(path)
        new_volumes = openmc.MeshMaterialVolumes.from_npz(path)

    assert new_volumes.by_element(0) == [(1, 0.5), (None, 0.5)]
    assert new_volumes.by_element(1) == [(None, 1.0)]
    assert new_volumes.by_element(2) == [(2, 0.5), (1, 0.5)]
    assert new_volumes.by_element(3) == [(2, 1.0)]


def test_mesh_material_volumes_boundary_conditions(sphere_model):
    """Test the material volumes method using a regular mesh
    that overlaps with a vacuum boundary condition."""

    mesh = openmc.SphericalMesh.from_domain(sphere_model.geometry, dimension=(1, 1, 1))
    # extend mesh beyond the outer sphere surface to test rays crossing the boundary condition
    mesh.r_grid[-1] += 5.0

    # add a new cell to the modelthat occupies the outside of the sphere
    sphere_surfaces = list(filter(lambda s: isinstance(s, openmc.Sphere),
                            sphere_model.geometry.get_all_surfaces().values()))
    outer_cell = openmc.Cell(region=+sphere_surfaces[0])
    sphere_model.geometry.root_universe.add_cell(outer_cell)

    volumes = mesh.material_volumes(sphere_model, (0, 100, 100))
    sphere_volume = 4/3*np.pi*25**3
    mats = sphere_model.materials
    expected_volumes = [(mats[0].id, 0.25*sphere_volume),
                        (mats[1].id, 0.25*sphere_volume),
                        (mats[2].id, 0.5*sphere_volume),
                        (None, 4/3*np.pi*mesh.r_grid[-1]**3 - sphere_volume)]

    for evaluated, expected in zip(volumes.by_element(0), expected_volumes):
        assert evaluated[0] == expected[0]
        assert evaluated[1] == pytest.approx(expected[1], rel=1e-2)


def test_raytrace_mesh_infinite_loop(run_in_tmpdir):
    # Create a model with one large spherical cell
    sphere = openmc.Sphere(r=100, boundary_type='vacuum')
    cell = openmc.Cell(region=-sphere)
    model = openmc.Model()
    model.geometry = openmc.Geometry([cell])

    # Create a regular mesh and associated tally
    mesh_surface = openmc.RegularMesh()
    mesh_surface.lower_left = (-30, -30, 30)
    mesh_surface.upper_right = (30, 30, 60)
    mesh_surface.dimension = (1, 1, 1)
    reg_filter = openmc.MeshSurfaceFilter(mesh_surface)
    mesh_surface_tally = openmc.Tally()
    mesh_surface_tally.filters = [reg_filter]
    mesh_surface_tally.scores = ['current']
    model.tallies = [mesh_surface_tally]

    # Define a source such that the z position is on a mesh boundary with a very
    # small directional cosine in the z direction
    polar = openmc.stats.delta_function(1.75e-7)
    azimuthal = openmc.stats.Uniform(0.0, 2.0*pi)
    model.settings.source = openmc.IndependentSource(
        angle=openmc.stats.PolarAzimuthal(polar, azimuthal)
    )
    model.settings.run_mode = 'fixed source'
    model.settings.particles = 10
    model.settings.batches =  1

    # Run the model; this should not cause an infinite loop
    model.run()


def test_filter_time_mesh(run_in_tmpdir):
    """Test combination of TimeFilter and MeshFilter"""

    # Define material
    mat = openmc.Material()
    mat.add_nuclide('Fe56', 1.0)
    mat.set_density('g/cm3', 7.8)

    # Define geometry
    surf_Z1 = openmc.XPlane(x0=-1e10, boundary_type="reflective")
    surf_Z2 = openmc.XPlane(x0=1e10, boundary_type="reflective")
    cell_F = openmc.Cell(fill=mat, region=+surf_Z1 & -surf_Z2)
    model = openmc.Model()
    model.geometry = openmc.Geometry([cell_F])

    # Define settings
    model.settings.run_mode = "fixed source"
    model.settings.particles = 1000
    model.settings.batches = 20
    model.settings.output = {"tallies": False}
    model.settings.cutoff = {"time_neutron": 1e-7}

    # Define tallies

    # Create a mesh filter that can be used in a tally
    mesh = openmc.RegularMesh()
    mesh.dimension = (21, 1, 1)
    mesh.lower_left = (-20.5, -1e10, -1e10)
    mesh.upper_right = (20.5, 1e10, 1e10)
    time_grid = np.linspace(0.0, 1e-7, 21)

    mesh_filter = openmc.MeshFilter(mesh)
    time_filter = openmc.TimeFilter(time_grid)

    # Now use the mesh filter in a tally and indicate what scores are desired
    tally1 = openmc.Tally(name="collision")
    tally1.estimator = "collision"
    tally1.filters = [time_filter, mesh_filter]
    tally1.scores = ["flux"]
    tally2 = openmc.Tally(name="tracklength")
    tally2.estimator = "tracklength"
    tally2.filters = [time_filter, mesh_filter]
    tally2.scores = ["flux"]
    model.tallies = openmc.Tallies([tally1, tally2])

    # Run and post-process
    model.run(apply_tally_results=True)

    # Get radial flux distribution
    flux_collision = tally1.mean.ravel()
    flux_collision_unc = tally1.std_dev.ravel()
    flux_tracklength = tally2.mean.ravel()
    flux_tracklength_unc = tally2.std_dev.ravel()

    # Construct arrays with uncertainties
    collision = uarray(flux_collision, flux_collision_unc)
    tracklength = uarray(flux_tracklength, flux_tracklength_unc)
    delta = collision - tracklength

    # Compute differences and standard deviations
    diff = nominal_values(delta)
    std_dev = std_devs(delta)

    # Exclude zero-uncertainty bins
    mask = std_dev > 0.0
    dof = int(np.sum(mask))

    # Global chi-square consistency test between collision and tracklength
    # estimators. Target false positive rate ~1e-4 (1 in 10,000)
    z = diff[mask] / std_dev[mask]
    chi2_stat = np.sum(z * z)
    alpha = 1.0e-4
    crit = chi2.ppf(1 - alpha, dof)
    assert chi2_stat < crit, (
        f"Collision vs tracklength tallies disagree: chi2={chi2_stat:.2f} "
        f">= {crit=:.2f} ({dof=}, {alpha=})"
    )