import lxml.etree as ET
import numpy as np
from uncertainties import ufloat
import openmc
import pytest

from tests.unit_tests import assert_unbounded
from openmc.data import atomic_mass, AVOGADRO


def test_contains():
    # Cell with specified region
    s = openmc.XPlane()
    c = openmc.Cell(region=+s)
    assert (1.0, 0.0, 0.0) in c
    assert (-1.0, 0.0, 0.0) not in c

    # Cell with no region
    c = openmc.Cell()
    assert (10.0, -4., 2.0) in c


def test_id():
    openmc.Cell(cell_id=0)
    with pytest.raises(ValueError):
        openmc.Cell(cell_id=-1)


def test_repr(cell_with_lattice):
    cells, mats, univ, lattice = cell_with_lattice
    repr(cells[0])  # cell with distributed materials
    repr(cells[1])  # cell with material
    repr(cells[2])  # cell with lattice

    # Empty cell
    c = openmc.Cell()
    repr(c)

    # Empty cell with volume
    c.volume = 3.0
    repr(c)

    # Empty cell with uncertain volume
    c.volume = ufloat(3.0, 0.2)
    repr(c)


def test_bounding_box():
    zcyl = openmc.ZCylinder()
    c = openmc.Cell(region=-zcyl)
    ll, ur = c.bounding_box
    assert ll == pytest.approx((-1., -1., -np.inf))
    assert ur == pytest.approx((1., 1., np.inf))

    # Cell with no region specified
    c = openmc.Cell()
    assert_unbounded(c)


def test_clone():
    m = openmc.Material()
    cyl = openmc.ZCylinder()
    c = openmc.Cell(fill=m, region=-cyl)

    # Check cloning with all optional params as the defaults
    c2 = c.clone()
    assert c2.id != c.id
    assert c2.fill != c.fill
    assert c2.region != c.region

    c3 = c.clone(clone_materials=False)
    assert c3.id != c.id
    assert c3.fill == c.fill
    assert c3.region != c.region

    c4 = c.clone(clone_regions=False)
    assert c4.id != c.id
    assert c4.fill != c.fill
    assert c4.region == c.region

    # Add optional properties to the original cell to ensure they're cloned successfully
    c.temperature = 650.
    c.translation = (1., 2., 3.)
    c.rotation = (4., 5., 6.)
    c.volume = 100

    c5 = c.clone(clone_materials=False, clone_regions=False)
    assert c5.id != c.id
    assert c5.fill == c.fill
    assert c5.region == c.region
    assert c5.temperature == c.temperature
    assert c5.volume == c.volume
    assert all(c5.translation == c.translation)
    assert all(c5.rotation == c.rotation)

    # Mutate the original to ensure the changes are not seen in the clones
    c.fill = openmc.Material()
    c.region = +openmc.ZCylinder()
    c.translation = (-1., -2., -3.)
    c.rotation = (-4., -5., -6.)
    c.temperature = 1
    c.volume = 1
    assert c5.fill != c.fill
    assert c5.region != c.region
    assert c5.temperature != c.temperature
    assert c5.volume != c.volume
    assert all(c5.translation != c.translation)
    assert all(c5.rotation != c.rotation)


def test_temperature(cell_with_lattice):
    # Make sure temperature propagates through universes
    m = openmc.Material()
    s = openmc.XPlane()
    c1 = openmc.Cell(fill=m, region=+s)
    c2 = openmc.Cell(fill=m, region=-s)
    u1 = openmc.Universe(cells=[c1, c2])
    c = openmc.Cell(fill=u1)

    c.temperature = 400.0
    assert c1.temperature == 400.0
    assert c2.temperature == 400.0
    with pytest.raises(ValueError):
        c.temperature = -100.
    c.temperature = None
    assert c1.temperature == None
    assert c2.temperature == None

    # distributed temperature
    cells, _, _, _ = cell_with_lattice
    c = cells[0]
    c.temperature = (300., 600., 900.)


def test_densities(cell_with_lattice):
    # Make sure density propagates through universes
    m = openmc.Material()
    s = openmc.XPlane()
    c1 = openmc.Cell(fill=m, region=+s)
    c2 = openmc.Cell(fill=m, region=-s)
    u1 = openmc.Universe(cells=[c1, c2])
    c = openmc.Cell(fill=u1)

    c.density = 1.
    assert c1.density == 1.
    assert c2.density == 1.
    with pytest.raises(ValueError):
        c.density = -1.
    c.density = None
    assert c1.density == None
    assert c2.density == None

    # distributed density
    cells, _, _, _ = cell_with_lattice
    c = cells[0]
    c.density = (1., 2., 3.)

def test_rotation():
    u = openmc.Universe()
    c = openmc.Cell(fill=u)
    c.rotation = (180.0, 0.0, 0.0)
    assert np.allclose(c.rotation_matrix, [
        [1., 0., 0.],
        [0., -1., 0.],
        [0., 0., -1.]
    ])

    c.rotation = (0.0, 90.0, 0.0)
    assert np.allclose(c.rotation_matrix, [
        [0., 0., -1.],
        [0., 1., 0.],
        [1., 0., 0.]
    ])


def test_get_nuclides(uo2):
    c = openmc.Cell(fill=uo2)
    nucs = c.get_nuclides()
    assert nucs == ['U235', 'O16']


def test_volume_setting():
    c = openmc.Cell()

    # Test ordinary volume and uncertain volume
    c.volume = 3
    c.volume = ufloat(3, 0.7)

    # Allow volume to be set to 0.0
    c.volume = 0.0
    c.volume = ufloat(0.0, 0.1)

    # Test errors for negative volume
    with pytest.raises(ValueError):
        c.volume = -1.0
    with pytest.raises(ValueError):
        c.volume = ufloat(-0.05, 0.1)


def test_atoms_material_cell(uo2, water):
    """ Test if correct number of atoms is returned.
    Also check if Cell.atoms still works after volume/material was changed
    """
    c = openmc.Cell(fill=uo2)
    c.volume = 2.0
    expected_nucs = ['U235', 'O16']

    # Precalculate the expected number of atoms
    M = (atomic_mass('U235') + 2 * atomic_mass('O16')) / 3
    expected_atoms = [
        1/3 * uo2.density/M * AVOGADRO * 2.0,  # U235
        2/3 * uo2.density/M * AVOGADRO * 2.0   # O16
    ]

    tuples = c.atoms.items()
    for nuc, atom_num, t in zip(expected_nucs, expected_atoms, tuples):
        assert nuc == t[0]
        assert atom_num == pytest.approx(t[1])

    # Change volume and check if OK
    c.volume = 3.0
    expected_atoms = [
        1/3 * uo2.density/M * AVOGADRO * 3.0,  # U235
        2/3 * uo2.density/M * AVOGADRO * 3.0   # O16
    ]

    tuples = c.atoms.items()
    for nuc, atom_num, t in zip(expected_nucs, expected_atoms, tuples):
        assert nuc == t[0]
        assert atom_num == pytest.approx(t[1])

    # Change material and check if OK
    c.fill = water
    expected_nucs = ['H1', 'O16']
    M = (2 * atomic_mass('H1') + atomic_mass('O16')) / 3
    expected_atoms = [
        2/3 * water.density/M * AVOGADRO * 3.0,  # H1
        1/3 * water.density/M * AVOGADRO * 3.0   # O16
    ]

    tuples = c.atoms.items()
    for nuc, atom_num, t in zip(expected_nucs, expected_atoms, tuples):
        assert nuc == t[0]
        assert atom_num == pytest.approx(t[1])


def test_atoms_distribmat_cell(uo2, water):
    """ Test if correct number of atoms is returned for a cell with
    'distribmat' fill
    """
    c = openmc.Cell(fill=[uo2, water])
    c.volume = 6.0

    # Calculate the expected number of atoms
    expected_nucs = ['U235', 'O16', 'H1']
    M_uo2 = (atomic_mass('U235') + 2 * atomic_mass('O16')) / 3
    M_water = (2 * atomic_mass('H1') + atomic_mass('O16')) / 3
    expected_atoms = [
        1/3 * uo2.density/M_uo2 * AVOGADRO * 3.0,        # U235
        (2/3 * uo2.density/M_uo2 * AVOGADRO * 3.0 +
         1/3 * water.density/M_water * AVOGADRO * 3.0),  # O16
        2/3 * water.density/M_water * AVOGADRO * 3.0     # H1
    ]

    tuples = c.atoms.items()
    for nuc, atom_num, t in zip(expected_nucs, expected_atoms, tuples):
        assert nuc == t[0]
        assert atom_num == pytest.approx(t[1])


def test_atoms_errors(cell_with_lattice):
    cells, mats, univ, lattice = cell_with_lattice

    # Material Cell with no volume
    with pytest.raises(ValueError):
        cells[1].atoms

    # Cell with lattice
    cells[2].volume = 3
    with pytest.raises(ValueError):
        cells[2].atoms

    # Cell with volume but with void fill
    cells[1].volume = 2
    cells[1].fill = None
    with pytest.raises(ValueError):
        cells[1].atoms


def test_nuclide_densities(uo2):
    c = openmc.Cell(fill=uo2)
    expected_nucs = ['U235', 'O16']
    expected_density = [1.0, 2.0]
    tuples = list(c.get_nuclide_densities().values())
    for nuc, density, t in zip(expected_nucs, expected_density, tuples):
        assert nuc == t[0]
        assert density == pytest.approx(t[1])

    # Empty cell
    c = openmc.Cell()
    assert not c.get_nuclide_densities()


def test_get_all_universes(cell_with_lattice):
    # Cell with nested universes
    c1 = openmc.Cell()
    u1 = openmc.Universe(cells=[c1])
    c2 = openmc.Cell(fill=u1)
    u2 = openmc.Universe(cells=[c2])
    c3 = openmc.Cell(fill=u2)
    univs = set(c3.get_all_universes().values())
    assert not (univs ^ {u1, u2})

    # Cell with lattice
    cells, mats, univ, lattice = cell_with_lattice
    univs = set(cells[-1].get_all_universes().values())
    assert not (univs ^ {univ})


def test_get_all_materials(cell_with_lattice):
    # Normal cell
    m = openmc.Material()
    c = openmc.Cell(fill=m)
    test_mats = set(c.get_all_materials().values())
    assert not(test_mats ^ {m})

    # Cell filled with distributed materials
    cells, mats, univ, lattice = cell_with_lattice
    c = cells[0]
    test_mats = set(c.get_all_materials().values())
    assert not (test_mats ^ set(m for m in c.fill if m is not None))

    # Cell filled with universe
    c = cells[-1]
    test_mats = set(c.get_all_materials().values())
    assert not (test_mats ^ set(mats))


def test_to_xml_element(cell_with_lattice):
    cells, mats, univ, lattice = cell_with_lattice

    c = cells[-1]
    root = ET.Element('geometry')
    elem = c.create_xml_subelement(root)
    assert elem.tag == 'cell'
    assert elem.get('id') == str(c.id)
    assert elem.get('region') is None
    surf_elem = root.find('surface')
    assert surf_elem.get('id') == str(cells[0].region.surface.id)

    c = cells[0]
    c.temperature = 900.0
    c.volume = 1.0
    elem = c.create_xml_subelement(root)
    assert elem.get('region') == str(c.region)
    assert elem.get('temperature') == str(c.temperature)
    assert elem.get('volume') == str(c.volume)


@pytest.mark.parametrize("rotation", [
    (90, 45, 0),
    [[1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, -1.0, 0.0]]
])
def test_rotation_from_xml(rotation):
    # Make sure rotation attribute (matrix) round trips through XML correctly
    s = openmc.ZCylinder(r=10.0)
    cell = openmc.Cell(region=-s)
    cell.rotation = rotation
    root = ET.Element('geometry')
    elem = cell.create_xml_subelement(root)
    new_cell = openmc.Cell.from_xml_element(
        elem, {s.id: s}, {'void': None}, openmc.Universe
    )
    np.testing.assert_allclose(new_cell.rotation, cell.rotation)


def test_dagmccell_from_xml_element():
    """DAGMCCell.from_xml_element parses material, temperature, density,
    and volume; rejects unsupported attributes."""
    mat = openmc.Material(1)
    mat.add_nuclide('U235', 1.0)
    mats = {'1': mat}
    # In practice, from_xml_element is always called from DAGMCUniverse
    # during XML parsing. A placeholder universe is used here so the test
    # can exercise the method directly without a real DAGMC model file.
    placeholder_univ = openmc.DAGMCUniverse('model.h5m')

    # material + temperature + density round-trip
    xml = '<cell id="5" name="fuel" material="1" temperature="900.0" density="10.5"/>'
    cell = openmc.DAGMCCell.from_xml_element(ET.fromstring(xml), mats, placeholder_univ)
    assert cell.id == 5
    assert cell.name == 'fuel'
    assert cell.fill is mat
    assert cell.density == 10.5
    assert cell.temperature == 900.0

    # volume round-trip
    xml = '<cell id="6" material="1" volume="42.0"/>'
    cell = openmc.DAGMCCell.from_xml_element(ET.fromstring(xml), mats, placeholder_univ)
    assert cell.volume == 42.0

    # forbidden: region, fill, universe
    for tag, val in [('region', '-1'), ('fill', '2'), ('universe', '0')]:
        xml = f'<cell id="7" material="1" {tag}="{val}"/>'
        with pytest.raises(ValueError, match=tag):
            openmc.DAGMCCell.from_xml_element(ET.fromstring(xml), mats, placeholder_univ)

    # forbidden: translation, rotation
    for tag, val in [('translation', '1 0 0'), ('rotation', '0 0 90')]:
        xml = f'<cell id="7" material="1" {tag}="{val}"/>'
        with pytest.raises(ValueError, match=tag):
            openmc.DAGMCCell.from_xml_element(ET.fromstring(xml), mats, placeholder_univ)

    # missing material raises
    xml = '<cell id="8"/>'
    with pytest.raises(ValueError, match='material'):
        openmc.DAGMCCell.from_xml_element(ET.fromstring(xml), mats, placeholder_univ)


def test_plot(run_in_tmpdir):
    zcyl = openmc.ZCylinder()
    c = openmc.Cell(region=-zcyl)

    # create a universe before the plot
    u_before = openmc.Universe()

    # create a plot of the cell
    c.plot()

    # create a universe after the plot
    u_after = openmc.Universe()

    # ensure that calling the plot method doesn't
    # affect the universe ID space
    assert u_before.id + 1 == u_after.id