import numpy as np
import openmc
import pytest


def test_reg_mesh_from_cell():
    """Tests a RegularMesh can be made from a Cell and the specified dimensions
    are propagated through. Cell is not centralized"""
    surface = openmc.Sphere(r=10, x0=2, y0=3, z0=5)
    cell = openmc.Cell(region=-surface)

    mesh = openmc.RegularMesh.from_domain(domain=cell, dimension=[7, 11, 13])
    assert isinstance(mesh, openmc.RegularMesh)
    assert np.array_equal(mesh.dimension, (7, 11, 13))
    assert np.array_equal(mesh.lower_left, cell.bounding_box[0])
    assert np.array_equal(mesh.upper_right, cell.bounding_box[1])


def test_reg_mesh_from_bounding_box():
    """Tests a RegularMesh can be made from a BoundingBox directly."""
    bb = openmc.BoundingBox([-8, -7, -5], [12, 13, 15])

    mesh = openmc.RegularMesh.from_domain(domain=bb, dimension=[7, 11, 13])
    assert isinstance(mesh, openmc.RegularMesh)
    assert np.array_equal(mesh.dimension, (7, 11, 13))
    assert np.array_equal(mesh.lower_left, bb[0])
    assert np.array_equal(mesh.upper_right, bb[1])


def test_rectilinear_mesh_from_bounding_box():
    """Tests a RectilinearMesh can be made from a BoundingBox directly."""
    bb = openmc.BoundingBox([-8, -7, -5], [12, 13, 15])

    mesh = openmc.RectilinearMesh.from_bounding_box(bb, dimension=[2, 4, 5])
    assert isinstance(mesh, openmc.RectilinearMesh)
    assert np.array_equal(mesh.dimension, (2, 4, 5))
    assert np.array_equal(mesh.lower_left, bb[0])
    assert np.array_equal(mesh.upper_right, bb[1])
    assert np.array_equal(mesh.x_grid, [-8., 2., 12.])
    assert np.array_equal(mesh.y_grid, [-7., -2., 3., 8., 13.])
    assert np.array_equal(mesh.z_grid, [-5., -1., 3., 7., 11., 15.])


def test_cylindrical_mesh_from_cell():
    """Tests a CylindricalMesh can be made from a Cell and the specified
    dimensions are propagated through."""
    # Cell is not centralized on Z axis
    cy_surface = openmc.ZCylinder(r=50)
    z_surface_1 = openmc.ZPlane(z0=40)
    z_surface_2 = openmc.ZPlane(z0=10)
    cell = openmc.Cell(region=-cy_surface & -z_surface_1 & +z_surface_2)
    mesh = openmc.CylindricalMesh.from_domain(domain=cell, dimension=[2, 4, 3])

    assert isinstance(mesh, openmc.CylindricalMesh)
    assert np.array_equal(mesh.dimension, (2, 4, 3))
    assert np.array_equal(mesh.r_grid, [0., 25., 50.])
    assert np.array_equal(mesh.phi_grid, [0., 0.5*np.pi, np.pi, 1.5*np.pi, 2.*np.pi])
    assert np.array_equal(mesh.z_grid, [0., 10., 20., 30.])
    assert np.array_equal(mesh.origin, [0., 0., 10.])

    # Cell is not centralized on Z or X axis
    cy_surface = openmc.ZCylinder(r=50, x0=100)
    cell = openmc.Cell(region=-cy_surface & -z_surface_1 & +z_surface_2)
    mesh = openmc.CylindricalMesh.from_domain(domain=cell, dimension=[1, 1, 1])

    assert isinstance(mesh, openmc.CylindricalMesh)
    assert np.array_equal(mesh.dimension, (1, 1, 1))
    assert np.array_equal(mesh.r_grid, [0., 50.])
    assert np.array_equal(mesh.origin, [100., 0., 10.])

    # Cell is not centralized on Z, X or Y axis
    cy_surface = openmc.ZCylinder(r=50, x0=100, y0=170)
    cell = openmc.Cell(region=-cy_surface & -z_surface_1 & +z_surface_2)
    mesh = openmc.CylindricalMesh.from_domain(domain=cell, dimension=[1, 1, 1])

    assert isinstance(mesh, openmc.CylindricalMesh)
    assert np.array_equal(mesh.r_grid, [0., 50.])
    assert np.array_equal(mesh.origin, [100., 170., 10.])


def test_reg_mesh_from_region():
    """Tests a RegularMesh can be made from a Region and the default dimensions
    are propagated through. Region is not centralized"""
    surface = openmc.Sphere(r=1, x0=-5, y0=-3, z0=-2)
    region = -surface

    mesh = openmc.RegularMesh.from_domain(domain=region)
    assert isinstance(mesh, openmc.RegularMesh)
    assert np.array_equal(mesh.dimension, (10, 10, 10))  # default values
    assert np.array_equal(mesh.lower_left, region.bounding_box[0])
    assert np.array_equal(mesh.upper_right, region.bounding_box[1])


def test_cylindrical_mesh_from_region():
    """Tests a CylindricalMesh can be made from a Region and the specified
    dimensions and phi_grid_bounds are propagated through. Cell is centralized"""
    cy_surface = openmc.ZCylinder(r=6)
    z_surface_1 = openmc.ZPlane(z0=30)
    z_surface_2 = openmc.ZPlane(z0=-30)
    cell = openmc.Cell(region=-cy_surface & -z_surface_1 & +z_surface_2)
    mesh = openmc.CylindricalMesh.from_domain(
        domain=cell,
        dimension=(6, 2, 3),
        phi_grid_bounds=(0., np.pi)
    )

    assert isinstance(mesh, openmc.CylindricalMesh)
    assert np.array_equal(mesh.dimension, (6, 2, 3))
    assert np.array_equal(mesh.r_grid, [0., 1., 2., 3., 4., 5., 6.])
    assert np.array_equal(mesh.phi_grid, [0., 0.5*np.pi, np.pi])
    assert np.array_equal(mesh.z_grid, [0.0, 20., 40., 60])
    assert np.array_equal(mesh.origin, (0.0, 0.0, -30.))


def test_spherical_mesh_from_domain():
    """Tests a SphericalMesh can be made from a Region and the specified
    dimensions are propagated through. Cell is not centralized"""
    sphere = openmc.Sphere(r=5, x0=2, y0=3, z0=4)
    region = -sphere

    geometry = openmc.Geometry(openmc.Universe(cells=[openmc.Cell(region=region)]))

    region_mesh = openmc.SphericalMesh.from_domain(
        domain=region, dimension=(4, 3, 4))
    universe_mesh = openmc.SphericalMesh.from_domain(
        domain=geometry.root_universe, dimension=(4, 3, 4))
    geometry_mesh = openmc.SphericalMesh.from_domain(
        domain=geometry, dimension=(4, 3, 4))


    for mesh in (region_mesh, universe_mesh, geometry_mesh):
        assert isinstance(mesh, openmc.SphericalMesh)
        assert np.array_equal(mesh.dimension, (4, 3, 4))
        assert np.array_equal(mesh.r_grid, [0., 1.25, 2.5, 3.75, 5.0])
        assert np.array_equal(mesh.theta_grid, [0., np.pi/3., 2*np.pi/3., np.pi])
        assert np.array_equal(mesh.phi_grid, [0., np.pi/2., np.pi, 3*np.pi/2., 2*np.pi])
        assert np.array_equal(mesh.origin, (2.0, 3.0, 4.0))

        for p in mesh.centroids.reshape(-1, 3):
            assert p in mesh.bounding_box


def test_reg_mesh_from_universe():
    """Tests a RegularMesh can be made from a Universe and the default
    dimensions are propagated through. Universe is centralized"""
    surface = openmc.Sphere(r=42)
    cell = openmc.Cell(region=-surface)
    universe = openmc.Universe(cells=[cell])

    mesh = openmc.RegularMesh.from_domain(universe)
    assert isinstance(mesh, openmc.RegularMesh)
    assert np.array_equal(mesh.dimension, (10, 10, 10))  # default values
    assert np.array_equal(mesh.lower_left, universe.bounding_box[0])
    assert np.array_equal(mesh.upper_right, universe.bounding_box[1])


def test_reg_mesh_from_geometry():
    """Tests a RegularMesh can be made from a Geometry and the default
    dimensions are propagated through. Geometry is centralized"""
    surface = openmc.Sphere(r=42)
    cell = openmc.Cell(region=-surface)
    universe = openmc.Universe(cells=[cell])
    geometry = openmc.Geometry(universe)

    mesh = openmc.RegularMesh.from_domain(geometry)
    assert isinstance(mesh, openmc.RegularMesh)
    assert np.array_equal(mesh.dimension, (10, 10, 10))  # default values
    assert np.array_equal(mesh.lower_left, geometry.bounding_box[0])
    assert np.array_equal(mesh.upper_right, geometry.bounding_box[1])


def test_error_from_unsupported_object():
    with pytest.raises(TypeError):
        openmc.RegularMesh.from_domain("vacuum energy")


def test_regularmesh_from_domain_error_from_small_dimensions():
    surface = openmc.Sphere(r=20)
    cell = openmc.Cell(region=-surface)
    with pytest.raises(
        ValueError, match='Unable to set "dimension" to "-2" since it is less than "1"'
    ):
        openmc.RegularMesh.from_domain(domain=cell, dimension=-2)


def test_dimensions_from_domain_dimensions_from_int():
    region = openmc.model.RectangularParallelepiped(
        xmin=-100,
        xmax=150,
        ymin=-50,
        ymax=200,
        zmin=300,
        zmax=400,
        boundary_type="vacuum",
    )
    cell = openmc.Cell(region=-region)
    mesh = openmc.RegularMesh.from_domain(domain=cell, dimension=1000)
    assert mesh.dimension == (14, 14, 5)