import os
import sys
root_folder = os.path.dirname(os.getcwd())
sys.path.append(root_folder)
import ResoFit
from ResoFit.calibration import Calibration
from ResoFit.fitresonance import FitResonance
from ResoFit.experiment import Experiment
from ResoFit._utilities import Layer
import numpy as np
import matplotlib.pyplot as plt

"""
Explanation: Tutorial
Getting started
Import library and packages
End of explanation
"""


energy_min = 7
energy_max = 150
energy_step = 0.01

"""
Explanation: Global paramters

min energy of 7 eV (has to be greater than 1 x 10<sup>-5</sup> eV)
max energy of 150 eV (has to be less than 3000 eV)
energy steps to interpolate database: 0.1 eV
End of explanation
"""


folder = 'data/_data_for_tutorial'
data_file = 'raw_data.csv'
spectra_file = 'spectra.txt'

"""
Explanation: File locations

data/* (directory to locate the file)
data_file (YOUR_DATA_FILE.txt or .csv)
spectra_file (YOUR_SPECTRA_FILE.txt or .csv)
End of explanation
"""


experiment = Experiment(data_file=data_file,
                        spectra_file=spectra_file,
                        folder=folder)

"""
Explanation: Preview data using Experiment()
Note this is meant to demenstrate the functionality of Experiment() class. These operations are not needed for general calibration and fitting.
Load the related files into Experiment()
End of explanation
"""


experiment.data

"""
Explanation: Data

Values are in transmission with image number as indexes after extracting with ImageJ
End of explanation
"""


experiment.spectra

"""
Explanation: Spectra


Column 0 time values (seconds) with image number as indexes


Column 1 neutron counts
End of explanation
"""


experiment.slice(slice_start=3, slice_end=2801, reset_index=False)

"""
Explanation: Remove unwanted data points
End of explanation
"""


experiment.x_raw(offset_us=0., source_to_detector_m=16)

"""
Explanation: Retrieve values for X-axis
In 'neutron energy'
End of explanation
"""


experiment.x_raw(angstrom=True, offset_us=0., source_to_detector_m=16)

"""
Explanation: In 'neutron wavelength'
End of explanation
"""


experiment.y_raw()

"""
Explanation: Retrieve values for Y-axis
In 'neutron attenuation'
End of explanation
"""


experiment.y_raw(transmission=True)

"""
Explanation: In 'neutron transmission'
End of explanation
"""


experiment.y_raw(baseline=True)

"""
Explanation: In 'neutron attenuation' & remove baseline
End of explanation
"""


experiment.xy_scaled(energy_min=energy_min, energy_max=energy_max, energy_step=energy_step,
                     angstrom=False, transmission=False,
                     offset_us=0, source_to_detector_m=15, 
                     baseline=True)

"""
Explanation: Retrieve interpolated values for both X-axis & Y-axis
X-axis: 'neutron energy'
Y-axis: 'neutron attenuation'
End of explanation
"""


experiment.plot_raw(x_axis='number')
plt.show()

"""
Explanation: Plot raw in various ways
Attenuation vs. image number
End of explanation
"""


experiment.plot_raw(x_axis='time')
plt.show()

experiment.plot_raw(x_axis='time', time_unit='s')
plt.show()

"""
Explanation: Attenuation vs. recorded time
Parameters need to be specified:
- time_unit from ['s', 'us', 'ns']. ('us') is default
End of explanation
"""


experiment.plot_raw(x_axis='lambda', lambda_xmax=0.12,
                   offset_us=0, source_to_detector_m=16)
plt.show()

"""
Explanation: Attenuation vs. neutron wavelength
Parameters need to be specified:
- lambda_xmax in (angstrom)
- offset_us in (us)
- source_to_detector_m in (m)
End of explanation
"""


experiment.plot_raw(x_axis='energy', energy_xmax=150,
                    offset_us=0, source_to_detector_m=16)
plt.show()

"""
Explanation: Attenuation vs. neutron energy
Parameters need to be specified:
- energy_xmax in (eV)
- offset_us in (us)
- source_to_detector_m in (m)
End of explanation
"""


experiment.plot_raw(offset_us=0, source_to_detector_m=16,
                    x_axis='energy', baseline=True, 
                    energy_xmax=150)
plt.show()

"""
Explanation: Remove baseline for plot
Note: this only changes the display of data, not the data values.
End of explanation
"""


layer_1 = 'U'
thickness_1 = 0.05 # mm
density_1 = None # g/cm^3 (if None or omitted, pure solid density will be used in fitting step)
layer_2 = 'Gd'
thickness_2 = 0.05 # mm
density_2 = None # g/cm^3 (if None or omitted, pure solid density will be used in fitting step)

"""
Explanation: Operation to experiment data
experiment.norm_to(file=background.csv)
This can be used to further normalize signal. (data divide background)
Note: this will change the data values.
Input sample info

Gd and U mixed foil
thickness neutron path within the sample in (mm)
density sample density in (g/cm<sup>3</sup>), if omitted, pure solid density will be used in fitting
repeat : reptition number if the data is summed of multiple runs (default: 1)
End of explanation
"""


layer = Layer()
layer.add_layer(layer=layer_1, thickness_mm=thickness_1, density_gcm3=density_1)
layer.add_layer(layer=layer_2, thickness_mm=thickness_2, density_gcm3=density_2)

"""
Explanation: Create sample layer

Note: Input each element as single layer.
End of explanation
"""


source_to_detector_m = 16.
offset_us = 0

"""
Explanation: Calibration
Estimated intrumental parameters

input estimated source to detector distance (m)
input estimated possible time offset in spectra file (us)

Note: a separated Experiment() will be created after the initialization of Calibration(). Calibration.slice() is needed if you would like to remove unwanted data points by indexes.
End of explanation
"""


calibration = Calibration(data_file=data_file,
                          spectra_file=spectra_file,
                          raw_layer=layer,
                          energy_min=energy_min,
                          energy_max=energy_max,
                          energy_step=energy_step,
                          folder=folder,
                          baseline=True)

"""
Explanation: Class initialization

Pass all the parameters definded above into the Calibration()
End of explanation
"""


calibration.calibrate(source_to_detector_m=source_to_detector_m,
                      offset_us=offset_us,
                      vary='all')

"""
Explanation: Equations for (time-wavelength-energy) conversion
$$E = \frac {81.787 }{ \lambda^2 }$$
$E$ : energy in (meV),
$\lambda$ : wavelength in (Å).
$$\lambda  = 0.3956\frac{t_{record} + t_{offset}}{L}$$
$t_{record}$ : recorded time in (µs),
$t_{offset}$ : recorded time offset in (µs),
$L$ : source to detector distance in (cm).
Calibrate instrumental parameters

using source_to_detector_m or offset_us or both to minimize the difference between the measured resonance signals and the simulated resonance signals from ImagingReso within the range specified in global parameters
vary can be one of ['source_to_detector', 'offset', 'all'] (default is 'all')
fitting parameters are displayed
End of explanation
"""


calibration.calibrated_offset_us

calibration.calibrated_source_to_detector_m

"""
Explanation: Retrieve calibrated parameters
End of explanation
"""


calibration.plot()

"""
Explanation: Plot calibration result

using the best fitted source_to_detector_m and offset_us to show the calibrated resonance signals from measured data and the expected resonance positions from ImagingReso
measured data before and after is ploted with raw data points instead of interpolated data points. However, the interpolated data was used during the calibration step above.

Default layout
End of explanation
"""


calibration.plot(table=False)

"""
Explanation: Hide table
End of explanation
"""


calibration.plot(table=False, grid=False)

"""
Explanation: Hide grid
End of explanation
"""


calibration.plot(table=False, grid=False, before=True)

"""
Explanation: Plot before to compare
End of explanation
"""


calibration.plot(table=False, grid=False, all_elements=True)

"""
Explanation: Show all elements
End of explanation
"""


calibration.plot(table=False, grid=False, all_isotopes=True)

"""
Explanation: Show all isotopes
End of explanation
"""


calibration.plot(table=False, grid=False, items_to_plot=['U', 'U-238', 'Gd-156'])

"""
Explanation: Show specified items

Format as list with each name in string
End of explanation
"""


calibration.plot(table=False, grid=False, items_to_plot=['U', 'U*', 'Gd-156'])

"""
Explanation: Note: 'U*' can be used to represent all the isotopes of uranium
End of explanation
"""


fit = FitResonance(spectra_file=spectra_file,
                   data_file=data_file,
                   folder=folder,
                   energy_min=energy_min,
                   energy_max=energy_max,
                   energy_step=energy_step,
                   calibrated_offset_us=calibration.calibrated_offset_us,
                   calibrated_source_to_detector_m=calibration.calibrated_source_to_detector_m,
                   baseline=True)

"""
Explanation: Fit resonances
Class initialization

Pass all the parameters definded and calibrated into the FitResonance()
End of explanation
"""


fit_result = fit.fit(layer, vary='density')

"""
Explanation: Fitting equations
Beer-Lambert Law:
$$T\left( E \right) =\frac { I\left( E \right)  }{ { I }{ 0 }\left( E \right)  } =exp\left[ -\sum\nolimits_i { { N }{ i }{ d }{ i } } \sum\nolimits_j { { \sigma  }{ ij }\left( E \right) { A }_{ ij } }  \right]$$
${ N }_{ i }$ : number of atoms per unit volume of element $i$,
${ d }_{ i }$ : effective thickness along the neutron path of element $i$,
${ \sigma  }_{ ij }\left( E \right)$  : energy-dependent neutron total cross-section for the isotope $j$ of element $i$,
${ A }_{ ij }$ : abundance for the isotope $j$ of element $i$.
$${N_i} = {N_A}{C_i} = \frac { {N_A}{\rho_i}} {\sum\nolimits_j {m_{ij}{A_{ij}}}}$$
${N_A}$ : Avogadro’s number,
${C_i}$ : molar concentration of element $i$,
${\rho_i}$ : density of the element $i$,
$m_{ij}$ : atomic mass values for the isotope $j$ of element $i$.
How to fit the resonance signals

using thickness (mm) or density (g/cm<sup>3</sup>) to minimize the difference between the measured resonance signals and the simulated resonance signals from ImagingReso within the range specified in global parameters
vary can be one of ['thickness', 'density'] (default is 'density')
fitting parameters are displayed
End of explanation
"""


fit.molar_conc()

"""
Explanation: Output fitted result in molar concentration

unit: mol/cm<sup>3</sup>
End of explanation
"""


fit.plot()

"""
Explanation: Plot fitting result

using the best fitted density to show the measured resonance signals and the fitted resonance signals from ImagingReso
measured data before and after is ploted with raw data points instead of interpolated data points. However, the interpolated data was used during the fitting step above.

Default layout
End of explanation
"""


fit.plot(error=False)

"""
Explanation: Hide difference plot
End of explanation
"""


fit.plot(table=False)

"""
Explanation: Hide table
End of explanation
"""


fit.plot(table=False, grid=False, before=True)

"""
Explanation: Hide grid & show before
End of explanation
"""


fit.plot(table=False, grid=False, all_elements=True)

"""
Explanation: Show all elements
End of explanation
"""


fit.plot(table=False, grid=False, all_isotopes=True)

"""
Explanation: Show all isotopes
End of explanation
"""


fit.plot(table=False, grid=False, items_to_plot=['U', 'U-238', 'Gd-156'])

fit.plot(table=False, grid=False, items_to_plot=['U', 'U*', 'Gd-156'])

"""
Explanation: Show specified items

Format as list with each name in string
End of explanation
"""


fit.fit_iso(layer='U')

fit.molar_conc()

fit.plot()

"""
Explanation: Fit isotopic ratios

specify the element you would like to perform isotopic (by at.%) fitting.

Note: Processing time of this step might be long.
End of explanation
"""


