## BISON depletion capability comparison
## Rep_Na_3 assessment case, fuel pellet mesh with 55 radial divisions.
## BISON solves the depletion for regions with the same radius but different volumes.
## The number of radial divisions in the mesh controls the number of the depletion regions

try:
    from mpi4py import MPI
    rank = MPI.COMM_WORLD.Get_rank()
except ModuleNotFoundError:
    rank = 0

import math, sys
import openmc
import openmc.deplete
import numpy as np
from matplotlib import pyplot

######################################################################
#                      Operational Parameters
######################################################################

fuel_temp = 978.31
gap_temp = 663.69
water_den = 0.7051
ppm_boron = 500
water_temp = 585.25

######################################################################
#                      Geometric Parameters
######################################################################

Rfo = 0.40956		## pellet outer radius
Rci = 0.41776		## clad inner radius
Rco = 0.47736		## clad outer radius
hpitch = 0.63		## half-pitch between rods
fheight = 1.3589	## fuel height
nr = 55				## number of radial divisions

######################################################################
#                      Chain Parameters
######################################################################

#chain_path = "/storage/home/lca5209/OpenMC/chains/chain_simple.xml"
#chain_path = "/storage/home/lca5209/OpenMC/chains/chain_casl_pwr.xml"
#chain_path = "/storage/home/lca5209/OpenMC/chains/chain_endfb71_pwr.xml"
chain_path = "/opt/xdata/endfb-vii.1-hdf5/chain_endfb71_pwr.xml"

######################################################################
#                      Settings
######################################################################

batches = 100
inactive = 20
particles = 50000
seed = np.random.randint(2147483647)*2+1

######################################################################
#                      Setting the materials
######################################################################
## fuel = 3.0 wt% enriched uranium dioxide

fuel = openmc.Material(material_id=1,name="UO2")

fuel.add_nuclide("U238",0.84223593,'wo')
fuel.add_nuclide("U235",0.03918626,'wo')
fuel.add_nuclide("O16",0.11857781,'wo')
fuel.set_density("g/cc",10.394075)
fuel.temperature = fuel_temp
fuel.volume = (math.pi*Rfo**2)*fheight

mat_list = []

for i in range(nr):
    mat = fuel.clone()
    mat.volume = (math.pi*((((i+1)/55)*Rfo)**2-(((i)/55)*Rfo)**2))*fheight
    mat_list.append(mat)

## gap = helium

gap = openmc.Material(name='He')
gap.set_density('g/cc',0.001598)
gap.add_element('He', 1.0,'wo')
gap.temperature = gap_temp
mat_list.append(gap)

## clad = zircaloy

clad = openmc.Material(name='Zircaloy_4')
clad.set_density('g/cc', 6.56)
clad.add_element('Sn', 0.014  , 'wo')
clad.add_element('Fe', 0.00165, 'wo')
clad.add_element('Cr', 0.001  , 'wo')
clad.add_element('Zr', 0.98335, 'wo')
mat_list.append(clad)

## coolant = borated water

coolant = openmc.Material(name='Borated_water')
coolant.set_density('g/cc', water_den)
coolant.add_nuclide('B10',ppm_boron*0.199E-06,'wo')
coolant.add_nuclide('B11',ppm_boron*0.801E-06,'wo')
coolant.add_element('H',0.1108612,'wo')
coolant.add_element('O',0.8886388,'wo')
coolant.remove_nuclide('O18')
coolant.temperature = water_temp
coolant.add_s_alpha_beta('c_H_in_H2O')
mat_list.append(coolant)

sys.stdout.flush()
try:
    MPI.COMM_WORLD.Barrier()
except NameError:
    pass
if rank == 0:
    print(mat_list)
sys.stdout.flush()

materials = openmc.Materials(mat_list)
materials.cross_sections="/opt/xdata/endfb-vii.1-hdf5/cross_sections.xml"
materials.export_to_xml()

######################################################################
#                      Building the geometry
######################################################################

## Fuel rings, gap, and cladding

cylinder = []

counter_1 = 0
for i in range(nr):
    c = openmc.ZCylinder(surface_id=(i+1), x0=0, y0=0, r=((i+1)/55)*Rfo, name='Fuel '+ str(i+1) + ' OR')
    cylinder.append(c)
    counter_1 +=1
cylinder_g = openmc.ZCylinder(surface_id=counter_1+1, x0=0, y0=0, r=Rci, name='Gap OR')
cylinder_c = openmc.ZCylinder(surface_id=counter_1+2, x0=0, y0=0, r=Rco, name='Clad OR')
cylinder.append(cylinder_g)
cylinder.append(cylinder_c)

## Box surrounding fuel rod
left = openmc.XPlane(surface_id=counter_1+3, x0=-hpitch, name='left')
right = openmc.XPlane(surface_id=counter_1+4, x0=hpitch, name='right')
bottom = openmc.YPlane(surface_id=counter_1+5, y0=-hpitch, name='bottom')
top = openmc.YPlane(surface_id=counter_1+6, y0=hpitch, name='top')

## Surfaces above/below fuel
fsouth = openmc.ZPlane(surface_id=counter_1+7, z0=-fheight/2.0, name='fsouth')
fnorth = openmc.ZPlane(surface_id=counter_1+8, z0=fheight/2.0, name='fnorth')

## Boundary Conditions
left.boundary_type = 'reflective'
right.boundary_type = 'reflective'
top.boundary_type = 'reflective'
bottom.boundary_type = 'reflective'
fsouth.boundary_type = 'reflective'
fnorth.boundary_type = 'reflective'

cell = []

counter_2 = 0
for i in range(nr):
    f = openmc.Cell(cell_id=(i+1), name='fuel '+str(i+1))
    cell.append(f)
    counter_2 += 1
gap_cell = openmc.Cell(cell_id=counter_2+1, name='Gap cell')
clad_cell = openmc.Cell(cell_id=counter_2+2, name='Clad cell')
water_cell = openmc.Cell(cell_id=counter_2+3, name='Water cell')
cell.append(gap_cell)
cell.append(clad_cell)
cell.append(water_cell)

for i in range(len(cell)-1):
    if i==0:
        cell[i].region = -cylinder[i] & -fnorth & +fsouth
    elif i==(len(cell)-2):
        cell[i+1].region = +cylinder[i] & +left & -right & +bottom & -top & -fnorth & +fsouth
        cell[i].region = +cylinder[(i-1)] & -cylinder[i] & -fnorth & +fsouth
    else:
        cell[i].region = +cylinder[(i-1)] & -cylinder[i] & -fnorth & +fsouth

count_3 = 0
for i in range(nr):
    cell[i].fill = mat_list[i]
    count_3 +=1
cell[count_3].fill = gap
cell[count_3+1].fill = clad
cell[count_3+2].fill = coolant

## Instantiate Universe
root = openmc.Universe(universe_id=0, name='root universe')

## Register Cells with Universe
root.add_cells(cell)

## Instantiate a Geometry, register the root Universe, and export to XML
geometry = openmc.Geometry(root)
geometry.export_to_xml()

plot = openmc.Plot.from_geometry(geometry)
plot.pixels = (250, 250)
plot.filename = 'pinplot_55_regions'
plot.width = (hpitch*2, hpitch*2)
plot.pixels = (500, 500)
plot.color_by = 'cell'
plot.colors = {cell[-3]: 'purple',cell[-2]: 'grey',
               cell[-1]: 'blue'}
plots = openmc.Plots([plot])
plots.export_to_xml()
openmc.plot_geometry()

###############################################################################
#                   Exporting to OpenMC settings.xml file
###############################################################################

## Instantiate a Settings object, set all runtime parameters, and export to XML
settings = openmc.Settings()
settings.batches = batches
settings.inactive = inactive
settings.particles = particles
settings.seed = seed
settings.verbosity = 7
settings.photon_transport = False

## Cross-section temperature information
settings.temperature = {
    'method': 'interpolation',
    'tolerance': 100 ,
    'multipole': True
}

## Create an initial uniform spatial source distribution over fissionable zones
bounds = [-hpitch, -hpitch, -fheight/2, hpitch, hpitch, fheight/2]
uniform_dist = openmc.stats.Box(bounds[:3], bounds[3:], only_fissionable=True)
settings.source = openmc.source.IndependentSource(space=uniform_dist)

entropy_mesh = openmc.RegularMesh()
entropy_mesh.lower_left = [-hpitch, -hpitch, -1.e50]
entropy_mesh.upper_right = [hpitch, hpitch, 1.e50]
entropy_mesh.dimension = [10, 10, 1]

settings.entropy_mesh = entropy_mesh
settings.export_to_xml()

###############################################################################
#                   Exporting to OpenMC tallies.xml file
###############################################################################

## Instantiate a tally mesh
mesh = openmc.RegularMesh()
mesh.type = 'regular'
mesh.dimension = [1, 1, 100]
mesh.lower_left = [-hpitch, -hpitch, -fheight/2]
mesh.upper_right = [hpitch, hpitch, fheight/2]

## Instantiate some tally Filters
energy_filter = openmc.EnergyFilter([0., 0.625, 20.e6])
energy_filter2 = openmc.EnergyFilter(np.logspace(-3, 7, num=100))
mesh_filter = openmc.MeshFilter(mesh)

## Instantiate the Tally
tally = openmc.Tally(tally_id=1, name='Spatial Tally 1')
tally.filters = [energy_filter, mesh_filter]
tally.scores = ['flux']#, 'fission', 'nu-fission']

## Instantiate the Tally
tallyspec = openmc.Tally(tally_id=2, name='Energy Spectrum Tally 1')
tallyspec.filters = [energy_filter2]
tallyspec.scores = ['flux']#, 'fission', 'nu-fission']

particle_filter = openmc.ParticleFilter(['neutron', 'photon'])
mat_filter = openmc.MaterialFilter(materials)

tallydep = openmc.Tally(tally_id=3, name='Spatial Tally 2')
tallydep.filters = [mat_filter, particle_filter]
tallydep.scores = ['fission', 'heating', 'heating-local']

## Instantiate a Tallies collection and export to XML
tallies = openmc.Tallies([tally, tallyspec,tallydep])
tallies.export_to_xml()

###############################################################################
#                      Depletion settings
###############################################################################

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

## setting the transport operator
operator = openmc.deplete.CoupledOperator(model,chain_path,diff_burnable_mats=False,normalization_mode='fission-q',fission_q={"U235": 202.27e6})

## setting the system linear power [W]
power = [45.50,91.01,136.51,182.02,227.52,273.03,324.45,324.45,273.03,273.03,227.52,182.02,136.51,91.01,45.50]
time_steps = [8640.00,8640.00,8640.00,8640.00,8640.00,8640.00,4268160.00,45792000.00,1728000.00,46915200.00,8640.00,8640.00,8640.00,8640.00,8640.00]

sys.stdout.flush()
try:
    MPI.COMM_WORLD.Barrier()
except NameError:
    pass
if rank == 0:
    print(time_steps)
sys.stdout.flush()

## depleting usin a first-order predictor algorithm
integrator = openmc.deplete.PredictorIntegrator(operator, time_steps, power, timestep_units =  's')

integrator.integrate()