# Copyright 2018 Clemson University
# 
# Author: Bradley S. Meyer
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# To obtain a copy of the GNU General Public License,
# see <https://www.gnu.org/licenses/>.

import argparse
import wnutils.xml as wx

parser = argparse.ArgumentParser(
            description='Push all beta-delayed neutron reactions to their\
                         non-beta-delayed neutron counterparts. Create \
                         the non-beta-delayed neutron counterpart if not \
                         present. Assumes all relevant beta-decay rates \
                         are of single_rate type.'
         )

parser.add_argument('in_xml', metavar='in_xml', help='input network xml')
parser.add_argument('out_xml', metavar='out_xml', help='output network xml')

args = parser.parse_args()

xml = wx.Xml(args.in_xml)

# Create the new XML and update the nuclide data.

new_xml = wx.New_Xml()
new_xml.set_nuclide_data(xml.get_nuclide_data())

# Retrieve reactions not including beta- decays

reactions = xml.get_reaction_data(
    "[not(product = 'electron' and product[contains(.,'neutrino')])]")

# Retrieve beta- decay reactions with no delayed neutrons

reactions_no_bdn = xml.get_reaction_data("[product = 'electron' and product[contains(.,'neutrino')] and not(product = 'n')]")

# Create a dictionary.  The key is the decaying species and the value
# is the no delayed neutron reaction.

work = {}
for reaction in reactions_no_bdn:
    work[reactions_no_bdn[reaction].reactants[0]] = reactions_no_bdn[reaction]

# Retrieve beta- decay reactions with delayed neutrons

reactions_bdn_only = xml.get_reaction_data("[product = 'electron' and product[contains(.,'neutrino')] and product = 'n']")

# Loop on the reactions with delayed neutrons.  Add the rate of the
# reaction with delayed neutrons to that without.  If the reaction without
# does not exist, create it first.

for reaction in reactions_bdn_only:
    reactant = reactions_bdn_only[reaction].reactants[0]
    if reactant in work:
        work[reactant].data["rate"] += \
            reactions_bdn_only[reaction].data["rate"]
    else:
        z, a, state = xml.get_z_a_state_from_nuclide_name(reactant)
        product = xml.create_nuclide_name(z+1, a, state)
        new_reaction = wx.Reaction()
        new_reaction.reactants = [reactant]
        new_reaction.products = [product, 'electron', 'anti_neutrino_e']
        new_reaction.data = \
           {'type': 'single_rate',
            'rate': reactions_bdn_only[reaction].data['rate']}
        reactions_no_bdn[new_reaction.get_string()] = new_reaction
        work[reactant] = new_reaction
        
# Merge the sets of reactions.

reactions = {**reactions, **reactions_no_bdn}

# Write the data to a new file.

new_xml.set_reaction_data(reactions)
new_xml.write(args.out_xml)