# 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 sys
import argparse
import wnutils.xml as wx

# Get arguments

parser = argparse.ArgumentParser(description='Select out data satisifying nuclide xpath from a Webnucleo xml file of any type')

parser.add_argument('input_xml', metavar='input_xml', type=str,
                    help='the input xml file')

parser.add_argument('nuc_xpath', metavar='nuc_xpath', type=str,
                    help='the nuclide xpath expression to apply')

parser.add_argument('output_xml', metavar='output_xml', type=str,
                    help='the output xml file')

# Parse data.

args = parser.parse_args()

xml = wx.Xml(args.input_xml)

# Check input.

allowed_types = {'nuclear_data', 'nuclear_network', 'libnucnet_input'}

if xml.get_type() not in allowed_types:
    print("Invalid input xml type for routine.")
    exit()

# Create output.

new_xml = wx.New_Xml(xml_type=xml.get_type())

nuclides = xml.get_nuclide_data(args.nuc_xpath)

new_xml.set_nuclide_data(nuclides)

if xml.get_type() != "nuclear_data":

    new_reactions = {}

    reactions = xml.get_reaction_data()

    new_reactions = {}

    other_reaction_elements = ['gamma', 'electron', 'positron',
                               'neutrino_e', 'anti-neutrino_e',
                               'neutrino_mu', 'anti-neutrino_mu',
                               'neutrino_tau', 'anti-neutrino_tau']

    for reaction in reactions:
        b_valid = True
        for reactant in reactions[reaction].reactants:
            if reactant not in nuclides and \
               reactant not in other_reaction_elements:
                   b_valid = False
        for product in reactions[reaction].products:
            if product not in nuclides and \
               product not in other_reaction_elements:
                   b_valid = False
        if b_valid:
            new_reactions[reaction] = reactions[reaction]

    new_xml.set_reaction_data(new_reactions)



    new_xml.set_reaction_data(new_reactions)

    if xml.get_type() != "nuclear_network":
        zones = xml.get_zone_data()

        new_zones = {}

        for zone in zones:
            props = zones[zone]['properties']
            x = {}
            for sp in zones[zone]['mass fractions']:
                if sp[0] in nuclides:
                    x[sp] = zones[zone]['mass fractions'][sp]
            new_zones[zone] = {'properties': props, 'mass fractions': x}

        new_xml.set_zone_data(new_zones)

new_xml.write(args.output_xml)