# Copyright 2021 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='Remove species from webnucleo XML.')

parser.add_argument('in_xml', metavar='in_xml', help='input xml')
parser.add_argument('out_xml', metavar='out_xml', help='output xml')
parser.add_argument('--species', metavar='species', nargs='*',
                    help='species to remove')

args = parser.parse_args()

xml = wx.Xml(args.in_xml)

my_type = xml.get_type()

new_xml = wx.New_Xml(xml_type = my_type)

# Remove nuclides

if my_type == 'nuclear_data' or my_type == 'nuclear_network' or \
    my_type == 'libnucnet_input':

    nuclides = xml.get_nuclide_data()

    for s in args.species:
        if s in nuclides:
            nuclides.pop(s)

    new_xml.set_nuclide_data(nuclides)

# Remove reactions

if my_type == 'reaction_data' or my_type == 'nuclear_network' or\
   my_type == 'libnucnet_input':

    reactions = xml.get_reaction_data()

    reac_xpath = '[(reactant = \'' + args.species[0] + \
                   '\' or product = \'' + args.species[0] + '\')'

    for i in range(1, len(args.species)):
        reac_xpath += ' or (reactant = \'' + args.species[i] + \
                      '\' or product = \'' + args.species[i] +'\')'

    reac_xpath += ']'

    reactions_to_remove = xml.get_reaction_data(reac_xpath)

    for s in reactions_to_remove:
        if s in reactions:
            reactions.pop(s)

    new_xml.set_reaction_data(reactions)

# Remove nuclides from zone data

if my_type == 'zone_data' or my_type == 'libnucnet_input':
    zones = xml.get_zone_data()

    t_list = []
    for s in args.species:
        z, a, state = xml.get_z_a_state_from_nuclide_name(s)
        t_list.append((s, z, a))

    for zone in zones:
        for t in t_list:
            if t in zones[zone]['mass fractions']:
                zones[zone]['mass fractions'].pop(t)

    new_xml.set_zone_data(zones)

# Write out new XML.

new_xml.write(args.out_xml)