""" ExternalRates depletion test suite """

from pathlib import Path
import shutil

import numpy as np
import pytest
import openmc
import openmc.deplete
from openmc.deplete import CoupledOperator

from tests.regression_tests import config, assert_reaction_rates_equal, \
    assert_atoms_equal

@pytest.fixture
def model():
    openmc.reset_auto_ids()
    f = openmc.Material(name="f")
    f.add_element("U", 1, percent_type="ao", enrichment=4.25)
    f.add_element("O", 2)
    f.set_density("g/cc", 10.4)

    w = openmc.Material(name="w")
    w.add_element("O", 1)
    w.add_element("H", 2)
    w.set_density("g/cc", 1.0)
    w.depletable = True

    radii = [0.42, 0.45]
    f.volume = np.pi * radii[0] ** 2
    w.volume = np.pi * (radii[1]**2 - radii[0]**2)

    materials = openmc.Materials([f, w])

    surf_f = openmc.Sphere(r=radii[0])
    surf_w = openmc.Sphere(r=radii[1], boundary_type='reflective')
    cell_f = openmc.Cell(fill=f, region=-surf_f)
    cell_w = openmc.Cell(fill=w, region=+surf_f & -surf_w)
    geometry = openmc.Geometry([cell_f, cell_w])

    settings = openmc.Settings()
    settings.particles = 150
    settings.inactive = 0
    settings.batches = 10

    return openmc.Model(geometry, materials, settings)


@pytest.mark.parametrize("rate, dest_mat, power, ref_result", [
    (1e-5, None, 0.0, 'no_depletion_only_removal'),
    (-1e-5, None, 0.0, 'no_depletion_only_feed'),
    (1e-5, None, 174.0, 'depletion_with_removal'),
    (-1e-5, None, 174.0, 'depletion_with_feed'),
    (-1e-5, 'w', 0.0, 'no_depletion_with_transfer'),
    (1e-5, 'w', 174.0, 'depletion_with_transfer'),
    (0.0, None, 174.0, 'depletion_with_redox'),
    (1e-5, None, 174.0, 'depletion_with_removal_and_redox'),
    (1e-5, 'w', 174.0, 'depletion_with_transfer_and_redox'),
    ])
def test_transfer_rates(run_in_tmpdir, model, rate, dest_mat, power, ref_result):
    """Tests transfer_rates depletion class with transfer rates"""

    chain_file = Path(__file__).parents[2] / 'chain_simple.xml'

    transfer_elements = ['Xe']
    os = {'I': -1, 'Xe':0, 'Cs': 1, 'Gd': 3, 'U': 4}

    op = CoupledOperator(model, chain_file)
    op.round_number = True
    integrator = openmc.deplete.PredictorIntegrator(
        op, [1], power, timestep_units = 'd')
    if rate != 0.0:
        integrator.add_transfer_rate('f', transfer_elements, rate,
                                    destination_material=dest_mat)
    if 'redox' in ref_result.split('_'):
        integrator.add_redox('f', {'Gd157':1}, os)

    integrator.integrate()

    # Get path to test and reference results
    path_test = op.output_dir / 'depletion_results.h5'
    path_reference = Path(__file__).with_name(f'ref_{ref_result}.h5')

    # If updating results, do so and return
    if config['update']:
        shutil.copyfile(str(path_test), str(path_reference))
        return

    # Load the reference/test results
    res_ref = openmc.deplete.Results(path_reference)
    res_test = openmc.deplete.Results(path_test)

    assert_atoms_equal(res_ref, res_test, tol=1e-3)
    assert_reaction_rates_equal(res_ref, res_test, tol=1e-3)

@pytest.mark.parametrize("rate, power, ref_result", [
    (1e-1, 0.0, 'no_depletion_with_ext_source'),
    (1e-1, 174., 'depletion_with_ext_source'),
])
def test_external_source_rates(run_in_tmpdir, model, rate, power, ref_result):
    """Tests external_rates depletion class with external source rates"""

    chain_file = Path(__file__).parents[2] / 'chain_simple.xml'

    external_source_vector = {'U': 1}

    op = CoupledOperator(model, chain_file)
    op.round_number = True
    integrator = openmc.deplete.PredictorIntegrator(
        op, [1], power, timestep_units='d')
    integrator.add_external_source_rate('f', external_source_vector, rate)
    integrator.integrate()

    # Get path to test and reference results
    path_test = op.output_dir / 'depletion_results.h5'
    path_reference = Path(__file__).with_name(f'ref_{ref_result}.h5')

    # If updating results, do so and return
    if config['update']:
        shutil.copyfile(str(path_test), str(path_reference))
        return

    # Load the reference/test results
    res_ref = openmc.deplete.Results(path_reference)
    res_test = openmc.deplete.Results(path_test)

    assert_atoms_equal(res_ref, res_test, tol=1e-3)
    assert_reaction_rates_equal(res_ref, res_test, tol=1e-3)

@pytest.mark.parametrize("external_source_rate, transfer_rate, power, ref_result", [
    (1e-1, 1e-1, 174., 'depletion_with_ext_source_and_transfer'),
])
def test_external_source_rates_with_transfer_rates(run_in_tmpdir, model, external_source_rate, 
    transfer_rate, power, ref_result):
    """Tests external_rates depletion class with external source rates and transfer rates"""

    chain_file = Path(__file__).parents[2] / 'chain_simple.xml'

    external_source_vector = {'U': 1}

    op = CoupledOperator(model, chain_file)
    op.round_number = True
    integrator = openmc.deplete.PredictorIntegrator(
        op, [1], power, timestep_units='d')
    integrator.add_external_source_rate('f', external_source_vector, external_source_rate)
    integrator.add_transfer_rate('f', ['U235'], transfer_rate, destination_material='w')
    integrator.integrate()

    # Get path to test and reference results
    path_test = op.output_dir / 'depletion_results.h5'
    path_reference = Path(__file__).with_name(f'ref_{ref_result}.h5')

    # If updating results, do so and return
    if config['update']:
        shutil.copyfile(str(path_test), str(path_reference))
        return

    # Load the reference/test results
    res_ref = openmc.deplete.Results(path_reference)
    res_test = openmc.deplete.Results(path_test)

    assert_atoms_equal(res_ref, res_test, tol=1e-3)
    assert_reaction_rates_equal(res_ref, res_test, tol=1e-3)