import numpy as np
import matplotlib.pyplot as plt
import scipy
from tqdm.auto import tqdm, trange
import os

cs_energy_peak = 661.7 #keV
co_energy_peak1 = 1173 # keV
co_energy_peak2 = 1332 # keV

import boost_histogram as bh
import hist
from hist import Hist
import pandas as pd
from scipy.optimize import curve_fit

def peak_function(x, ampliude, center,sigma, slope, intercept):
    return ampliude*np.exp(-(x-center)**2/(2*sigma**2)) + slope*x + intercept


def fit_peak(x,y,guess,plot = False):
    params, cov = curve_fit(peak_function,x,y,p0=guess,  maxfev = 12000)
    if plot:
        ax = plt.gca()
        ax.scatter(x,y)
        ax.plot(x,peak_function(x,*params))
        _,_,_,sl,it = params
        ax.plot(x, sl*x+it, ls='--')
    return params

def fit_two_peaks(x,y,guess,plot = False):
    def two_peaks_func(x, amp, cen, sig, amp2, cen2, sig2, sl, inte):
        return amp*np.exp(-(x-cen)**2/(2*sig**2)) + amp2*np.exp(-(x-cen2)**2/(2*sig2**2))  + sl*x + inte
    params, cov = curve_fit(two_peaks_func,x,y,p0=guess,
                            #bounds=((0,x.min(),0,0,x.min(),0,None,None),(y.max(),x.max(),x.max(),y.max(),x.max(),x.max(),0,None)),
                            maxfev = 12000)
    if plot:
        ax = plt.gca()
        ax.scatter(x,y)
        ax.plot(x,two_peaks_func(x,*params))
        _,_,_,_,_,_,sl,it = params
        ax.plot(x, sl*x+it, ls='--')
        
        ax.plot(x,two_peaks_func(x,*guess),ls='--')
    return params

def load_csv_to_histogram(csv_path, value_col='V', weight_col='hit', plot = True, ax = None):
    df = pd.read_csv(csv_path, names=[value_col, weight_col], skiprows=1)
    
    centers = df[value_col].to_numpy()
    edges = np.concatenate([
        [2*centers[0] - centers[1]],  # First edge
        (centers[:-1] + centers[1:]) / 2,  # Midpoints between centers
        [2*centers[-1] - centers[-2]]  # Last edge
    ])

    h = Hist.new.Variable(edges=edges).Double()
    h.fill(df[value_col].to_numpy(), weight=df[weight_col].to_numpy())
    
    if plot and ax is None:
        fig, ax = plt.subplots()
    if plot:
        h.plot(ax=ax)
    
    return h

def plot_histogram_with_slice(hist, slice_range, rebin=4, plot_original=True, ax=None):
    """
    Plots the original histogram and a sliced+rebinned version.
    Returns the sliced histogram.
    """
    if ax is None:
        fig, ax = plt.subplots()
    if plot_original:
        hist.plot(ax=ax, label='Original', alpha=0.5)
    # Slicing and rebinning
    start, stop = slice_range
    h_slice = hist[start*1j:stop*1j:rebin*1j]
    h_slice.plot(ax=ax, label=f'Slice {start}:{stop}', color='C1')
    ax.legend()
    return h_slice

def fit_histogram_peaks(hist, n_peaks=1, guess=None, plot=True):
    """
    Fits one or two peaks to the histogram.
    Returns peak centers and FWHMs.
    """
    x = hist.axes[0].centers
    y = hist.values()
    if n_peaks == 1:
        params = fit_peak(x, y, guess, plot=plot)
        _, center, sigma, _, _ = params
        FWHM = 2.355 * sigma
        return [center], [FWHM]
    elif n_peaks == 2:
        params = fit_two_peaks(x, y, guess, plot=plot)
        _, center1, sigma1, _, center2, sigma2, _, _ = params
        FWHM1 = 2.355 * sigma1
        FWHM2 = 2.355 * sigma2
        return [center1, center2], [FWHM1, FWHM2]
    else:
        raise ValueError("Only n_peaks=1 or 2 supported.")

L3_Cs_hist = load_csv_to_histogram("LYSO_3_Cs.csv",'V', 'hit', False)

L3_Cs_hist_4fit = plot_histogram_with_slice(L3_Cs_hist,(0.125,0.175),rebin=4)

LYSO3_Cs,LYSO3_Cs_FWHM = fit_histogram_peaks(L3_Cs_hist_4fit, 1,
                                                         (670,0.15,0.002,0,0), True)

L3_Co_hist = load_csv_to_histogram("LYSO_3_Co.csv",'V', 'hit', False)
L3_Co_hist_4fit = plot_histogram_with_slice(L3_Co_hist,(0.21,0.32),rebin=4)

(LYSO3_Co1, LYSO3_Co2),(LYSO3_Co1_FWHM, LYSO3_Co2_FWHM) = fit_histogram_peaks(L3_Co_hist_4fit, 2,
                                                         (50,0.245,0.002,30,0.29,0.002,-80/0.3,80), True)

E_peaks = [cs_energy_peak, co_energy_peak1, co_energy_peak2]

L3 = [LYSO3_Cs[0], LYSO3_Co1, LYSO3_Co2]
L3FWHM = [LYSO3_Cs_FWHM[0], LYSO3_Co1_FWHM, LYSO3_Co2_FWHM]

def linear_function(V, a, b):
    # E as function of V
    return a * V + b

fig, ax = plt.subplots(figsize=(8, 5))

datasets = [
    (E_peaks, L3, L3FWHM, 'LYSO-3'),
]

calib_params = dict()

# Plot data and fit: now x = V (voltage), y = E (keV)
for E, V, V_err, label in datasets:
    V_arr = np.array(V)
    E_arr = np.array(E)
    V_err_arr = np.abs(np.array(V_err))
    
    # Linear fit E = a*V + b
    params, _ = curve_fit(linear_function, V_arr, E_arr)
    a, b = params
    
    calib_params[label] = params
    # plot fit with same color
    V_fit = np.linspace(0, V_arr.max()*1.05, 200)
    ln = ax.plot(V_fit, linear_function(V_fit, a, b), '--', alpha=0.8)
    color = ln[0].get_color()
    
    # Convert FWHM (in V) to keV using E = a*V + b -> dE = a * dV
    E_errs_keV = np.abs(a) * V_err_arr
    
        # plot data with x-errors (voltage)
    eb = ax.errorbar(V_arr, E_arr, xerr=V_err_arr, yerr=E_errs_keV,  marker='o', linestyle='None', label=f"{label}: E(keV) = {a:.3f} * V + {b:.3f}", capsize=5, color = color)
    #color = eb[0].get_color()  # use same color for fit
    
    # Print fit and converted FWHMs
    print(f"{label}: E = {a:.3f} * V + {b:.3f}")

ax.set_xlabel('Voltage (V)')
ax.set_ylabel('Energy (keV)')
ax.set_xlim(left=0)
ax.set_ylim(bottom=0)
ax.legend()
plt.tight_layout()
plt.savefig("L3.png")

LYSO-3: E = 4592.465 * V + -5.567

LYSO3¶