import json
import numpy as np
import hyperspy.api as hs

indir = '/***/'

si_1 = hs.load(indir + '***.dm3')
si_2 = hs.load(indir + '***.dm3')

si_metadata_1 = json.load(open(indir + '***.json', 'r'))
si_metadata_2 = json.load(open(indir + '***.json', 'r'))

si_zlp_1 = si_1[1]
si_inelastic_1 = si_1[2]
si_ssd_1 = si_1[3]

si_zlp_2 = si_2[1]
si_inelastic_2 = si_2[2]
si_ssd_2 = si_2[3]

eels_cal_1 = si_metadata_1['spatial_calibrations'][0]['scale']
eels_offset_1 = si_metadata_1['spatial_calibrations'][0]['offset'] - ****

eels_cal_2 = si_metadata_2['spatial_calibrations'][0]['scale']
eels_offset_2 = si_metadata_2['spatial_calibrations'][0]['offset'] - ****

for signal in [si_zlp_1, si_inelastic_1, si_ssd_1]:
    signal.axes_manager[0].offset = eels_offset_1
    signal.set_microscope_parameters(beam_energy=**, convergence_angle=**, collection_angle=**)

for signal in [si_zlp_2, si_inelastic_2, si_ssd_2]:
    signal.axes_manager[0].offset = eels_offset_2
    signal.set_microscope_parameters(beam_energy=**, convergence_angle=**, collection_angle=**)

x_1 = np.linspace(
    eels_offset_1,
    eels_offset_1 + np.shape(si_ssd_1)[0] * eels_cal_1,
    np.shape(si_ssd_1)[0]
)
x_2 = np.linspace(
    eels_offset_2,
    eels_offset_2 + np.shape(si_ssd_2)[0] * eels_cal_2,
    np.shape(si_ssd_2)[0]
)

si_ssd_1.smooth_savitzky_golay(polynomial_order=2, window_length=9, differential_order=0)
si_inelastic_1.smooth_savitzky_golay(polynomial_order=2, window_length=9, differential_order=0)
si_ssd_2.smooth_savitzky_golay(polynomial_order=2, window_length=9, differential_order=0)
si_inelastic_2.smooth_savitzky_golay(polynomial_order=2, window_length=9, differential_order=0)

### Calculating thickness (t1,t2) of the MNP for points (18,18) and (47,39)
### The center of this sphere is at (34,28.5), calculated from HAADF (SubScan) image.
### D is the diameter. The value multiplied at the end (6.7136) is spatial calibration.

t1 = 2 * np.sqrt(24.5**2 - ((34 - 18)**2 + (28.5 - 18)**2)) * 6.7136
t2 = 2 * np.sqrt(24.5**2 - ((34 - 47)**2 + (28.5 - 39)**2)) * 6.7136
D = 2 * np.sqrt(24.5**2 - ((34 - 34)**2 + (28.5 - 28.5)**2)) * 6.7136

epsilon_1, output_1 = si_ssd_1.kramers_kronig_analysis(
    zlp=si_zlp_1,
    iterations=50,
    n=None,
    t=t1,
    delta=0.05,
    full_output=True
)

x_out_1 = np.linspace(
    0,
    np.shape(output_1['surface plasmon estimation'].data)[0] * eels_cal_1,
    np.shape(output_1['surface plasmon estimation'].data)[0]
)

epsilon_2, output_2 = si_ssd_2.kramers_kronig_analysis(
    zlp=si_zlp_2,
    iterations=50,
    n=None,
    t=t2,
    delta=0.05,
    full_output=True
)

x_out_2 = np.linspace(
    0,
    np.shape(output_2['surface plasmon estimation'].data)[0] * eels_cal_2,
    np.shape(output_2['surface plasmon estimation'].data)[0]
)

n_1 = np.sqrt((np.sqrt(epsilon_1.real * epsilon_1.real + epsilon_1.imag * epsilon_1.imag) + epsilon_1.real) / 2)
k_1 = np.sqrt((np.sqrt(epsilon_1.real * epsilon_1.real + epsilon_1.imag * epsilon_1.imag) - epsilon_1.real) / 2)

n_2 = np.sqrt((np.sqrt(epsilon_2.real * epsilon_2.real + epsilon_2.imag * epsilon_2.imag) + epsilon_2.real) / 2)
k_2 = np.sqrt((np.sqrt(epsilon_2.real * epsilon_2.real + epsilon_2.imag * epsilon_2.imag) - epsilon_2.real) / 2)

np.savetxt(
    '***.csv',
    np.column_stack((x_out_1, n_1.data, k_1.data)),
    delimiter=',',
    header='Energy (eV), n, k',
    comments=''
)

np.savetxt(
    '***.csv',
    np.column_stack((x_out_2, n_2.data, k_2.data)),
    delimiter=',',
    header='Energy (eV), n, k',
    comments=''
)

print('t1 =', t1)
print('t2 =', t2)
print('D =', D)
print('SI_SSD_1 shape =', np.shape(si_ssd_1.data))
print('SI_SSD_2 shape =', np.shape(si_ssd_2.data))
print('n_1 shape =', np.shape(n_1))
print('n_2 shape =', np.shape(n_2))
print('Done.')