from scipy.optimize import curve_fit
import numpy as np
import matplotlib.pyplot as plt

# ---------- DEFINE ---------------
datapath = '/Volumes/TOSHIBA EXT/Test-Pulses4(22deg)/'
datapath2 = '/Users/marialauraperez/Desktop/SummerCERN/ChunkedSources/'
data_path_output = '/Users/marialauraperez/Desktop/'
txt_output = '/Users/marialauraperez/Desktop/'
first_charges = np.arange(500, 1600, 100)
second_charges = np.arange(2000, 12000, 1000)
charges = np.concatenate((first_charges, second_charges))
ToT_TP = list()
ToT_sources = list()
energies_TP = charges*3.62/1000
del first_charges, second_charges
energies_sources = [4.512, 5.9, 6.405, 8.046, 15.775, 17.48, 25.271]#, 59.5]

# --------- IMPORT --------------
for q in range(len(charges)):
    matrix = np.genfromtxt(datapath+str(charges[q])+'e.txt_iToT.txt')
    ToT_TP.append(matrix)
for source in ['Ti', 'Mn', 'Fe', 'Cu', 'Zr', 'Mo', 'Sn']:#, 'Am241']:
    matrix = np.genfromtxt(datapath2+source+'_FinalPeakMatrix.txt').reshape([256, 256])
    ToT_sources.append(matrix)

# -------- FUNCTIONS -----------


def surrogate_function_1(x, a, b, c, t, g):
    y = a*g*x + b - c/(g*x-t)
    y[np.where(g*x < t)] = 0
    y *= y > 0
    return y


def surrogate_function_2(A, B, C, T):
    def surrogate_function_g(x, g):
        y = g*A*x + B - C/(g*x-T)
        y[np.where(g * x < T)] = 0
        y *= y > 0
        return y
    return surrogate_function_g


def normal_surrogate_function(x, a, b, c, t):
    y = a*x + b - c/(x-t)
    y[np.where(x < t)] = 0
    y *= y > 0
    return y


def get_calibration_matrices(energies_TP, ToT_mats_TP, energies_sources, ToT_mats_sources):
    energies = [energies_TP, energies_sources]
    ToTs = [ToT_mats_TP, ToT_mats_sources]
    a_mat = np.zeros(np.shape(ToT_mats_TP[0]))
    b_mat = np.zeros(np.shape(ToT_mats_TP[0]))
    c_mat = np.zeros(np.shape(ToT_mats_TP[0]))
    t_mat = np.zeros(np.shape(ToT_mats_TP[0]))
    g_mats = [np.zeros(np.shape(ToT_mats_TP[0])), np.zeros(np.shape(ToT_mats_TP[0]))]
    px125 = list()
    for w in range(2):
        ToT_matrices = ToTs[w]
        energy_vals = energies[w]
        g = g_mats[w]
        for i in range(256):
            for j in range(256):
                entry = list()
                for k in range(len(ToT_matrices)):
                    if w is 0:
                        entry.append(ToT_matrices[k][i, j]/100)
                    else:
                        entry.append(ToT_matrices[k][i, j])
                if i is 125 and j is 125:
                    px125.append(entry)
                try:
                    if w is 0:
                        popt, pcov = curve_fit(surrogate_function_1, energy_vals, entry)
                        a_mat[i, j] = popt[0]
                        b_mat[i, j] = popt[1]
                        c_mat[i, j] = popt[2]
                        t_mat[i, j] = popt[3]
                        g[i, j] = popt[4]
                    else:
                        try:
                            popt, pcov = curve_fit(surrogate_function_2(a_mat[i, j], b_mat[i, j], c_mat[i, j],
                                                                        t_mat[i, j]), energy_vals, entry)
                            g[i, j] = popt[0]
                        except ValueError:
                            print('Error: calibration not performed in pixel (', str(i)+', '+str(j)+' )')
                except RuntimeError:
                    print('Error: calibration not performed in pixel (', str(i)+', '+str(j)+' )')
                    # plt.figure()
                    # plt.scatter(energy_vals, entry)
                    # plt.show()
    return a_mat, b_mat, c_mat, t_mat, g_mats, px125

# ------------- EXCECUTION ------------------


a, b, c, t, [g1, g2], [px1, px2] = get_calibration_matrices(energies_TP, ToT_TP, energies_sources,
                                                                              ToT_sources)

new_a = g2*a
new_b = b
new_c = c/g2
new_t = t/g2

x1 = np.linspace(0, 70, 100)
plt.figure()
plt.scatter(energies_TP, px1, label='Test pulse points')
plt.plot(x1, surrogate_function_1(x1, a[125, 125], b[125, 125], c[125, 125], t[125, 125], g1[125, 125]),
         label='Test pulse fit: g = '+str(np.round(g1[125, 125], 3)))
plt.scatter(energies_sources, px2, label='Source points')
plt.plot(x1, surrogate_function_1(x1, a[125, 125], b[125, 125], c[125, 125], t[125, 125], g2[125, 125]),
         label='Source fit: g = '+str(np.round(g2[125, 125], 3)))
plt.legend()
plt.title('Energy calibration curve for pixel (125, 125)', fontweight='bold')
plt.xlabel('Energy [keV]')
plt.ylabel('ToT value')
plt.savefig(data_path_output+'Pixel125_CalibGain_NoAm.png')
# plt.show()

np.savetxt(txt_output+'A_STP_NoAm.txt', new_a, fmt='%.5f')
np.savetxt(txt_output+'B_STP_NoAm.txt', new_b, fmt='%.5f')
np.savetxt(txt_output+'C_STP_NoAm.txt', new_c, fmt='%.5f')
np.savetxt(txt_output+'T_STP_NoAm.txt', new_t, fmt='%.5f')

np.savetxt(txt_output+'G_TestPulse_NoAm.txt', g1, fmt='%.5f')
np.savetxt(txt_output+'G_Source_NoAm.txt', g2, fmt='%.5f')