#!/usr/bin/env python3

import numpy as np
import openmc.deplete
from beavrs.builder import BEAVRS
import argparse, json, openmc, os, shutil, sys
import matplotlib.pyplot as plt
try:
    from mpi4py import MPI
except ModuleNotFoundError:
    print("Cannot use MPI")

try:
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
except NameError:
    comm = None
    rank = 0

parser = argparse.ArgumentParser(description="Program to make BEAVRS reactor and run" \
                                             " simulations with resulting geometry")
parser.add_argument('-d', '--2d', action='store_true', dest='is_2d',
             default=False, help='Create 2D BEAVRS input files, 3D by default')
parser.add_argument('-s', '--symmetric', action='store_true', dest='is_symmetric',
             default=False, help='Create octant-symmetric input files, not symmetric by default')
parser.add_argument('-e', '--deplete', action='store_true', dest='run_deplete', default=False,
             help='Run the depletion code for this setup')
parser.add_argument('-i', '--insert', action='store', type=int, dest='rod_loc', default=200,
             help='Set where the control rods should be. Acceptable values are 15-573')
parser.add_argument('-o', '--optimal', action='store_true', default=False,
             help='Find the best combination of particles, inactive and active batches to run')
parser.add_argument('-r', '--run', action='store_true', default=False,
             help='Run the model once after building it')
parser.add_argument('--shutdown', action='store_true', default=False,
             help='Set reactor to be in shutdown, all control rods fully inserted.')
args = parser.parse_args()

if rank == 0:
    try:
        os.mkdir(os.path.dirname(os.path.realpath(__file__)) + '/build')
    except OSError: pass
comm.Barrier()

os.chdir(os.path.dirname(os.path.realpath(__file__)) + '/build')

def no_overwrite(rank, filename):
    if os.path.isfile(filename) and rank == 0:
        j = 0
        new_filepath = filename*1
        path, new_name = os.path.split(new_filepath)
        name_len = len(new_name)
        while os.path.isfile(new_filepath):
            if new_name[-2] == 'h':
                new_name = f'#{new_name}.{j}'
            else:
                j += 1
                new_name = f'{new_name[:name_len+1]}.{j}'

            new_filepath = f"{path}/{new_name}"

        print(f"Backed up file {filename} to {new_filepath}")
        os.replace(filename, new_filepath)

    comm.Barrier()

def run_loop(rank, val):
    b.set_S(rank, val)
    return b.export_model()

if args.shutdown:
    b = BEAVRS(rank, 573, 228, is_symmetric=args.is_symmetric, is_2d=args.is_2d)
else:
    b = BEAVRS(rank, args.rod_loc, 0, is_symmetric=args.is_symmetric, is_2d=args.is_2d)

# openmc.calculate_volumes()
# openmc.plot_geometry()

if args.optimal:

    particles = [5000, 10000, 20000, 50000, 100000]
    inactive = [10, 20, 50, 100, 150]
    active = [50, 100, 150, 300, 500]

    for p in particles:
        for i in inactive:
            for a in active:
                if rank == 0:
                    print(f"Particles: {p} Inactive Cycles: {i} Active Cycles: {a}\n\n")
                b.set_params(batches=i+a,inactive=i,particles=p)
                b.export_xml(rank)
                openmc.run()
                if rank == 0:
                    print("\n\n\n")
                sys.stdout.flush()

###############################################################################
#                      Depletion Settings
###############################################################################

elif args.run_deplete:
    b.set_params(batches=350,inactive=50,particles=20000)
    b.set_full_depletion_mats()
    # b.set_fuel_otf_files()
    # b.set_distributed_materials()
    # b.set_fuel_distrib_tallies()
    # b.set_fuel_depletion_tallies()

    ## Get fission Q values from JSON file generated by get_fission_qvals.py
    with open('../beavrs/serpent_fissq.json', 'r') as f:
        serpent_fission_q = json.load(f)

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

    if rank == 0:
        try:
            os.makedirs('depletion_results')
        except FileExistsError:
            pass

    ## setting the system linear power [W]
    """powers = [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]"""

    ## cumulative steps in days
    # time_steps_cum = np.array([3/24, 6/24, 9/24, 12/24, 15/24, 18/24, 21/24, 1, 10, 30, 60, 90, 365, 548])
    #                      i =  298    300    300    300    300    300    300  300 300 300 300 300 300  300
    time_steps_cum = np.array([12/24, 24/24, 36/24, 48/24, 60/24, 72/24, 84/24, 4, 10, 30, 60, 90, 365, 548])
    time_steps = np.diff(time_steps_cum, prepend=0.0)

    ## Power in Watts
    full_power = 3411e6
    powers = [full_power if x > 4 else full_power*x/4 for x in time_steps_cum]

    ## setting the transport operator
    model = b.export_model()
    comm.Barrier()
    # openmc.plot_geometry()
    operator = openmc.deplete.CoupledOperator(model, chain_file,
                                        #diff_burnable_mats=False, normalization_mode='fission-q',
                                        fission_q=serpent_fission_q, fission_yield_mode="average")

    ## depleting using a first-order predictor algorithm
    ## Working algorithms: CELI, CECM, SILEQI
    integrator = openmc.deplete.SILEQIIntegrator(operator, time_steps, powers, timestep_units='d')
    integrator.integrate()

    """
    if rank == 0:
        print(time_steps)

    i = 0
    for power, time_step in zip(powers, time_steps):

        if rank == 0:
            print(f'Current loop info: {power} {time_step}')

        ## setting the transport operator
        model = b.export_model()
        comm.Barrier()
        openmc.plot_geometry()
        operator = openmc.deplete.CoupledOperator(model, chain_file,
                                        #diff_burnable_mats=False, normalization_mode='fission-q',
                                        fission_q=serpent_fission_q, fission_yield_mode="average")

        ## depleting using a first-order predictor algorithm
        ## Working algorithms: CELI, CECM, SILEQI
        integrator = openmc.deplete.SILEQIIntegrator(operator, time_steps, powers, timestep_units='d')
        integrator.integrate()

        filename = f'depletion_results/depletion_results_t{i}.h5'
        no_overwrite(rank, filename)
        os.replace('depletion_results.h5', filename)
        shutil.copy2('materials.xml', f'depletion_results/materials_{i}.xml')
        i += 1
    """

    results = openmc.deplete.Results('depletion_results.h5')

    # Get materials at the end of the last simulation
    if rank == 0:
        model.materials = results.export_to_materials(len(time_steps))
        model.export_to_xml()
    comm.Barrier()

    # Obtain K_eff as a function of time
    time, keff = results.get_keff(time_units='d')

    n_U235 = 0
    Xe_capture = 0
    n_Xe135 = 0
    tmp = [[],[],[]]
    # ['Fuel 1.6%' 'Fuel 2.4%' 'Fuel 3.1%' 'Fuel 3.2%' 'Fuel 3.4%']
    for mat in b.mats:
        # Obtain U235 concentration as a function of time
        _, tmp[0] = results.get_atoms(mat, 'U235')

        # Obtain Xe135 capture reaction rate as a function of time
        _, tmp[1] = results.get_reaction_rate(mat, 'Xe135', '(n,gamma)')

        # Obtain U235 concentration as a function of time
        _, tmp[2] = results.get_atoms(mat, 'Xe135')

        n_U235 += tmp[0]
        Xe_capture += tmp[1]
        n_Xe135 += tmp[2]

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

    fig, ax = plt.subplots()
    ax.errorbar(time,keff[:,0], keff[:,1], label="k-effective")
    ax.set_xlabel("Time [d]")
    ax.set_ylabel("Keff")
    if rank == 0:
        plt.savefig("k-effective.png", dpi=150)
        plt.draw()

    fig, ax = plt.subplots()
    ax.plot(time, n_U235, label="U235")
    ax.set_xlabel("Time [d]")
    ax.set_ylabel("U235 atoms")
    if rank == 0:
        plt.savefig("U235", dpi=150)
        plt.draw()

    fig, ax = plt.subplots()
    ax.plot(time, Xe_capture, label="Xe135 capture")
    ax.set_xlabel("Time [d]")
    ax.set_ylabel("Xe135 capture rate")
    if rank == 0:
        plt.savefig("Xe135_capture", dpi=150)
        plt.draw()

    fig, ax = plt.subplots()
    ax.plot(time, n_Xe135, label="Xe135")
    ax.set_xlabel("Time [d]")
    ax.set_ylabel("Xe135 atoms")
    if rank == 0:
        plt.savefig("Xe135", dpi=150)
        plt.draw()

elif args.run:
    b.set_params(batches=350,inactive=50,particles=20000)

    """rod_loc, guess_list, result = openmc.search_for_keff(run_loop, initial_guess=10, target=1.0,
                                                     tol=6e-4, print_iterations=True, run_args=(rank))"""

    b.export_xml(rank)
    openmc.run()

else:
    b.set_params(batches=350,inactive=50,particles=20000)
    b.export_xml(rank)






