import openmc
import openmc.deplete
import numpy as np
import matplotlib.pyplot as plt

###############################################################################
#                      Simulation Input File Parameters
###############################################################################

## OpenMC simulation parameters
batches = 300
inactive = 50
particles = 100000

## Depletion simulation parameters
ts = 40*24*60*60 # s
td = 10*24*60*60 # s

time_steps = np.array([ts,ts,ts,ts,ts,ts,ts,ts,ts,ts,td,ts,ts,ts,ts,ts,ts,ts,ts,ts,ts,td,ts,ts,ts,ts,ts,ts,ts,ts,ts,ts,td,ts,ts,ts,ts,ts,ts,ts,ts,ts,ts,td])

chain_file = '../chain_endfb71.xml'
power = 17373 # W/cm, for 2D simulations only (use W for 3D)

###############################################################################
#                              Define materials
###############################################################################

## Instantiate some Materials and register the appropriate Nuclides
zircaloy = openmc.Material(material_id=3, name='Zircaloy 4')
zircaloy.set_density('g/cm3', 6.49012)
zircaloy.temperature = 600.0
zircaloy.add_nuclide('Sn112', 1.37041E-04, 'wo')
zircaloy.add_nuclide('Sn114', 9.34709E-05, 'wo')
zircaloy.add_nuclide('Sn115', 4.93221E-05, 'wo')
zircaloy.add_nuclide('Sn116', 2.12757E-03, 'wo')
zircaloy.add_nuclide('Sn117', 1.13348E-03, 'wo')
zircaloy.add_nuclide('Sn118', 3.60520E-03, 'wo')
zircaloy.add_nuclide('Sn119', 1.28950E-03, 'wo')
zircaloy.add_nuclide('Sn120', 4.93338E-03, 'wo')
zircaloy.add_nuclide('Sn122', 7.12568E-04, 'wo')
zircaloy.add_nuclide('Sn124', 9.05713E-04, 'wo')
zircaloy.add_nuclide('Fe54', 7.34249E-05, 'wo')
zircaloy.add_nuclide('Fe56', 1.19419E-03, 'wo')
zircaloy.add_nuclide('Fe57', 2.80871E-05, 'wo')
zircaloy.add_nuclide('Fe58', 3.77460E-06, 'wo')
zircaloy.add_nuclide('Cr50', 4.17870E-05, 'wo')
zircaloy.add_nuclide('Cr52', 8.37035E-04, 'wo')
zircaloy.add_nuclide('Cr53', 9.67293E-05, 'wo')
zircaloy.add_nuclide('Cr54', 2.44824E-05, 'wo')
zircaloy.add_nuclide('Ni58', 4.03133E-04, 'wo')
zircaloy.add_nuclide('Ni60', 1.60602E-04, 'wo')
zircaloy.add_nuclide('Ni61', 7.09947E-06, 'wo')
zircaloy.add_nuclide('Ni62', 2.29762E-05, 'wo')
zircaloy.add_nuclide('Ni64', 6.07653E-06, 'wo')
zircaloy.add_nuclide('Zr90', 4.97990E-01, 'wo')
zircaloy.add_nuclide('Zr91', 1.09810E-01, 'wo')
zircaloy.add_nuclide('Zr92', 1.69691E-01, 'wo')
zircaloy.add_nuclide('Zr94', 1.75712E-01, 'wo')
zircaloy.add_nuclide('Zr96', 2.89111E-02, 'wo')


uo2 = openmc.Material(name='Fuel Batch 1')
uo2.set_density('g/cc' ,10.07)
uo2.temperature = 900.0
uo2.add_nuclide('O16'  ,1.18530E-01,'wo')
uo2.add_nuclide('U234' ,3.92254E-04,'wo')
uo2.add_nuclide('U235' ,4.36326E-02,'wo')
uo2.add_nuclide('U238' ,8.37445E-01,'wo')

uo2.depletable = True


borated_water = openmc.Material(material_id=4, name='Borated water')
borated_water.set_density('g/cm3', 0.7245)
borated_water.temperature = 600.0
borated_water.add_nuclide('B10',8.40178E-05,'wo')
borated_water.add_nuclide('B11',3.71839E-04,'wo')
borated_water.add_nuclide('H1', 1.11832E-01,'wo')
borated_water.add_nuclide('H2', 3.35290E-05,'wo')
borated_water.add_nuclide('O16',8.85414E-01,'wo' )
borated_water.add_nuclide('O17',2.26386E-03,'wo' )
borated_water.add_s_alpha_beta('c_H_in_H2O')



###############################################################################
#                             Create geometry
###############################################################################

## Instantiate ZCylinder surfaces
fuel_or = openmc.ZCylinder(surface_id=1, x0=0, y0=0, r=0.41266, name='Fuel OR')
clad_ir = openmc.ZCylinder(surface_id=2, x0=0, y0=0, r=0.474364, name='Clad IR')


left = openmc.XPlane(surface_id=4, x0=-0.632458, name='left')
right = openmc.XPlane(surface_id=5, x0=0.632458, name='right')
bottom = openmc.YPlane(surface_id=6, y0=-0.632458, name='bottom')
top = openmc.YPlane(surface_id=7, y0=0.632458, name='top')
up = openmc.ZPlane(surface_id=8, z0=50.0, name='up')
down = openmc.ZPlane(surface_id=9, z0=-50.0, name='down')

left.boundary_type   = 'reflective'
right.boundary_type  = 'reflective'
top.boundary_type    = 'reflective'
bottom.boundary_type = 'reflective'
up.boundary_type     = 'reflective'
down.boundary_type   = 'reflective'



## Instantiate Cells
fuel = openmc.Cell(cell_id=1, name='cell 1')
clad = openmc.Cell(cell_id=3, name='cell 3')
water = openmc.Cell(cell_id=4, name='cell 4')

## Use surface half-spaces to define regions
fuel.region = -fuel_or & -up & +down
clad.region = +fuel_or & -clad_ir  & -up & +down
water.region = +clad_ir & +left & -right & +bottom & -top & -up & +down

## Register Materials with Cells
fuel.fill = uo2
clad.fill = zircaloy
water.fill = borated_water

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

## Register Cells with Universe
root.add_cells([fuel, clad, water])

## Instantiate a Geometry, register the root Universe
geometry = openmc.Geometry(root)

###############################################################################
#                     Set volumes of depletable materials
###############################################################################

## Compute cell areas
## area = {}
## area[fuel] = np.pi * fuel_or.coefficients['r'] ** 2

## Set materials volume for depletion. Set to an area for 2D simulations
uo2.volume = 53.49763956

###############################################################################
#                     Transport calculation settings
###############################################################################

## Instantiate a Settings object, set all runtime parameters, and export to XML
settings_file = openmc.Settings()
settings_file.batches = batches
settings_file.inactive = inactive
settings_file.particles = particles
settings_file.photon_transport = True
## Create an initial uniform spatial source distribution over fissionable zones
bounds = [-0.632458, -0.632458, -50, 0.632458, 0.632458, 50]
# uniform_dist = openmc.stats.Box(bounds[:3], bounds[3:])
uniform_dist = openmc.stats.Box(bounds[:3], bounds[3:], only_fissionable=True)
settings_file.source = openmc.source.IndependentSource(space=uniform_dist)
## settings_file.electron_treatment = 'ttb'
## settings_file.cutoff = {'energy_photon' : 1000.0}

#entropy_mesh = openmc.RegularMesh()
#entropy_mesh.lower_left = [-0.39218, -0.39218, -50]
#entropy_mesh.upper_right = [0.39218, 0.39218, 50]
#entropy_mesh.dimension = [10, 10, 1]
#settings_file.entropy_mesh = entropy_mesh

###############################################################################
#                   Initialize and run depletion calculation
###############################################################################
tallies_file = openmc.Tallies()
t1 = openmc.Tally(name='edep')
t1.filters = [openmc.MaterialFilter([uo2,zircaloy,borated_water])]
#t1.filters.append(openmc.ParticleFilter(['neutron']))
#,'photon','electron','positron']))
t1.scores = ['flux','fission', 'nu-fission','heating-local','heating']
tallies_file.append(t1)
tallies_file.export_to_xml()


fission_q={"U232":     193.0442774*1E6,
           "U233":     199.7961836*1E6,
           "U234":     200.632851*1E6,
           "U235":     202.27*1E6,
           "U236":     203.4043119*1E6,
           "U237":     196.429643*1E6,
           "U238":     206.8513817*1E6,
           "Np237":    205.3704803*1E6,
           "Np238":    208.6993709*1E6,
           "Pu238":    209.5400125*1E6,
           "Pu239":    208.0185328*1E6,
           "Pu240":    208.6125667*1E6,
           "Pu241":    211.2377153*1E6,
           "Pu242":    212.0721865*1E6,
           "Am241":    211.2167986*1E6,
           "Am242_m1": 215.1453706*1E6,
           "Am243":    212.952779*1E6,
           "Cm242":    212.7864913*1E6,
           "Cm243":    213.375296*1E6,
           "Cm244":    217.9267669*1E6,
           "Cm245":    214.6240222*1E6,
           "Cm246":    220.179494*1E6,
           "Th232":    197.1083894*1E6,
           "Pa231":    194.0999425*1E6}

## op = openmc.deplete.Operator(geometry, settings_file, chain_file,fission_q=fission_q)
op = openmc.deplete.Operator(geometry, settings_file, chain_file, energy_mode="energy-deposition")

## Perform simulation using the predictor algorithm
integrator = openmc.deplete.PredictorIntegrator(op, time_steps, power)
integrator.integrate()

## openmc.run(mpi_args=['mpiexec', '-n', '32', '--bind-to', 'core'])

###############################################################################
#                    Read depletion calculation results
###############################################################################

## Open results file
#results = openmc.deplete.ResultsList.from_hdf5("depletion_results.h5")

## Obtain K_eff as a function of time
#time, keff = results.get_eigenvalue()

## Obtain U235 concentration as a function of time
#time, n_U235 = results.get_atoms('1', 'U235')

## Obtain Xe135 absorption as a function of time
#time, Xe_gam = results.get_reaction_rate('1', 'Xe135', '(n,gamma)')
# t_days, n_cm3 = res.get_atoms("1", "Xe135", nuc_units="atom/cm3", time_units="d")

###############################################################################
#                            Generate plots
###############################################################################

#plt.figure()
#plt.plot(time/(24*60*60), keff, label="K-effective")
#plt.xlabel("Time (days)")
#plt.ylabel("Keff")
#plt.show()

#plt.figure()
#plt.plot(time/(24*60*60), n_U235, label="U 235")
#plt.xlabel("Time (days)")
#plt.ylabel("n U5 (-)")
#plt.show()

#plt.figure()
#plt.plot(time/(24*60*60), Xe_gam, label="Xe135 absorption")
#plt.xlabel("Time (days)")
#plt.ylabel("RR (-)")
#plt.show()
#plt.close('all')
#tallies_file = openmc.Tallies()
#energy_filter = openmc.EnergyFilter([0., 20.0e6])

## Instantiate flux Tally in moderator and fuel
#tally = openmc.Tally(name='flux')
#tally.filters = [openmc.CellFilter(fuel)]
#tally.filters.append(energy_filter)
#tally.scores = ['flux']
#tallies_file.append(tally)

## Instantiate reaction rate Tally in moderator
#tally = openmc.Tally(name='fuel rxn rates')
#tally.filters = [openmc.CellFilter(fuel)]
#tally.filters.append(energy_filter)
#tally.scores = ['fission']
#tally.nuclides = ['O16', 'H1']
#tallies_file.append(tally)

#tallies_file.export_to_xml()
