import numpy as np
import matplotlib.pyplot as plt
import scipy.odr as sodr
import scipy.stats as scs
from hist_parse import histParse
from sigfig import round as sigfigRound

def round(value, sigma, notation_scientific = False):
    return sigfigRound(float(value), float(sigma), cutoff = 29, notation = 'sci' if notation_scientific else 'std')

rate_empty_coincident = 2.0481
sigma_rate_empty_coincident = 0.0010
rate_empty_single = 4388.7
sigma_rate_empty_single = 0.4

rate_max_coincident = 58.902
sigma_rate_max_coincident = 0.005
rate_max_single = 5952.3
sigma_rate_max_single = 0.5

E_0 = 511 #kev

log_flag = False

def gauss(B, x):
    # B[0] = my ; B[1] = sigma
    return B[2]/np.sqrt(2*np.pi*B[1]**2) * np.exp(-(x-B[0])**2/(2*B[1]**2))

def gaussian_fit(xdata, ydata, xerr, yerr, my0 = 0, sigma0 = 1, A0 = 1):
    model = sodr.Model(gauss)
    data = sodr.RealData(xdata, ydata, sx = xerr, sy = yerr)
    odr = sodr.ODR(data, model, beta0=[my0, sigma0, A0])
    return odr.run()

# data collection
data_empty_single = histParse('Energy_Single_Empty.log')
data_empty_coincident = histParse('Energy_Coincidence_Empty.log')
data_max_single = histParse('Energy_Single_Maximum_2_2.log')
data_max_coincident = histParse('Energy_Coincidence_Maximum_2_2.log')

# general characteristics
sigma_data_empty_single = np.sqrt(data_empty_single[:, 1])
sigma_data_empty_coincident = np.sqrt(data_empty_coincident[:, 1])
sigma_data_max_single = np.sqrt(data_max_single[:, 1])
sigma_data_max_coincident = np.sqrt(data_max_coincident[:, 1])

n_empty_single = np.sum(data_empty_single[:, 1])
n_max_single = np.sum(data_max_single[:, 1])
n_empty_coincident = np.sum(data_empty_coincident[:, 1])
n_max_coincident = np.sum(data_max_coincident[:, 1])

print(f'\n\t{n_empty_single=}')
print(f'\t{n_max_single=}')
print(f'\t{n_empty_coincident=}')
print(f'\t{n_max_coincident=}\n')

t_empty_single = n_empty_single / rate_empty_single
t_max_single = n_max_single / rate_max_single
t_empty_coincident = n_empty_coincident / rate_empty_coincident
t_max_coincident = n_max_coincident / rate_max_coincident

sigma_t_syst_empty_single = n_empty_single * sigma_rate_empty_single / rate_empty_single**2
sigma_t_syst_max_single = n_max_single * sigma_rate_max_single / rate_max_single**2
sigma_t_syst_empty_coincident = n_empty_coincident * sigma_rate_empty_coincident / rate_empty_coincident**2
sigma_t_syst_max_coincident = n_max_coincident * sigma_rate_max_coincident / rate_max_coincident**2

# calculate rates
data_rate_empty_single = data_empty_single.copy()
data_rate_max_single = data_max_single.copy()
data_rate_empty_coincident = data_empty_coincident.copy()
data_rate_max_coincident = data_max_coincident.copy()

data_rate_empty_single[:, 1] = data_empty_single[:, 1] / t_empty_single
data_rate_max_single[:, 1] = data_max_single[:, 1] / t_max_single
data_rate_empty_coincident[:, 1] = data_empty_coincident[:, 1] / t_empty_coincident
data_rate_max_coincident[:, 1] = data_max_coincident[:, 1] / t_max_coincident

sigma_data_rate_empty_single = sigma_data_empty_single / t_empty_single
sigma_data_rate_max_single = sigma_data_max_single / t_max_single
sigma_data_rate_empty_coincident = sigma_data_empty_coincident / t_empty_coincident
sigma_data_rate_max_coincident = sigma_data_max_coincident / t_max_coincident

# denoise rates
data_rate_max_single_denoised = data_rate_max_single.copy()
data_rate_max_coincident_denoised = data_rate_max_coincident.copy()

data_rate_max_single_denoised[:, 1] = data_rate_max_single[:, 1] - data_rate_empty_single[:, 1]
data_rate_max_coincident_denoised[:, 1] = data_rate_max_coincident[:, 1] - data_rate_empty_coincident[:, 1]

sigma_data_rate_max_single_denoised = np.sqrt(sigma_data_rate_max_single**2 + sigma_data_rate_empty_single**2)
sigma_data_rate_max_coincident_denoised = np.sqrt(sigma_data_rate_max_coincident**2 + sigma_data_rate_empty_coincident**2)

#'''
my0 = 480
sigma0 = 90
i_min = np.max(np.where(data_rate_max_single_denoised[:, 0] < 425))
i_max = np.max(np.where(data_rate_max_single_denoised[:, 0] < 520))
gaussian = gaussian_fit(data_rate_max_single_denoised[i_min:i_max:, 0], data_rate_max_single_denoised[i_min:i_max:, 1], np.ones_like(data_rate_max_single_denoised[i_min:i_max:, 0])*2/np.sqrt(12), sigma_data_rate_max_single_denoised[i_min:i_max:], my0, sigma0, np.max(data_rate_max_single_denoised[i_min:i_max:, 1])*np.sqrt(2*np.pi*sigma0**2))
print('////////// singles //////////')
print(f'{gaussian.beta=}')
print(f'errors: {np.sqrt(gaussian.cov_beta[0, 0])}, {np.sqrt(gaussian.cov_beta[1, 1])}, {np.sqrt(gaussian.cov_beta[2, 2])}')
chi_sq = np.sum(((data_rate_max_single_denoised[i_min:i_max:, 1] - gauss(gaussian.beta, data_rate_max_single_denoised[i_min:i_max:, 0]))/sigma_data_rate_max_single_denoised[i_min:i_max:])**2)
dof = i_max - i_min
print(f'chiq/dof: {chi_sq/dof}')
p = 1 - scs.chi2.cdf(chi_sq, dof)
print(f'{p=}')

scale = E_0 / gaussian.beta[0]
print(f'\n\tscaling factor singles: {scale}\n')
data_rate_max_single_denoised *= scale
sigma_data_rate_max_single_denoised *= scale

gaussian = gaussian_fit(data_rate_max_single_denoised[i_min:i_max:, 0], data_rate_max_single_denoised[i_min:i_max:, 1], np.ones_like(data_rate_max_single_denoised[i_min:i_max:, 0])*2/np.sqrt(12), sigma_data_rate_max_single_denoised[i_min:i_max:], my0, sigma0, np.max(data_rate_max_single_denoised[i_min:i_max:, 1])*np.sqrt(2*np.pi*sigma0**2))
print('////////// singles corrected //////////')
print(f'{gaussian.beta=}')
print(f'errors: {np.sqrt(gaussian.cov_beta[0, 0])}, {np.sqrt(gaussian.cov_beta[1, 1])}, {np.sqrt(gaussian.cov_beta[2, 2])}')
chi_sq = np.sum(((data_rate_max_single_denoised[i_min:i_max:, 1] - gauss(gaussian.beta, data_rate_max_single_denoised[i_min:i_max:, 0]))/sigma_data_rate_max_single_denoised[i_min:i_max:])**2)
dof = i_max - i_min
print(f'chiq/dof: {chi_sq/dof}')
p = 1 - scs.chi2.cdf(chi_sq, dof)
print(f'{p=}')

FWHM = 2 * np.sqrt(2*np.log(2)) * gaussian.beta[1]
sigma_FWHM = 2 * np.sqrt(2*np.log(2)) * np.sqrt(gaussian.cov_beta[1, 1])
print(f'\n\tFWHM singles: {round(FWHM, sigma_FWHM)}\n')

plt.errorbar(data_rate_max_single_denoised[:, 0], data_rate_max_single_denoised[:, 1], label='data', linestyle = '', marker = ',', xerr = 2/np.sqrt(12), yerr = sigma_data_rate_max_single_denoised)
#plt.plot(data_rate_max_single_denoised[i_min:i_max:, 0], gauss(gaussian.beta, data_rate_max_single_denoised[i_min:i_max:, 0]), label='fit', color='C1')
plt.vlines([511, 1275] , np.min(data_rate_max_single_denoised[:, 1]), np.max(data_rate_max_single_denoised[:, 1]), color = 'C1', linestyles = 'dashed')
if log_flag:
    plt.yscale('log')
plt.grid()
plt.legend()
plt.ylabel('# Counts per second')
plt.xlabel('Energy [keV]')
plt.title('Energy spectrum singles (maximum, corrected)')
plt.tight_layout()
plt.savefig(f'./output/max_rate_single_corrected_errorbar.pdf')
plt.cla()

fig, ax = plt.subplots(nrows = 2, sharex = True)
ax[0].plot(data_rate_max_single_denoised[i_min:i_max:, 0], gauss(gaussian.beta, data_rate_max_single_denoised[i_min:i_max:, 0]), label='fit', color='C1')
ax[0].errorbar(data_rate_max_single_denoised[i_min:i_max:, 0], data_rate_max_single_denoised[i_min:i_max:, 1], label='data', linestyle = '', marker = ',', xerr = 2/np.sqrt(12), yerr = sigma_data_rate_max_single_denoised[i_min:i_max:])
ax[1].errorbar(data_rate_max_single_denoised[i_min:i_max:, 0], data_rate_max_single_denoised[i_min:i_max:, 1] - gauss(gaussian.beta, data_rate_max_single_denoised[i_min:i_max:, 0]), yerr = sigma_data_rate_max_single_denoised[i_min:i_max:], label='residuals', linestyle = '', marker = '.')
ax[1].hlines(0, data_rate_max_single_denoised[i_min, 0], data_rate_max_single_denoised[i_max, 0], color = 'C1')
ax[0].set(ylabel = '# Counts per second', title = 'Energy spectrum singles fit 511 keV (maximum, corrected)')
ax[1].set(xlabel = 'Energy [keV]', ylabel = r'Data - Fit [$\frac{\#}{s}$]')
ax[0].text(0.5, 0.05, f'µ: {round(gaussian.beta[0], np.sqrt(gaussian.cov_beta[0, 0]))}\nsigma: {round(gaussian.beta[1], np.sqrt(gaussian.cov_beta[1, 1]))}\nchiq/dof: {np.around(chi_sq/dof, 2)}\np: {sigfigRound(float(p), sigfigs = 3)}', horizontalalignment='center', verticalalignment='bottom', transform=ax[0].transAxes, bbox=dict(facecolor='white', edgecolor='gray'))
ax[0].grid()
ax[1].grid()
ax[0].legend()
ax[1].legend()
fig.tight_layout()
fig.savefig(f'./output/max_rate_single_corrected_residuals_511.pdf')
plt.clf()

# check for 1275 keV peak
my0 = 1200
sigma0 = 200
i_min = np.max(np.where(data_rate_max_single_denoised[:, 0] < 1150))
i_max = np.max(np.where(data_rate_max_single_denoised[:, 0] < 1400))
gaussian = gaussian_fit(data_rate_max_single_denoised[i_min:i_max:, 0], data_rate_max_single_denoised[i_min:i_max:, 1], np.ones_like(data_rate_max_single_denoised[i_min:i_max:, 0])*2/np.sqrt(12), sigma_data_rate_max_single_denoised[i_min:i_max:], my0, sigma0, np.max(data_rate_max_single_denoised[i_min:i_max:, 1])*np.sqrt(2*np.pi*sigma0**2))
print('////////// singles corrected - 1275 keV peak //////////')
print(f'{gaussian.beta=}')
print(f'errors: {np.sqrt(gaussian.cov_beta[0, 0])}, {np.sqrt(gaussian.cov_beta[1, 1])}, {np.sqrt(gaussian.cov_beta[2, 2])}')
chi_sq = np.sum(((data_rate_max_single_denoised[i_min:i_max:, 1] - gauss(gaussian.beta, data_rate_max_single_denoised[i_min:i_max:, 0]))/sigma_data_rate_max_single_denoised[i_min:i_max:])**2)
dof = i_max - i_min
print(f'chiq/dof: {chi_sq/dof}')
p = 1 - scs.chi2.cdf(chi_sq, dof)
print(f'{p=}')

fig, ax = plt.subplots(nrows = 2, sharex = True)
ax[0].plot(data_rate_max_single_denoised[i_min:i_max:, 0], gauss(gaussian.beta, data_rate_max_single_denoised[i_min:i_max:, 0]), label='fit', color='C1')
ax[0].errorbar(data_rate_max_single_denoised[i_min:i_max:, 0], data_rate_max_single_denoised[i_min:i_max:, 1], label='data', linestyle = '', marker = ',', xerr = 2/np.sqrt(12), yerr = sigma_data_rate_max_single_denoised[i_min:i_max:])
ax[1].errorbar(data_rate_max_single_denoised[i_min:i_max:, 0], data_rate_max_single_denoised[i_min:i_max:, 1] - gauss(gaussian.beta, data_rate_max_single_denoised[i_min:i_max:, 0]), yerr = sigma_data_rate_max_single_denoised[i_min:i_max:], label='residuals', linestyle = '', marker = '.')
ax[1].hlines(0, data_rate_max_single_denoised[i_min, 0], data_rate_max_single_denoised[i_max, 0], color = 'C1')
ax[0].set(ylabel = '# Counts per second', title = 'Energy spectrum singles fit 1275 keV (maximum, corrected)')
ax[1].set(xlabel = 'Energy [keV]', ylabel = r'Data - Fit [$\frac{\#}{s}$]')
ax[0].text(0.5, 0.05, f'µ: {round(gaussian.beta[0], np.sqrt(gaussian.cov_beta[0, 0]))}\nsigma: {round(gaussian.beta[1], np.sqrt(gaussian.cov_beta[1, 1]))}\nchiq/dof: {np.around(chi_sq/dof, 2)}\np: {sigfigRound(float(p), sigfigs = 3)}', horizontalalignment='center', verticalalignment='bottom', transform=ax[0].transAxes, bbox=dict(facecolor='white', edgecolor='gray'))
ax[0].grid()
ax[1].grid()
ax[0].legend()
ax[1].legend()
fig.tight_layout()
fig.savefig(f'./output/max_rate_single_corrected_residuals_1275.pdf')
plt.clf()
#'''

#'''
my0 = 480
sigma0 = 90
i_min = np.max(np.where(data_rate_max_coincident_denoised[:, 0] < 425))
i_max = np.max(np.where(data_rate_max_coincident_denoised[:, 0] < 520))
gaussian = gaussian_fit(data_rate_max_coincident_denoised[i_min:i_max:, 0], data_rate_max_coincident_denoised[i_min:i_max:, 1], np.ones_like(data_rate_max_coincident_denoised[i_min:i_max:, 0])*2/np.sqrt(12), sigma_data_rate_max_coincident_denoised[i_min:i_max:], my0, sigma0, np.max(data_rate_max_coincident_denoised[i_min:i_max:, 1])*np.sqrt(2*np.pi*sigma0**2))
print('////////// coincidents //////////')
print(f'{gaussian.beta=}')
print(f'errors: {np.sqrt(gaussian.cov_beta[0, 0])}, {np.sqrt(gaussian.cov_beta[1, 1])}, {np.sqrt(gaussian.cov_beta[2, 2])}')
chi_sq = np.sum(((data_rate_max_coincident_denoised[i_min:i_max:, 1] - gauss(gaussian.beta, data_rate_max_coincident_denoised[i_min:i_max:, 0]))/sigma_data_rate_max_coincident_denoised[i_min:i_max:])**2)
dof = i_max - i_min
print(f'chiq/dof: {chi_sq/dof}')
p = 1 - scs.chi2.cdf(chi_sq, dof)
print(f'{p=}')

scale = E_0 / gaussian.beta[0]
print(f'\n\tscaling factor coincidents: {scale}\n')
data_rate_max_coincident_denoised *= scale
sigma_data_rate_max_coincident_denoised *= scale

gaussian = gaussian_fit(data_rate_max_coincident_denoised[i_min:i_max:, 0], data_rate_max_coincident_denoised[i_min:i_max:, 1], np.ones_like(data_rate_max_coincident_denoised[i_min:i_max:, 0])*2/np.sqrt(12), sigma_data_rate_max_coincident_denoised[i_min:i_max:], my0, sigma0, np.max(data_rate_max_coincident_denoised[i_min:i_max:, 1])*np.sqrt(2*np.pi*sigma0**2))
print('////////// coincidents corrected //////////')
print(f'{gaussian.beta=}')
print(f'errors: {np.sqrt(gaussian.cov_beta[0, 0])}, {np.sqrt(gaussian.cov_beta[1, 1])}, {np.sqrt(gaussian.cov_beta[2, 2])}')
chi_sq = np.sum(((data_rate_max_coincident_denoised[i_min:i_max:, 1] - gauss(gaussian.beta, data_rate_max_coincident_denoised[i_min:i_max:, 0]))/sigma_data_rate_max_coincident_denoised[i_min:i_max:])**2)
dof = i_max - i_min
print(f'chiq/dof: {chi_sq/dof}')
p = 1 - scs.chi2.cdf(chi_sq, dof)
print(f'{p=}')

FWHM = 2 * np.sqrt(2*np.log(2)) * gaussian.beta[1]
sigma_FWHM = 2 * np.sqrt(2*np.log(2)) * np.sqrt(gaussian.cov_beta[1, 1])
print(f'\n\tFWHM coincidents: {round(FWHM, sigma_FWHM)}\n')

plt.errorbar(data_rate_max_coincident_denoised[:, 0], data_rate_max_coincident_denoised[:, 1], label='data', linestyle = '', marker = ',', xerr = 2/np.sqrt(12), yerr = sigma_data_rate_max_coincident_denoised)
plt.vlines(511 , np.min(data_rate_max_coincident_denoised[:, 1]), np.max(data_rate_max_coincident_denoised[:, 1]), color = 'C1', linestyles = 'dashed')
#plt.plot(data_rate_max_coincident_denoised[i_min:i_max:, 0], gauss(gaussian.beta, data_rate_max_coincident_denoised[i_min:i_max:, 0]), label='fit', color='C1')
if log_flag:
    plt.yscale('log')
plt.grid()
plt.legend()
plt.ylabel('# Counts per second')
plt.xlabel('Energy [keV]')
plt.title('Energy spectrum coincidents (maximum, corrected)')
plt.tight_layout()
plt.savefig(f'./output/max_rate_coincident_corrected_errorbar.pdf')
plt.cla()

fig, ax = plt.subplots(nrows = 2, sharex = True)
ax[0].plot(data_rate_max_coincident_denoised[i_min:i_max:, 0], gauss(gaussian.beta, data_rate_max_coincident_denoised[i_min:i_max:, 0]), label='fit', color='C1')
ax[0].errorbar(data_rate_max_coincident_denoised[i_min:i_max:, 0], data_rate_max_coincident_denoised[i_min:i_max:, 1], label='data', linestyle = '', marker = ',', xerr = 2/np.sqrt(12), yerr = sigma_data_rate_max_coincident_denoised[i_min:i_max:])
ax[1].errorbar(data_rate_max_coincident_denoised[i_min:i_max:, 0], data_rate_max_coincident_denoised[i_min:i_max:, 1] - gauss(gaussian.beta, data_rate_max_coincident_denoised[i_min:i_max:, 0]), yerr = sigma_data_rate_max_coincident_denoised[i_min:i_max:], label='residuals', linestyle = '', marker = '.')
ax[1].hlines(0, data_rate_max_coincident_denoised[i_min, 0], data_rate_max_coincident_denoised[i_max, 0], color = 'C1')
ax[0].set(ylabel = '# Counts per second', title = 'Energy spectrum coincidents fit 511 keV (maximum, corrected)')
ax[1].set(xlabel = 'Energy [keV]', ylabel = r'Data - Fit [$\frac{\#}{s}$]')
ax[0].text(0.5, 0.05, f'µ: {round(gaussian.beta[0], np.sqrt(gaussian.cov_beta[0, 0]))}\nsigma: {round(gaussian.beta[1], np.sqrt(gaussian.cov_beta[1, 1]))}\nchiq/dof: {np.around(chi_sq/dof, 2)}\np: {sigfigRound(float(p), sigfigs = 3)}', horizontalalignment='center', verticalalignment='bottom', transform=ax[0].transAxes, bbox=dict(facecolor='white', edgecolor='gray'))
ax[0].grid()
ax[1].grid()
ax[0].legend()
ax[1].legend()
fig.tight_layout()
fig.savefig(f'./output/max_rate_coincident_corrected_residuals_511.pdf')
plt.clf()
#'''

#'''
plt.bar(data_rate_empty_single[:, 0], data_rate_empty_single[:, 1], width = 1, label = 'data')
plt.vlines([597, 395, 290, 88] , np.min(data_rate_empty_single[:, 1]), np.max(data_rate_empty_single[:, 1]), color = 'C1', linestyles = 'dashed')
if log_flag:
    plt.yscale('log')
plt.grid()
plt.legend()
plt.ylabel('# Counts per second')
plt.xlabel('Energy [keV]')
plt.title('Energy spectrum singles (no sample)')
plt.tight_layout()
plt.savefig(f'./output/empty_rate_single_bar.pdf')
plt.cla()
plt.bar(data_rate_empty_coincident[:, 0], data_rate_empty_coincident[:, 1], width = 1, label = 'data')
plt.vlines([290, 88] , np.min(data_rate_empty_coincident[:, 1]), np.max(data_rate_empty_coincident[:, 1]), color = 'C1', linestyles = 'dashed')
if log_flag:
    plt.yscale('log')
plt.grid()
plt.legend()
plt.ylabel('# Counts per second')
plt.xlabel('Energy [keV]')
plt.title('Energy spectrum coincidents (no sample)')
plt.tight_layout()
plt.savefig(f'./output/empty_rate_coincident_bar.pdf')
plt.cla()
plt.bar(data_rate_max_single[:, 0], data_rate_max_single[:, 1], width = 1, label = 'data')
plt.vlines([511, 1275] , np.min(data_rate_max_single[:, 1]), np.max(data_rate_max_single[:, 1]), color = 'C1', linestyles = 'dashed')
if log_flag:
    plt.yscale('log')
plt.grid()
plt.legend()
plt.ylabel('# Counts per second')
plt.xlabel('Energy [keV]')
plt.title('Energy spectrum singles (maximum)')
plt.tight_layout()
plt.savefig(f'./output/max_rate_single_bar.pdf')
plt.cla()
plt.bar(data_rate_max_coincident[:, 0], data_rate_max_coincident[:, 1], width = 1, label = 'data')
plt.vlines([511] , np.min(data_rate_max_coincident[:, 1]), np.max(data_rate_max_coincident[:, 1]), color = 'C1', linestyles = 'dashed')
if log_flag:
    plt.yscale('log')
plt.grid()
plt.legend()
plt.ylabel('# Counts per second')
plt.xlabel('Energy [keV]')
plt.title('Energy spectrum coincidents (maximum)')
plt.tight_layout()
plt.savefig(f'./output/max_rate_coincident_bar.pdf')
plt.cla()
#'''

#'''
plt.bar(data_rate_max_single_denoised[:, 0], data_rate_max_single_denoised[:, 1], width = 1, label = 'data')
if log_flag:
    plt.yscale('log')
plt.grid()
plt.legend()
plt.ylabel('# Counts per second')
plt.xlabel('Energy [keV]')
plt.title('Energy spectrum singles (maximum, corrected)')
plt.tight_layout()
plt.savefig(f'./output/max_rate_single_corrected_bar.pdf')
plt.cla()
plt.bar(data_rate_max_coincident_denoised[:, 0], data_rate_max_coincident_denoised[:, 1], width = 1, label = 'data')
if log_flag:
    plt.yscale('log')
plt.grid()
plt.legend()
plt.ylabel('# Counts per second')
plt.xlabel('Energy [keV]')
plt.title('Energy spectrum coincidents (maximum, corrected)')
plt.tight_layout()
plt.savefig(f'./output/max_rate_coincident_corrected_bar.pdf')
plt.cla()
#'''