import numpy as np
import matplotlib.pyplot as plt
import Fit

data_path_disk = '/Users/marialauraperez/Desktop/SummerCERN/npz histos 22 deg/'
output_data_path = '/Users/marialauraperez/Desktop/SummerCERN/ChunkedSources/'
sources = ['Ti', 'Mn', 'Fe', 'Cu', 'Zr', 'Mo', 'Sn']#, 'Am241']
colors = ['red', 'orange', 'yellowgreen', 'green', 'dodgerblue', 'navy', 'purple', 'pink']

# for i in range(4, 7):
#     print('processing data from ', sources[i])
#     matrix_path = output_data_path + ('%s_FinalPeakMatrix.txt' % sources[i])
#     data = np.load(data_path_disk+sources[i]+'XRF_cluster_sp_hist2.npz')
#     data_hist = data['hist']  # 65536 x 200 array. make a histogram for every index
#     data_tot = data['ToT']
#     p0 = None
#     f = Fit.FitGauss()
#     peaks = np.zeros(len(data_hist))
#     for j in range(len(data_hist)):
#         try:
#             gauss = f.fit(x=data_tot, y=data_hist[j], p0=p0)
#             p0 = gauss[0]
#             mu = gauss[0][1]
#             if gauss[2].R2 < 0.8:
#                 p0 = None
#         except:
#             mu = 0.0
#         peaks[j] = mu
#         if j%1000 is 0:
#             print(j)
#     np.savetxt(matrix_path, peaks, fmt='%.4f', delimiter='\n')

# ----------- FOR AMERICIUM, GADOLINIUM AND TIN ---------------------------------
# print('processing data from ', sources[7])
# matrix_path = output_data_path + ('%s_FinalPeakMatrix.txt' % sources[7])
# data = np.load(data_path_disk+sources[7]+'XRF_cluster_sp_hist2.npz')
# data_hist = data['hist']  # 65536 x 200 array. make a histogram for every index
# data_tot = data['ToT']
# p0 = None
# p2 = None
# f = Fit.FitGauss()
# peaks_am = np.zeros(len(data_hist))
# peaks_sn = np.zeros(len(data_hist))
# for j in range(len(data_hist)):
#     try:
#         argmax_am = np.argmax(data_hist[j][100:200])+100
#         argmax_sn = np.argmax(data_hist[j][40:90])+40
#
#         gauss_am = f.fit(x=data_tot[argmax_am-20:argmax_am+20], y=data_hist[j][argmax_am-20:argmax_am+20], p0=p0)
#         gauss_sn = f.fit(x=data_tot[argmax_sn-20:argmax_sn+20], y=data_hist[j][argmax_sn-20:argmax_sn+20], p0=p2)
#         mu_am = gauss_am[0][1]
#         mu_sn = gauss_sn[0][1]
#         p0 = gauss_am[0]
#         p2 = gauss_sn[0]
#         if gauss_am[2].R2 < 0.8:
#             p0 = None
#         if gauss_sn[2].R2 < 0.8:
#             p2 = None
#     except:
#         mu_am = 0.0
#         mu_sn = 0.0
#     peaks_am[j] = mu_am
#     peaks_sn[j] = mu_sn
#     if j % 1000 is 0:
#         print(j)
# np.savetxt(output_data_path+('%s_FinalPeakMatrix.txt' % sources[7]), peaks_am, fmt='%.4f', delimiter='\n')
# np.savetxt(output_data_path+'(Check)Sn_FinalPeakMatrix.txt', peaks_sn, fmt='%.4f', delimiter='\n')


def normalize(v):
    norm = np.linalg.norm(v)
    if norm == 0:
        return v
    return v / norm


mus = np.zeros(len(sources))
plt.figure(figsize=(9,5))
plt.xlabel('ToT value')
plt.ylabel('Normalized counts')
plt.title('Total spectra with Timepix3 (T=22ºC)', fontweight='bold')
#plt.ticklabel_format(style='sci', axis='y', scilimits=(0, 5))
plt.xlim(0, 90)
for i in range(len(sources)):
    data = np.load(data_path_disk + sources[i] + 'XRF_cluster_sp_hist2.npz')
    data_hist = data['hist']  # 65536 x 200 array. make a histogram for every index
    data_tot = data['ToT']
    mean = np.zeros(np.shape(data_tot))
    for j in range(len(data_tot)):
        mean[j] = np.sum(np.transpose(data_hist)[:][j])
    mean = normalize(mean)
    f = Fit.FitGauss()
    if sources[i] != 'Am241':
        gauss = f.fit(x=data_tot, y=mean)
        y1 = Fit.gauss(data_tot, gauss[0])
        p0 = gauss[0]
        mu = gauss[0][1]
        mus[i] = mu
        plt.scatter(data_tot, mean, label='_nolegend_', s=10, color=colors[i])
        plt.plot(data_tot, y1, label=sources[i]+r': $\mu=$ ' + str(np.round(mu, 3)), linestyle='--', color=colors[i])
    else:
        gauss_am = f.fit(x=data_tot[100:150], y=mean[100:150])
        gauss_gd = f.fit(x=data_tot[75:100], y=mean[75:100])
        gauss_sn = f.fit(x=data_tot[40:75], y=mean[40:75])
        y1 = Fit.gauss(data_tot, gauss_am[0])
        y2 = Fit.gauss(data_tot, gauss_gd[0])
        y3 = Fit.gauss(data_tot, gauss_sn[0])
        mu_am = gauss_am[0][1]
        mu_gd = gauss_gd[0][1]
        mu_sn = gauss_sn[0][1]
        mus[i] = mu_am
        plt.scatter(data_tot, mean, label='_nolegend_', s=10, color=colors[i])
        plt.plot(data_tot[80:150], y1[80:150], label=r'Am: $\mu=$ ' + str(np.round(mu_am, 3)), linestyle='--', color='hotpink')
        #plt.plot(data_tot, y2, label='_nolegend_', linestyle='--')
        #plt.plot(data_tot, y3, label='_nolegend_', linestyle='--')
plt.legend()
plt.savefig('/Users/marialauraperez/Desktop/totalfitssources_noam.png')

# -------- FOR CALCIUM AND COPPER --------------

# data = np.load(data_path_disk + 'CaXRF_cluster_sp_hist2.npz')
# data_hist = data['hist']  # 65536 x 200 array. make a histogram for every index
# data_tot = data['ToT']
# mean = np.zeros(np.shape(data_tot))
# for j in range(len(data_tot)):
#     mean[j] = np.sum(np.transpose(data_hist)[:][j])
# f = Fit.FitGauss()
# gauss_ca = f.fit(x=data_tot[0:13], y=mean[0:13])
# gauss_cu = f.fit(x=data_tot[13:25], y=mean[13:25])
# y1 = Fit.gauss(data_tot, gauss_ca[0])
# y2 = Fit.gauss(data_tot, gauss_cu[0])
# mu_ca = gauss_ca[0][1]
# mu_cu = gauss_cu[0][1]
# plt.figure()
# plt.plot(data_tot, mean, label='Data')
# plt.plot(data_tot, y1, label=r'Fit Ca: $\mu=$ ' + str(np.round(mu_ca, 3)), linestyle='--')
# plt.plot(data_tot, y2, label=r'Fit Cu: $\mu=$ ' + str(np.round(mu_cu, 3)), linestyle='--')
# plt.xlabel('ToT value')
# plt.ylabel('Total counts')
# plt.title('Ca total spectrum with Timepix3 (T=22ºC)', fontweight='bold')
# plt.ticklabel_format(style='sci', axis='y', scilimits=(0, 5))
# plt.legend()
# plt.savefig('/Users/marialauraperez/Desktop/' + 'Ca_totalfit.png')