from collections.abc import Mapping
import os

import numpy as np
import pytest
import openmc
import openmc.exceptions as exc
import openmc.capi

from tests import cdtemp


@pytest.fixture(scope='module')
def pincell_model():
    """Set up a model to test with and delete files when done"""
    openmc.reset_auto_ids()
    pincell = openmc.examples.pwr_pin_cell()
    pincell.settings.verbosity = 1

    # Add a tally
    filter1 = openmc.MaterialFilter(pincell.materials)
    filter2 = openmc.EnergyFilter([0.0, 1.0, 1.0e3, 20.0e6])
    mat_tally = openmc.Tally()
    mat_tally.filters = [filter1, filter2]
    mat_tally.nuclides = ['U235', 'U238']
    mat_tally.scores = ['total', 'elastic', '(n,gamma)']
    pincell.tallies.append(mat_tally)

    # Add an expansion tally
    zernike_tally = openmc.Tally()
    filter3 = openmc.ZernikeFilter(5, r=.63)
    cells = pincell.geometry.root_universe.cells
    filter4 = openmc.CellFilter(list(cells.values()))
    zernike_tally.filters = [filter3, filter4]
    zernike_tally.scores = ['fission']
    pincell.tallies.append(zernike_tally)

    # Write XML files in tmpdir
    with cdtemp():
        pincell.export_to_xml()
        yield


@pytest.fixture(scope='module')
def capi_init(pincell_model):
    openmc.capi.init()
    yield
    openmc.capi.finalize()


@pytest.fixture(scope='module')
def capi_simulation_init(capi_init):
    openmc.capi.simulation_init()
    yield


@pytest.fixture(scope='module')
def capi_run(capi_simulation_init):
    openmc.capi.run()


def test_cell_mapping(capi_init):
    cells = openmc.capi.cells
    assert isinstance(cells, Mapping)
    assert len(cells) == 3
    for cell_id, cell in cells.items():
        assert isinstance(cell, openmc.capi.Cell)
        assert cell_id == cell.id


def test_cell(capi_init):
    cell = openmc.capi.cells[1]
    assert isinstance(cell.fill, openmc.capi.Material)
    cell.fill = openmc.capi.materials[1]
    assert str(cell) == 'Cell[0]'


def test_new_cell(capi_init):
    with pytest.raises(exc.AllocationError):
        openmc.capi.Cell(1)
    new_cell = openmc.capi.Cell()
    new_cell_with_id = openmc.capi.Cell(10)
    assert len(openmc.capi.cells) == 5


def test_material_mapping(capi_init):
    mats = openmc.capi.materials
    assert isinstance(mats, Mapping)
    assert len(mats) == 3
    for mat_id, mat in mats.items():
        assert isinstance(mat, openmc.capi.Material)
        assert mat_id == mat.id


def test_material(capi_init):
    m = openmc.capi.materials[3]
    assert m.nuclides == ['H1', 'O16', 'B10', 'B11']

    old_dens = m.densities
    test_dens = [1.0e-1, 2.0e-1, 2.5e-1, 1.0e-3]
    m.set_densities(m.nuclides, test_dens)
    assert m.densities == pytest.approx(test_dens)

    assert m.volume is None
    m.volume = 10.0
    assert m.volume == 10.0

    with pytest.raises(exc.InvalidArgumentError):
        m.set_density(1.0, 'goblins')

    rho = 2.25e-2
    m.set_density(rho)
    assert sum(m.densities) == pytest.approx(rho)


def test_new_material(capi_init):
    with pytest.raises(exc.AllocationError):
        openmc.capi.Material(1)
    new_mat = openmc.capi.Material()
    new_mat_with_id = openmc.capi.Material(10)
    assert len(openmc.capi.materials) == 5


def test_nuclide_mapping(capi_init):
    nucs = openmc.capi.nuclides
    assert isinstance(nucs, Mapping)
    assert len(nucs) == 12
    for name, nuc in nucs.items():
        assert isinstance(nuc, openmc.capi.Nuclide)
        assert name == nuc.name


def test_settings(capi_init):
    settings = openmc.capi.settings
    assert settings.batches == 10
    settings.batches = 10
    assert settings.inactive == 5
    assert settings.generations_per_batch == 1
    assert settings.particles == 100
    assert settings.seed == 1
    settings.seed = 11

    assert settings.run_mode == 'eigenvalue'
    settings.run_mode = 'volume'
    settings.run_mode = 'eigenvalue'


def test_tally_mapping(capi_init):
    tallies = openmc.capi.tallies
    assert isinstance(tallies, Mapping)
    assert len(tallies) == 2
    for tally_id, tally in tallies.items():
        assert isinstance(tally, openmc.capi.Tally)
        assert tally_id == tally.id


def test_tally(capi_init):
    t = openmc.capi.tallies[1]
    assert t.type == 'volume'
    assert len(t.filters) == 2
    assert isinstance(t.filters[0], openmc.capi.MaterialFilter)
    assert isinstance(t.filters[1], openmc.capi.EnergyFilter)

    # Create new filter and replace existing
    with pytest.raises(exc.AllocationError):
        openmc.capi.MaterialFilter(uid=1)
    mats = openmc.capi.materials
    f = openmc.capi.MaterialFilter([mats[2], mats[1]])