from pathlib import Path
import os
import openmc
import openmc.deplete
import matplotlib.pyplot as plt
from openmc.data import NATURAL_ABUNDANCE
from openmc.deplete import REACTIONS
import pandas as pd 
import numpy as np
import tqdm 
# PR to change this to DADZ
from openmc.deplete import REACTIONS
from openmc.data import REACTION_MT
from openmc.data import REACTION_NAME

chain_file = Path(os.environ['OPENMC_CHAIN_FILE'])
chain = openmc.deplete.Chain.from_xml(chain_file)

cross_section_file = Path(os.environ['OPENMC_CROSS_SECTIONS'])
cross_section = openmc.data.DataLibrary().from_xml(cross_section_file)


def find_threshold_energy(cross_sections, mt):
    if mt not in cross_sections.reactions.keys():
        print(f'{cross_sections.name} mt {mt} not found')
        return False
    xs = cross_sections.reactions[mt].xs['1200K']
    energy = xs.x
    barns = xs.y
    index_of_threshold = np.where(barns > 0)[0]
    if len(index_of_threshold) > 0:
        first_index_over_threshold = index_of_threshold[0]
        energy_threshold = energy[first_index_over_threshold]
        return energy_threshold
    else:
        return 0.0

transmutation_reaction_names = list(REACTIONS.keys())

nuclide_names = []
reaction_types = []
reaction_targets = []
threshold_energies_needed = []

for nuclide in tqdm.tqdm(chain.nuclides): # 3820 nuclides
    if nuclide.name in NATURAL_ABUNDANCE.keys(): # 289 nuclides
        filename = cross_section.get_by_material(nuclide.name)['path']
        cross_sections = openmc.data.IncidentNeutron.from_hdf5(filename)
        for reaction in nuclide.reactions:
            # check if reaction results in a change to the proton or neutorn numbers or meta stable
            # this looks like is a duplicate check as we already check that a stable nuc is the target and unstable is the product. However it removes elastic heating dpa and other scores
            if reaction.type in transmutation_reaction_names:
                # check if target is unstable
                if reaction.target not in NATURAL_ABUNDANCE.keys():
                    # ruling out fission as they produce a wide range of products
                    if reaction.type != 'fission':
                        mt= REACTION_MT[reaction.type]
                        threshold_energy_needed = find_threshold_energy(cross_sections, mt)
                        if threshold_energy_needed: # can be false if mt not found
                            print(nuclide.name, reaction.type, reaction.target, threshold_energy_needed)
                            nuclide_names.append(nuclide.name)
                            reaction_types.append(reaction.type)
                            reaction_targets.append(reaction.target)
                            threshold_energies_needed.append(threshold_energy_needed)


dict = {'nuclide_names': nuclide_names, 'reaction_types': reaction_types, 'reaction_targets': reaction_targets, 'threshold_energies_needed': threshold_energies_needed} 
df = pd.DataFrame(dict)


all_products=df['reaction_targets'].unique()

threshold_energy = 5e6

options = []
unique_reactions = {}
for product in all_products[:40]:
    tmp_df = df.loc[df['reaction_targets'] == product]
    options_below_threshold = tmp_df.loc[tmp_df['threshold_energies_needed'] < threshold_energy]
    options_above_threshold = tmp_df.loc[tmp_df['threshold_energies_needed'] > threshold_energy]
    if len(options_above_threshold) >= 1 and len(options_below_threshold) == 0:
        options.append(tmp_df['reaction_targets'].values[0])
        all_reaction_paths = []
        for index, row in options_above_threshold.iterrows():
            all_reaction_paths.append(row.reaction_types)
        unique_reactions[row.nuclide_names] = all_reaction_paths

import matplotlib.pyplot as plt
openmc.plotter.plot_xs(
    reactions = unique_reactions
)
plt.xscale('linear')
plt.title('Isotopes that can only be produced with 5MeV neutrons')
plt.ylim(10e-9, 10e0)
plt.savefig(f'isotopes_{threshold_energy}.png')