# This example plots the probability of scattering angle for different incident neutron energies.

import openmc
import matplotlib.pyplot as plt
from pprint import pprint

# change this path to point to your nuclide of choice
nuc_path = "/home/jshimwell/ENDF-B-VIII.0-NNDC/h5_files/neutron/C12.h5"

gold = openmc.data.IncidentNeutron.from_hdf5(nuc_path)

# prints all the nuclear reaactions avaialble, MT numbers
print('\n All the reactions available')
pprint(list(gold.reactions.values()))

# selects the elastic cross section (MT number 2) 
gold_elastic = gold[2]

# prints out the products of the reaction
print('\n Products of the elastic reaction', gold_elastic.products)

neutron_out = gold_elastic.products[0]

distribution = neutron_out.distribution[0]

neutron_energy = distribution.angle.energy

fig, ax = plt.subplots()

# This loops through the incident neutron energies and the angular distributions
# This loop also uses some python list slicing syntax
# As the nuclear data goes to high energies (150MeV) we skip the last 15 entries of the list with the [:-15] part
# As there are a lot of data points for lots of energies we are plotting every 50th neutron energy with [:::50] part
# mu is the cosine of the scattering angle
for energy, angle in zip(distribution.angle.energy[:-15][::50], distribution.angle.mu[:-15][::50]):
    ax.plot(angle.x, angle.p, label=f"neutron energy={energy/1e6}MeV")


plt.xlabel("Cosine of scattering angle")
plt.ylabel("Probability")
plt.title("Neutron energy angle distribution of C12")
plt.legend(bbox_to_anchor=(1.01, 0.9))
plt.savefig("angle_energy_c12.png", bbox_inches="tight", dpi=400)
print('written angle_energy_c12.png')

# This example plots the exiting neutron energy distribution for different incident neutron energies.

import openmc
import matplotlib.pyplot as plt
from pprint import pprint

# change this path to point to your nuclide of choice
nuc_path = "/home/jshimwell/ENDF-B-VIII.0-NNDC/h5_files/neutron/Be9.h5"

beryllium = openmc.data.IncidentNeutron.from_hdf5(nuc_path)

# prints all the nuclear reactions available, MT numbers
print('\n All the reactions available')
pprint(list(beryllium.reactions.values()))

# selects the neuron multiplication cross section (MT number 16) 
beryllium_elastic = beryllium[16]

# prints out the products of the reaction
print('\n Products of the elastic reaction', beryllium_elastic.products)

neutron_out = beryllium_elastic.products[0]

distribution = neutron_out.distribution[0]

incident_neutron_energy = distribution.energy
outgoing_neutron_energy = distribution.energy_out

fig, ax = plt.subplots()

# # This loops through the incident neutron energies and outgoing neutron energies
# # This loop also uses some python list slicing syntax
# # As there are a lot of data points for lots of energies we are plotting every 4th neutron energy with [::4] part
for incident_energy, outgoing_energy in zip(incident_neutron_energy[::4], outgoing_neutron_energy[::4]):
    ax.semilogy(outgoing_energy.x/1e6, outgoing_energy.p, label=f"incident neutron energy={incident_energy/1e6}MeV")


plt.xlabel("Outgoing energy [MeV]")
plt.ylabel("Probability/eV")
plt.title("Neutron energy distribution of Be9(n,2n) reactions")
plt.legend(bbox_to_anchor=(1.01, 0.9))
plt.savefig("angle_energy_be9.png", bbox_inches="tight", dpi=400)
print('written angle_energy_be9.png')