import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit
import pandas as pd 
from PIL import Image
import glob


# generate meta data
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']]

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

def open_tiff(shot, diag):
    fn = grab_data_path(shot, diag)
    return np.array(Image.open(fn))

def im_roi(image, roi):
    x, y, w, h = roi
    return image[y:y+h, x:x+w]

def mean_roi(shot, diag, roi=[400, 200, 1200, 800]):
    im = open_tiff(shot, diag)
    return np.mean(im_roi(im, roi), axis=(0, 1))

def std_roi(shot, diag, roi=[400, 200, 1200, 800]):
    im = open_tiff(shot, diag)
    return np.std(im_roi(im, roi), axis=(0, 1))

# Generate Data from RAW

flags = ['CD25']
shots = [[*np.arange(67, 75), *np.arange(75, 82)]]

s = shots[0]
energy = np.zeros(len(s))
g1start = np.zeros((len(s)))
g1width = np.zeros((len(s)))
g2start = np.zeros((len(s)))
g2width = np.zeros((len(s)))
neutron_yield = np.zeros((len(s), 4))
for n, shot in enumerate(s):
    row = df3[df3['Shot No.'] == shot]
    energy[n] = row['B8 Energy'].values[0]
    g1start[n] = row['G1 Start'].values[0]
    g1width[n] = row['G1 Width'].values[0]
    g2start[n] = row['G2 Start'].values[0]
    g2width[n] = row['G2 Width'].values[0]

    neutron_yield[n, 0] = mean_roi(shot, 'GatedImagerEarly')
    neutron_yield[n, 1] = mean_roi(shot, 'GatedImagerLate')
    neutron_yield[n, 2] = std_roi(shot, 'GatedImagerEarly')
    neutron_yield[n, 3] = std_roi(shot, 'GatedImagerLate')

offset = np.min([*g1start, *g2start])
background = 0.2
x = g2start-offset + (g2width/2)
y = (neutron_yield[:, 1]/neutron_yield[:, 0]) - background

# Remove NaN values from the data
mask = ~np.isnan(y)
x = x[mask]
y = y[mask]

# Define the Gaussian function
def fn(var, arr):
    return np.argmin((var-arr)**2)

def gaussian(x, amplitude, mean, stddev):
    return amplitude * np.exp(-((x - mean) / stddev)**2 / 2)

def halfgaussian(x, amplitude, stddev):
    mean = 101
    g1 = gaussian(x, amplitude, mean, stddev)
    edge = np.zeros_like(x)
    edge[:fn(mean, x)] += 1
    return g1 * edge

# Initial guesses for the parameters
initial_guess1 = [1, 10, np.sqrt(7**2+1.6**2)]

# Fit the first Gaussian to x <= 25
params1, covariance1 = curve_fit(
    gaussian, x[x <= 25], y[x <= 25], p0=initial_guess1)
errs1 = np.sqrt(np.diag(covariance1))

# Initial guesses for the parameters
initial_guess2 = [0.6, 30]

# Fit the second Gaussian to x >= 25
params2, covariance2 = curve_fit(halfgaussian,
                                 [*np.linspace(10, 20, 6), *x[x >= 30]],
                                 [*np.zeros(6), *y[x >= 30]],
                                 p0=initial_guess2)
errs2 = np.sqrt(np.diag(covariance2))


new_x = np.linspace(0, 100, 101)
fit1 = gaussian(new_x, *params1)
fit1_l = gaussian(new_x, *(params1 - errs1))
fit1_u = gaussian(new_x, *(params1 + errs1))


fit2 = halfgaussian(new_x, *params2)
fit2_l = halfgaussian(new_x, *(params2 - errs2))
fit2_u = halfgaussian(new_x, *(params2 + errs2))

# Plot the data and the fitted curves
plt.figure(figsize=(6, 4))
plt.errorbar(x=g2start-offset + (g2width/2),
             y=(neutron_yield[:, 1]/neutron_yield[:, 0]) - background,
             yerr=0.5*neutron_yield[:, 3]/neutron_yield[:, 0],
             fmt='o', color='k', label='Data')
plt.plot(new_x, fit1, color='r', label=r'Fit ($\tau<25ns$)', ls='--')
plt.plot(new_x, fit2, color='b', label=r'Fit ($\tau>25ns$)', ls='--')

plt.fill_between(new_x, fit1_u, fit1_l, color='r',
                 lw=0, alpha=0.1, label='Uncertainty')

fwhm1 = 2 * np.sqrt(2 * np.log(2)) * params1[2]
fwhm1_err = 2 * np.sqrt(2 * np.log(2)) * errs1[2]
fwhm2 = 2 * np.sqrt(2 * np.log(2)) * params2[1]

plt.ylim([-0.05, 1.15])
plt.text(params1[1] + fwhm1 / 2 + 1, 1,
         f'$t_{{{1/2}}} = {fwhm1:.0f} \pm {fwhm1_err:.0f}$ ns')
plt.text(params1[1] + fwhm1 / 2 + 1, 0.9,
         f'   $\sigma = {params1[2]:.0f} \pm {errs1[2]:.0f}$ ns')

plt.xlabel('Gate time (ns)')
plt.ylabel('Ratio of signal (arb.)')
plt.legend(bbox_to_anchor=(1.05, 1.15), ncol=4, fontsize=9)
plt.xlim([0, 97])
perr = np.sqrt(np.diag(covariance1))

plt.show()
