import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy.signal import savgol_filter
from scipy.optimize import curve_fit
from scipy.interpolate import interp1d
from mpl_toolkits.axes_grid1 import make_axes_locatable
import glob 
from scipy.constants import c 
### supporting functions
def expfit(x, a, sig):
    return a*np.exp(-x/sig)


def grab_data_path(shot, diag):
    path = f'../data/'
    fmt = f'*_00{shot}_{diag}.*'
    return glob.glob(path+fmt)[0]

def baseline_data(shot, diag):
    fn = grab_data_path(shot, diag)
    data = pd.read_csv(fn, usecols=[3, 4], names=['time', 'volts'])
    xx, yy = data['time'].values, data['volts'].values

    idx = np.argmin(yy)

    return (xx[idx-10000:idx+20000], -yy[idx-10000:idx+20000])

### load shotsheet
file_name = f'../data/Shotsheet_forDOI.xlsx'
df = pd.read_excel(file_name, sheet_name='forpython', skiprows=[
                   1], usecols=[3, *np.arange(6, 32)])
df2 = df.loc[df['Shot No.'].isin([*np.arange(20, 97)])]
df3 = df2[['Shot No.', 'DT Material', 'DT Thickness', 'CT Material',
           'CT Thickness', 'B7 Energy', 'B8 Energy', 'B8 Pulse', 'G1 Start', 'G1 Width', 'G2 Start', 'G2 Width']]

nLY = pd.read_csv(f'../data/neutronlightyield.csv',
                  names=['Energy', 'Normalised Yield'])
nLY_interp = interp1d(nLY['Energy'].values, nLY['Normalised Yield'].values, fill_value=(
    np.nan, np.max(nLY['Normalised Yield'].values)), bounds_error=False)

### load background
shot = 23
(xx_bkg, yy_bkg) = baseline_data(shot, 'nTOF10')
yy_bkg = savgol_filter(yy_bkg, 101, 4)
bkg_tof = xx_bkg*1e9 - 200
bkg_signal = yy_bkg
### load nTOF
shot = 75
(xx, yy) = baseline_data(shot, 'nTOF10')
yy2 = savgol_filter(yy, 1001, 4)
neutron_tof = xx*1e9 - 200
neutron_tof_trimmed = neutron_tof[neutron_tof > 0]
neutron_signal = yy2
distance_to_detector = 3.6
neutron_beta = (distance_to_detector/(neutron_tof_trimmed*1e-9))/c
neutron_gamma = 1/np.sqrt(1-neutron_beta**2)
neutron_energy = np.squeeze((neutron_gamma-1)*939)

popt, pcov = curve_fit(
    expfit, neutron_tof[neutron_tof > 0], neutron_signal[neutron_tof > 0])

neutron_signal_gammasubbed = neutron_signal[neutron_tof > 0] - expfit(
    neutron_tof[neutron_tof > 0], *popt) - 0.1*bkg_signal[neutron_tof > 0]

x = neutron_energy
y = neutron_signal_gammasubbed
q = neutron_signal_gammasubbed/nLY_interp(neutron_energy)

yp = savgol_filter(y, 1000, 1)
qp = savgol_filter(q, 1000, 1)
xn = np.geomspace(0.04, 25, 501)
ypn = interp1d(x, yp, bounds_error=False)(xn)
qpn = interp1d(x, qp, bounds_error=False)(xn)

xrayimage = pd.read_csv('../simulation/200keV_x-rays/totalenergy.txt',
                        skiprows=2, delimiter='\t', names=['Volume', 'MeV', 'Joules'], nrows=3250)
im_xray = np.reshape(xrayimage['MeV'].values, (50, 65))

neutronimage = pd.read_csv('../simulation/3MeV_neutrons/totalenergy.txt',
                           skiprows=2, delimiter='\t', names=['Volume', 'MeV', 'Joules'], nrows=3250)
im_neutron = np.reshape(neutronimage['MeV'].values, (50, 65))

xpx = np.arange(65)
ypx = np.arange(50)


idx = (neutron_energy < 25) * (neutron_energy > 0.8)
popt, pcov = curve_fit(expfit, x[idx], q[idx], p0=[1, 5])


fig, axes = plt.subplots(1, 6, figsize=(
    11, 3), gridspec_kw={'height_ratios': [1], 'width_ratios': [1.2, 0, 1.1, 1.1, 0.075, 0.075]})

axes[0].plot(xn, ypn, color='tab:blue', label='Measured Signal', alpha=0.3)
axes[0].plot(xn, qpn, color='tab:blue', label='Light Yield Corrected')
axes[0].plot(xn, expfit(xn, *popt), color='k', ls='--')
axes[0].plot(xn, 1e2*np.exp(-xn/0.2), color='tab:orange',
             label='Expected X-ray Signal')

axes[0].text(1.8, 0.04, f'$k_bT=0.2$ MeV', color='tab:orange')
axes[0].text(4, 1, f'$k_bT$={popt[1]:.1f} MeV ', color='tab:blue')
axes[0].set_title('Radiation Spectra')
# axes[0].set_aspect(0.4)

axes[0].set_xlim([0.4, 20])
axes[0].set_ylim([0.01, 10])
axes[0].set_yscale('log')
axes[0].set_xscale('log')
axes[0].set_xlabel('Radiation Energy (MeV)')
axes[0].set_ylabel('Intensity ($dN_n/dE$)')
axes[0].legend(fontsize=8)

axes[1].axis('off')

im1 = axes[2].pcolormesh(ypx, xpx, im_xray.T/np.max(im_xray),
                         vmin=0.1, vmax=1, cmap='Reds')
axes[2].set_title('X-rays')
axes[2].set_xlabel('X (mm)')
# axes[2].set_yticks([-100, -50, 0, 50, 100])


im = axes[3].pcolormesh(ypx, xpx, im_neutron.T/np.max(im_neutron),
                        vmin=0.1, vmax=1, cmap='Blues')
axes[3].set_title('Neutron')
axes[3].set_xlabel('X (mm)')
axes[3].set_yticks([])
# axes[4].axis('off')
# cax0 = make_axes_locatable(axes[2]).append_axes(
#     "right", size="5%", pad=0.05)  # Adjust the padding
cbar0 = plt.colorbar(im, cax=axes[4])
cbar0.set_ticks([])
# cbar0.set_label('Energy Deposited (Arb.)')
# cax1 = make_axes_locatable(axes[1]).append_axes(
#     "right", size="5%", pad=0.05)  # Adjust the padding
cbar1 = plt.colorbar(im1, cax=axes[5])
cbar1.set_label('Energy Deposited (Arb.)')

# cbar1.set_ticks([])
# cbar1.set_ticklabels([])
# cbar1.set_label('Energy Deposited (Arb.)')
# for ax in axes:
#     ax.tick_params(axis='both', which='major')
# plt.tight_layout()

axes[0].text(0.2, 2e1, 'a)', fontweight='bold')
axes[2].text(-5, 70, 'b)', fontweight='bold')
axes[3].text(-5, 70, 'c)', fontweight='bold')

plt.show()