#!/usr/bin/env python
# coding: utf-8

# In[1]:


import openmc
import matplotlib.pyplot as plt


# In[2]:


# Get ACE file and assign variables for NXS, JXS, and XSS arrays
ace_table = openmc.data.ace.get_table('dos-irdff2-1125.acef')
nxs = ace_table.nxs
jxs = ace_table.jxs
xss = ace_table.xss


# In[3]:


# Get MT values and locators for cross section blocks
lmt = jxs[3]
nmt = nxs[4]
lxs = jxs[6]
mts = xss[lmt : lmt+nmt].astype(int)
locators = xss[lxs : lxs+nmt].astype(int)
print(f'MTs: {mts}')


# In[4]:


# Create dictionary mapping MT to Tabulated1D object
cross_sections = {}
for mt, loca in zip(mts, locators):
    # Determine starting index on energy grid
    nr = int(xss[jxs[7] + loca - 1])
    if nr == 0:
        breakpoints = None
        interpolation = None
    else:
        breakpoints = xss[jxs[7] + loca : jxs[7] + loca + nr].astype(int)
        interpolation = xss[jxs[7] + loca + nr : jxs[7] + loca + 2*nr].astype(int)

    # Determine number of energies in reaction
    ne = int(xss[jxs[7] + loca + 2*nr])

    # Read reaction cross section
    start = jxs[7] + loca + 1 + 2*nr
    energy = xss[start : start + ne] * 1e6
    xs = xss[start + ne : start + 2*ne]

    cross_sections[mt] = openmc.data.Tabulated1D(energy, xs, breakpoints, interpolation)


# In[5]:


plt.loglog(cross_sections[102].x, cross_sections[102].y)
plt.xlabel('Energy [eV]')
plt.ylabel('Cross section [b]')


# Now we can create an energy function filter and use it in a tally:

# In[6]:


tally = openmc.Tally()
multiplier = openmc.EnergyFunctionFilter.from_tabulated1d(cross_sections[102])
tally.filters = [multiplier]
tally.scores = ['flux']