import openmc
import openmc_source_plotter  # extends openmc.Source with plotting functions
import pint

total_energy = pint.Quantity(1, 'GJ')
energy_per_dd = pint.Quantity(12.5, 'MeV/particle')
energy_per_dt = pint.Quantity(17.6, 'MeV/particle')
energy_per_fission = pint.Quantity(200., 'MeV/particle')

my_dt_source = openmc.Source()
my_dt_source.energy = openmc.stats.Muir(e0=14080000.0, m_rat=5.0, kt=20000.0)
my_dt_source.strength=(total_energy/energy_per_dt).to('particle').magnitude
print(my_dt_source.strength)


my_dd_source = openmc.Source()
my_dd_source.energy = openmc.stats.Muir(e0=2080000.0, m_rat=2.0, kt=20000.0)
my_dd_source.strength=(total_energy/energy_per_dd).to('particle').magnitude
print(my_dd_source.strength)

my_fission_source = openmc.Source()
my_fission_source.energy = openmc.stats.Watt(a=988000.0, b=2.249e-06)
my_fission_source.strength=(total_energy/energy_per_fission).to('particle').magnitude*2.5  # 2.5 neutrons per fission on average
print(my_fission_source.strength)

figure1 = my_dd_source.plot_source_energy(n_samples=1000000, name=f'DD {my_dd_source.strength:.2e} neutrons per second')
figure2 = my_dt_source.plot_source_energy(figure=figure1, n_samples=1000000, name=f'DT {my_dt_source.strength:.2e}  neutrons per second')
figure3 = my_fission_source.plot_source_energy(figure=figure2, n_samples=1000000, name=f'Fission {my_fission_source.strength:.2e}  neutrons per second')

figure3.layout.yaxis['title']['text'] = 'neutrons per second'
figure3.layout.title = 'Energy distribution of neutrons from a 1GW thermal energy fission reactor, DD fusion reactor or DT fusion'

figure3.write_html('neutrons_from_different_sources.html')
figure3.show()
# hand calc for number of neutrons from each source
((1e9 * 6.242e+18)/200e6)*2.5 # Fission
((1e9 * 6.242e+18)/17.6e6) # DT Fusion
((1e9 * 6.242e+18)/12.5e6) # DD Fusion