%matplotlib inline
import os
from pprint import pprint
import shutil
import subprocess
import urllib.request

import h5py
import numpy as np
import matplotlib.pyplot as plt

import openmc.data

"""
Explanation: In this notebook we will explore features of the Python API that allow us to import and manipulate resonance covariance data. A full description of the ENDF-VI and ENDF-VII formats can be found in the ENDF102 manual.
End of explanation
"""


# Download ENDF file
url = 'https://t2.lanl.gov/nis/data/data/ENDFB-VII.1-neutron/Gd/157'
filename, headers = urllib.request.urlretrieve(url, 'gd157.endf')

# Load into memory
gd157_endf = openmc.data.IncidentNeutron.from_endf(filename, covariance=True)
gd157_endf

"""
Explanation: ENDF: Resonance Covariance Data
Let's download the ENDF/B-VII.1 evaluation for $^{157}$Gd and load it in:
End of explanation
"""


gd157_endf.resonance_covariance.ranges[0].parameters[:5]

"""
Explanation: We can access the parameters contained within File 32 in a similar manner to the File 2 parameters from before.
End of explanation
"""


covariance = gd157_endf.resonance_covariance.ranges[0].covariance

"""
Explanation: The newly created object will contain multiple resonance regions within gd157_endf.resonance_covariance.ranges. We can access the full covariance matrix from File 32 for a given range by:
End of explanation
"""


plt.imshow(covariance,cmap='seismic',vmin=-0.08, vmax=0.08)
plt.colorbar()

"""
Explanation: This covariance matrix currently only stores the upper triangular portion as covariance matrices are symmetric. Plotting the covariance matrix:
End of explanation
"""


corr = np.zeros([len(covariance),len(covariance)])
for i in range(len(covariance)):
    for j in range(len(covariance)):
        corr[i, j]=covariance[i, j]/covariance[i, i]**(0.5)/covariance[j, j]**(0.5)
plt.imshow(corr, cmap='seismic',vmin=-1.0, vmax=1.0)
plt.colorbar()


"""
Explanation: The correlation matrix can be constructed using the covariance matrix and also give some insight into the relations among the parameters.
End of explanation
"""


rm_resonance = gd157_endf.resonances.ranges[0]
n_samples = 5
samples = gd157_endf.resonance_covariance.ranges[0].sample(n_samples)
type(samples[0])


"""
Explanation: Sampling and Reconstruction
The covariance module also has the ability to sample a new set of parameters using the covariance matrix. Currently the sampling uses numpy.multivariate_normal(). Because parameters are assumed to have a multivariate normal distribution this method doesn't not currently guarantee that sampled parameters will be positive.
End of explanation
"""


print('Sample 1')
samples[0].parameters[:5]

print('Sample 2')
samples[1].parameters[:5]

"""
Explanation: The sampling routine requires the incorporation of the openmc.data.ResonanceRange for the same resonance range object. This allows each sample itself to be its own openmc.data.ResonanceRange with a new set of parameters. Looking at some of the sampled parameters below:
End of explanation
"""


gd157_endf.resonances.ranges

energy_range = [rm_resonance.energy_min, rm_resonance.energy_max]
energies = np.logspace(np.log10(energy_range[0]),
                       np.log10(energy_range[1]), 10000)
for sample in samples:
    xs = sample.reconstruct(energies)
    elastic_xs = xs[2]
    plt.loglog(energies, elastic_xs)
plt.xlabel('Energy (eV)')
plt.ylabel('Cross section (b)')

"""
Explanation: We can reconstruct the cross section from the sampled parameters using the reconstruct method of  openmc.data.ResonanceRange. For more on reconstruction see the Nuclear Data example notebook.
End of explanation
"""


lower_bound = 2;  # inclusive
upper_bound = 2;  # inclusive
rm_res_cov_sub = gd157_endf.resonance_covariance.ranges[0].subset('J',[lower_bound,upper_bound])
rm_res_cov_sub.file2res.parameters[:5]

"""
Explanation: Subset Selection
Another capability of the covariance module is selecting a subset of the resonance parameters and the corresponding subset of the covariance matrix. We can do this by specifying the value we want to discriminate and the bounds within one energy region. Selecting only resonances with J=2:
End of explanation
"""


rm_res_cov_sub.covariance
gd157_endf.resonance_covariance.ranges[0].covariance.shape


"""
Explanation: The subset method will also store the corresponding subset of the covariance matrix
End of explanation
"""


old_n_parameters = gd157_endf.resonance_covariance.ranges[0].parameters.shape[0]
old_shape = gd157_endf.resonance_covariance.ranges[0].covariance.shape
new_n_parameters = rm_res_cov_sub.file2res.parameters.shape[0]
new_shape = rm_res_cov_sub.covariance.shape
print('Number of parameters\nOriginal: '+str(old_n_parameters)+'\nSubet: '+str(new_n_parameters)+'\nCovariance Size\nOriginal: '+str(old_shape)+'\nSubset: '+str(new_shape))


"""
Explanation: Checking the size of the new covariance matrix to be sure it was sampled properly:
End of explanation
"""


samples_sub = rm_res_cov_sub.sample(n_samples)
samples_sub[0].parameters[:5]

"""
Explanation: And finally, we can sample from the subset as well
End of explanation
"""


