##############################################################################
#                          ALFRED Fuel Assembly
##############################################################################

import openmc
import numpy as np
import math

###############################################################################
#                      Simulation Input File Parameters
###############################################################################
batches = 20
inactive = 1
particles = 1000
generations = 10

###############################################################################
#                 Exporting to OpenMC materials.xml File
###############################################################################
mox_fuel = openmc.Material(name='MOX fuel', temperature = 1473)
mox_fuel.add_nuclide('O16'  ,	4.6204e-02)
mox_fuel.add_nuclide('U234' ,   5.5962e-07)
mox_fuel.add_nuclide('U235' ,   7.5067e-05)
mox_fuel.add_nuclide('U236' ,   1.8499e-06)
mox_fuel.add_nuclide('U238' ,   1.8262e-02)
mox_fuel.add_nuclide('Pu238',	1.1698e-04)
mox_fuel.add_nuclide('Pu239',	2.8413e-03)
mox_fuel.add_nuclide('Pu240',	1.3433e-03)
mox_fuel.add_nuclide('Pu241',	3.0251e-04)
mox_fuel.add_nuclide('Pu242',	3.7955e-04)
mox_fuel.add_nuclide('Am241',	6.4786e-05)
mox_fuel.volume = math.pi * (0.450 - 0.100)**2 * 127

ti1515 = openmc.Material(name='Ti15-15', temperature = 673)
ti1515.add_element('C'    , 3.5874e-04)
ti1515.add_nuclide('Cr50' , 5.8098e-04)
ti1515.add_nuclide('Cr52' , 1.1187e-02)
ti1515.add_nuclide('Cr53' , 1.2680e-03)
ti1515.add_nuclide('Cr54' , 3.1532e-04)
ti1515.add_nuclide('Ni58' , 8.6075e-03)
ti1515.add_nuclide('Ni60' , 3.3160e-03)
ti1515.add_nuclide('Ni61' , 1.4371e-04)
ti1515.add_nuclide('Ni62' , 4.5894e-04)
ti1515.add_nuclide('Ni64' , 1.1724e-04)
ti1515.add_nuclide('Mn55' , 1.3072e-03)
ti1515.add_nuclide('Mo92' , 1.1596e-04)
ti1515.add_nuclide('Mo94' , 7.0739e-05)
ti1515.add_nuclide('Mo95' , 1.2046e-04)
ti1515.add_nuclide('Mo96' , 1.2490e-04)
ti1515.add_nuclide('Mo97' , 7.0772e-05)
ti1515.add_nuclide('Mo98' , 1.7699e-04)
ti1515.add_nuclide('Mo100', 6.9221e-05)
ti1515.add_nuclide('Ti46' , 3.4381e-05)
ti1515.add_nuclide('Ti47' , 3.0346e-05)
ti1515.add_nuclide('Ti48' , 2.9444e-04)
ti1515.add_nuclide('Ti49' , 2.1166e-05)
ti1515.add_nuclide('Ti50' , 1.9862e-05)
ti1515.add_nuclide('Si28' , 1.3363e-03)
ti1515.add_nuclide('Si29' , 6.7833e-05)
ti1515.add_nuclide('Si30' , 4.4803e-05)
ti1515.add_nuclide('B10'  , 5.7090e-06)
ti1515.add_nuclide('B11'  , 2.0900e-05)
ti1515.add_nuclide('P31'  , 6.9556e-05)
ti1515.add_nuclide('N14'  , 5.1096e-05)
ti1515.add_nuclide('N15'  , 1.7618e-07)
ti1515.add_nuclide('S32'  , 2.1323e-05)
ti1515.add_nuclide('S33'  , 1.6553e-07)
ti1515.add_nuclide('S34'  , 9.0698e-07)
ti1515.add_nuclide('S36'  , 3.9933e-09)
ti1515.add_nuclide('Al27' , 2.6616e-05)
ti1515.add_nuclide('Zr90' , 8.2194e-06)
ti1515.add_nuclide('Zr91' , 1.7727e-06)
ti1515.add_nuclide('Zr92' , 2.6802e-06)
ti1515.add_nuclide('Zr94' , 2.6582e-06)
ti1515.add_nuclide('Zr96' , 4.1932e-07)
ti1515.add_nuclide('V51'  , 2.8193e-05)
ti1515.add_nuclide('W182' , 2.0944e-06)
ti1515.add_nuclide('W183' , 1.1248e-06)
ti1515.add_nuclide('W184' , 2.3952e-06)
ti1515.add_nuclide('W186' , 2.1985e-06)
ti1515.add_nuclide('Nb93' , 7.7297e-06)
ti1515.add_nuclide('Ta181', 3.9688e-06)
ti1515.add_nuclide('Cu63' , 1.5787e-05)
ti1515.add_nuclide('Cu65' , 6.8200e-06)
ti1515.add_nuclide('Co59' , 2.4371e-05)
ti1515.add_nuclide('Ca40' , 3.4843e-05)
ti1515.add_nuclide('Ca42' , 2.2147e-07)
ti1515.add_nuclide('Ca43' , 4.5136e-08)
ti1515.add_nuclide('Ca44' , 6.8162e-07)
ti1515.add_nuclide('Ca48' , 5.6010e-08)
ti1515.add_nuclide('Fe54' , 3.2774e-03)
ti1515.add_nuclide('Fe56' , 5.1407e-02)
ti1515.add_nuclide('Fe57' , 1.1870e-03)
ti1515.add_nuclide('Fe58' , 1.5662e-04)
ti1515.volume = math.pi * (0.525**2 - 0.465**2) * 127

fmst91 = openmc.Material(name='FMST91', temperature = 673)
fmst91.add_element('C'    ,  3.8786e-04)
fmst91.add_nuclide('Si28' ,  7.6490e-04)
fmst91.add_nuclide('Si29' ,  3.8827e-05)
fmst91.add_nuclide('Si30' ,  2.5645e-05)
fmst91.add_nuclide('V51'  ,  1.8289e-04)
fmst91.add_nuclide('Cr50' ,  3.5090e-04)
fmst91.add_nuclide('Cr52' ,  6.7565e-03)
fmst91.add_nuclide('Cr53' ,  7.6585e-04)
fmst91.add_nuclide('Cr54' ,  1.9044e-04)
fmst91.add_nuclide('Mn55' ,  5.0879e-04)
fmst91.add_nuclide('Mo92' ,  7.5223e-05)
fmst91.add_nuclide('Mo94' ,  4.5890e-05)
fmst91.add_nuclide('Mo95' ,  7.8147e-05)
fmst91.add_nuclide('Mo96' ,  8.1025e-05)
fmst91.add_nuclide('Mo97' ,  4.5911e-05)
fmst91.add_nuclide('Mo98' ,  1.1482e-04)
fmst91.add_nuclide('Mo100',  4.4905e-05)
fmst91.add_nuclide('Ni58' ,  1.0807e-04)
fmst91.add_nuclide('Ni60' ,  4.1634e-05)
fmst91.add_nuclide('Ni61' ,  1.8044e-06)
fmst91.add_nuclide('Ni62' ,  5.7624e-06)
fmst91.add_nuclide('Ni64' ,  1.4720e-06)
fmst91.add_nuclide('Nb93' ,  5.0144e-05)
fmst91.add_nuclide('Fe54' ,  4.3089e-03)
fmst91.add_nuclide('Fe56' ,  6.7586e-02)
fmst91.add_nuclide('Fe57' ,  1.5606e-03)
fmst91.add_nuclide('Fe58' ,  2.0592e-04)

lead = openmc.Material(name='Lead', temperature = 673)
lead.add_nuclide('Pb204',  4.5465e-04)
lead.add_nuclide('Pb206',  7.7163e-03)
lead.add_nuclide('Pb207',  7.0759e-03)
lead.add_nuclide('Pb208',  1.6771e-02)

materials = openmc.Materials([mox_fuel, ti1515, fmst91, lead])
materials.export_to_xml()

###############################################################################
#                 Exporting to OpenMC geometry.xml file
###############################################################################
fuel_ir = openmc.ZCylinder(r=0.100, name='Fuel IR')
fuel_or = openmc.ZCylinder(r=0.450, name='Fuel OR')
clad_ir = openmc.ZCylinder(r=0.465, name='Clad IR')
clad_or = openmc.ZCylinder(r=0.525, name='Clad OR')

inner_duct = openmc.model.hexagonal_prism(edge_length=9.12213425, orientation='y', boundary_type='transmission')
outer_duct = openmc.model.hexagonal_prism(edge_length=9.58401447, orientation='y', boundary_type='transmission')
boundary   = openmc.model.hexagonal_prism(edge_length=9.87268960, orientation='y', boundary_type='reflective')

hole    = openmc.Cell(fill=None    , region=-fuel_ir)
fuel    = openmc.Cell(fill=mox_fuel, region=+fuel_ir & -fuel_or)
gap     = openmc.Cell(fill=None    , region=+fuel_or & -clad_ir)
clad    = openmc.Cell(fill=ti1515  , region=+clad_ir & -clad_or)
coolant = openmc.Cell(fill=lead    , region=+clad_or)
pin = openmc.Universe(cells=[hole, fuel, gap, clad, coolant])

lead_universe = openmc.Universe(cells=[openmc.Cell(fill = lead)])
#####################################################################
#                          Lattice
#####################################################################
lattice = openmc.HexLattice()
lattice.center = (0., 0.)
lattice.pitch = [1.386]
lattice.orientation = 'y'
lattice.outer = lead_universe
ring_0 = [pin]
ring_1 = [pin]*6
ring_2 = [pin]*12
ring_3 = [pin]*18
ring_4 = [pin]*24
ring_5 = [pin]*30
ring_6 = [pin]*36
lattice.universes = [ring_6, ring_5, ring_4, ring_3, ring_2, ring_1, ring_0]

#####################################################################
#                           Assembly
#####################################################################)
fuel_cell     = openmc.Cell(fill = lattice, region =  inner_duct)
duct_cell     = openmc.Cell(fill = fmst91 , region = ~inner_duct & outer_duct)
outer_cell    = openmc.Cell(fill = lead   , region = ~outer_duct & boundary)
assembly = openmc.Universe(cells=[fuel_cell,duct_cell,outer_cell])

geometry = openmc.Geometry(assembly)
geometry.export_to_xml()

###############################################################################
#                   Exporting to OpenMC settings.xml file
###############################################################################
settings = openmc.Settings()
settings.run_mode = 'eigenvalue'
settings.batches = batches
settings.inactive = inactive
settings.generations_per_batch = generations
settings.particles = particles
settings.temperature = {'method': 'interpolation'}

bounds = [-20, -20, 0, 20, 20, 1]
uniform_dist = openmc.stats.Box(bounds[:3], bounds[3:], only_fissionable=True)
settings.source = openmc.source.Source(space=uniform_dist)
settings.ptables = True

settings.export_to_xml()

###############################################################################
#                   Exporting to OpenMC plots.xml file
###############################################################################
plot_xy = openmc.Plot()
plot_xy.filename = 'plot_xy'
plot_xy.basis = 'xy'
plot_xy.origin = [0, 0, 0]
plot_xy.width = [30, 30]
plot_xy.pixels = [2000, 2000]
plot_xy.color_by = 'material'
plot_xy.colors = {}

plots = openmc.Plots([plot_xy])
plots.export_to_xml()

#openmc.plot_geometry()

###############################################################################
#                   Exporting to OpenMC tallies.xml file
###############################################################################
energy_deposition = openmc.Tally(name='Energy Deposition')
energy_deposition.scores = ['heating-local']

clad_damage_energy = openmc.Tally(name='Damage Energy')
clad_damage_energy.filters = [openmc.MaterialFilter(ti1515)]
clad_damage_energy.scores = ['damage-energy']
clad_damage_energy.nuclides = ['C0', 'Si28', 'Si29', 'Si30', 'Ti46', 'Ti47', 'Ti48', 'Ti49', 'Ti50', 'Cr50', 'Cr52', 'Cr53', 'Cr54', 'Mn55', 'Fe54', 'Fe56', 'Fe57', 'Fe58', 'Ni58', 'Ni60', 'Ni61', 'Ni62', 'Ni64', 'Mo92', 'Mo94', 'Mo95', 'Mo96', 'Mo97', 'Mo98', 'Mo100']

tallies = openmc.Tallies([energy_deposition, clad_damage_energy])
tallies.export_to_xml()

###############################################################################
#                    Running the model
###############################################################################

openmc.run()

###############################################################################
#                    Computing DPA
###############################################################################

# Elements considering in the DPA calculations
elements = ['C', 'Si', 'Ti', 'Cr', 'Mn', 'Fe', 'Ni', 'Mo']
# The Atomic Displacement Energy needed to compute DPA  [eV], see NJOY manual @p.20
E_d = {'C'  : 31,
       'Si' : 25,
       'Ti' : 40,
       'Cr' : 40,
       'Mn' : 40,
       'Fe' : 40,
       'Ni' : 40,
       'Mo' : 60}

# Initialize sets and constants for a later use
number_of_atoms = {'C' : 0, 'Si' : 0,'Ti' : 0., 'Cr' : 0., 'Mn' : 0., 'Fe' : 0., 'Ni' : 0., 'Mo' : 0.}
deps = {}
dps = {}
dpa = {}
average_power = 300e6/171/60
q = 1.60217663e-19
irradiation_time = 365 * 5 * 24 * 3600 # [s]

# Calcuclate number of atoms in the cladding for chosen elements
for e in elements:
    nuclides = ti1515.get_nuclides(e)
    for n in nuclides:
        nuclide_number = ti1515.get_nuclide_atom_densities(n)[n] * ti1515.volume * 1e24 
        number_of_atoms[e] += nuclide_number

with openmc.StatePoint(f'statepoint.{batches}.h5') as sp:
    # Get the heating tally to normalize the source
    heating_tally = sp.get_tally(name='Energy Deposition') 
    heating_value = heating_tally.mean.ravel()[0]

    # Get the damage energy tally  
    damage_tally = sp.get_tally(name='Damage Energy') 
    nuclides = damage_tally.nuclides

    # Get the number of neutrons per second via power [src/s] = [W/eV*src/J*eV]
    nps = average_power/heating_value/q

    for e in elements:
        # Get the damage energy for each element per src [eV/src]
        deps[e] = damage_tally.get_slice(nuclides = [n for n in nuclides if e in n]).mean.sum()
        
        # The number of displacements per src [1/src]
        dps[e] = 0.8*deps[e]/ (2*E_d[e]) 

        # Normalize N_d to get the number of displacements per second [1/s] = [(1/src) * (src/s)]
        dps[e] = dps[e] * nps 

        # Get total dpa
        dpa[e] = dps[e] / number_of_atoms[e] * irradiation_time

print(dpa)
print('Total DPA', sum(dps.values()) / sum(number_of_atoms.values()) * irradiation_time)