import openmc
import openmc.deplete
import math
import matplotlib.pyplot as plt

# chain and cross section paths have been set on the docker image but you may want to change them
#openmc.config['chain_file']=/home/j/openmc_data/chain-endf-b8.0.xml
#openmc.config['cross_sections']=/home/j/nndc-b8.0-hdf5/endfb-viii.0-hdf5/cross_sections.xml

# Creates a simple material
my_material = openmc.Material() 
my_material.add_element('Ag', 1, percent_type='ao')
my_material.set_density('g/cm3', 10.49)

# As we are doing a depletion simulation we must set the material volume and the .depletion to True
sphere_radius = 100
volume_of_sphere = (4/3) * math.pi * math.pow(sphere_radius, 3)
my_material.volume = volume_of_sphere  # a volume is needed so openmc can find the number of atoms in the cell/material
my_material.depletable = True  # depletable = True is needed to tell openmc to update the material with each time step

materials = openmc.Materials([my_material])

# makes a simple sphere surface and cell
sph1 = openmc.Sphere(r=sphere_radius, boundary_type='vacuum')
shield_cell = openmc.Cell(region=-sph1)
shield_cell.fill = my_material
geometry = openmc.Geometry([shield_cell])

# creates a simple point source
source = openmc.IndependentSource()
source.space = openmc.stats.Point((0, 0, 0))
source.angle = openmc.stats.Isotropic()
source.energy = openmc.stats.Discrete([14e6], [1])
source.particles = 'neutron'

settings = openmc.Settings()
settings.batches = 10
settings.inactive = 0
settings.particles = 1000
settings.source = source
settings.run_mode = 'fixed source'

mesh = openmc.RegularMesh.from_domain(shield_cell, [100,100,100])

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

# this does perform transport but just to get the flux and micro xs
flux_in_each_group, micro_xs = openmc.deplete.get_microxs_and_flux(
    model=model,
    domains=mesh,
    energies='CCFE-709',
)

# constructing the operator, note we pass in the flux and micro xs
operator = openmc.deplete.IndependentOperator(
    materials=materials,
    fluxes=flux_in_each_group,
    micros=micro_xs,
    reduce_chain=True,  # reduced to only the isotopes present in depletable materials and their possible progeny
    reduce_chain_level=5,
    normalization_mode="source-rate"
)

# We define timesteps together with the source rate to make it clearer
timesteps_and_source_rates = [
    (24, 1e20),
    (24, 1e20),
    (24, 1e20),
    (24, 1e20),
    (24, 1e20),  # should saturate Ag110 here as it has been irradiated for over 5 halflives
    (24, 1e20),
    (24, 1e20),
    (24, 1e20),
    (24, 1e20),
    (24, 0),
    (24, 0),
    (24, 0),
    (24, 0),
    (24, 0),
    (24, 0),
    (24, 0),
    (24, 0),
    (24, 0),
    (24, 0),
    (24, 0),
]

# Uses list Python comprehension to get the timesteps and source_rates separately
timesteps = [item[0] for item in timesteps_and_source_rates]
source_rates = [item[1] for item in timesteps_and_source_rates]

# construct the integrator
integrator = openmc.deplete.PredictorIntegrator(
    operator=operator,
    timesteps=timesteps,
    source_rates=source_rates,
    timestep_units='s'
)

# this runs the depletion calculations for the timesteps
integrator.integrate()

# Loads up the results
results = openmc.deplete.ResultsList.from_hdf5("depletion_results.h5")

# make the source
for i, j, k in mesh.indices:
    ijk = (i-1, j-1, k-1)
    # find the materials in the voxel 
    # find the volume fraction of these materials in the voxel
    energy = materials.decay_photon_energy
    strength = energy.integral() * volume fraction
    # could be multiple materials
    sources[ijk] = openmc.IndependentSource(
        energy=energy,
        angle=angle,
        strength=strength,
        particle="photon",
        domains=[material]
    )

mesh_source = openmc.MeshSource(mesh, sources)
model.settings.source = mesh_source

# then do a gamma simulation with a dose tally