#!/usr/bin/env python

import openmc, os
import openmc.deplete
from beavrs.builder import BEAVRS
from optparse import OptionParser

try:
    os.mkdir('build')
except OSError: pass
os.chdir('build')

usage = """usage: %prog [options]"""
p = OptionParser(usage=usage)
p.add_option('-d', '--2d', action='store_true', dest='is_2d',
             default=False, help='Create 2D BEAVRS input files,' \
             + ' 3D by default')
p.add_option('-s', '--symmetric', action='store_true', dest='is_symmetric',
             default=False, help='Create octant-symmetric input files,' \
             + ' not symmetric by default')
p.add_option('-e', '--deplete', action='store_true', dest='run_deplete', default=False,
             help='Run the depletion code for this setup')
(options, args) = p.parse_args()

if not len(args) == 0:
    p.print_help()

chain_path = "/opt/xdata/endfb-vii.1-hdf5/chain_endfb71_pwr.xml"

b = BEAVRS(is_symmetric=options.is_symmetric, is_2d=options.is_2d)
b.write_openmc_geometry()
b.write_openmc_materials()
b.write_openmc_plots()
b.write_openmc_settings()

# openmc.calculate_volumes()

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

if options.run_deplete:
    model = openmc.model.Model(geometry=b.openmc_geometry, materials=b.materials, settings=b.settings, tallies=b.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]

    print(time_steps)

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

    integrator.integrate()
