import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from PeakOTron import PeakOTron
import pandas as pd
import numpy as np
import argparse


def float_or_none(value):
    return None if value.lower() == 'none' else float(value)


parser = argparse.ArgumentParser(description='Fit SiPM data')
parser.add_argument('-V_bd_hmt', type=float,
                    default=26.1, help='V_bd_hmt value')
parser.add_argument('-V_0_hmt', type=float, default=1.4, help='V_0_hmt value')
parser.add_argument('-tau', type=float, default=20.0,
                    help='SLOW COMPONENT OF SIPM PULSE')
parser.add_argument('-t_0', type=float, default=100.0,
                    help='PRE-INTEGRATION TIME')
parser.add_argument('-t_gate', type=float, default=100.0, help='GATE LENGTH')
parser.add_argument('-bin_0', type=float_or_none, default=-100.0,
                    help='SELECT FIRST BIN OF SPECTRUM')
parser.add_argument('-truncate_nsigma0_up', type=float_or_none,
                    default=2.0, help='SCAN SPECTRUM FROM Q < Q_0 - 4 sigma_0')
parser.add_argument('-truncate_nsigma0_do', type=float_or_none, default=2.0,
                    help='EVALUATE SPECTRUM CHI2 IN Q_0 - x*sigma_0 < Q < Q_0 + 2*sigma_0')
parser.add_argument('-prefit_only', action='store_true',
                    help='FIT THE WHOLE SPECTRUM')
parser.add_argument('-folder', type=str, default='data/sumlab',
                    help='Directory containing the data files')

args = parser.parse_args()


def C_tau(V, V_bd, V_0): return (V - V_bd)/V_0


def f_tau(V, V_bd, V_0): return -1/np.log((1-np.exp(C_tau(V, V_bd, V_0)
                                                    * np.exp(-1)))/(1 - np.exp(C_tau(V, V_bd, V_0))))


V_bd_hmt = args.V_bd_hmt
V_0_hmt = args.V_0_hmt
tau = args.tau  # SLOW COMPONENT OF SIPM PULSE
t_0 = args.t_0  # PRE-INTEGRATION TIME
t_gate = args.t_gate  # GATE LENGTH
bin_0 = args.bin_0  # SELECT FIRST BIN OF SPECTRUM (CAN BE AUTOMATIC)
# SCAN SPECTRUM FROM Q < Q_0 - 4 sigma_0
truncate_nsigma0_up = args.truncate_nsigma0_up
# EVALUATE SPECTRUM CHI2 IN Q_0 - x*sigma_0  < Q < Q_0 + 2*sigma_0
truncate_nsigma0_do = args.truncate_nsigma0_do
prefit_only = args.prefit_only  # FIT THE WHOLE SPECTRUM

out_dict = {}
files_to_fit = []

# Find all histograms in directory
folder = args.folder
for root, dirs, files in os.walk(folder):
    for file in files:
        if file.endswith(".txt"):
            files_to_fit.append([file, os.path.join(root, file)])

print("\033[95m\n=======================================")
print("   PeakOTron")
print("=======================================\033[0m")

# Loop thorough files
for i, (file, path) in enumerate(files_to_fit):
    print("\033[95mFitting: {:s}\033[0m".format(file))
    V = float(file.split('deg')[1].split('V')[0].replace('_', '.'))
    f_tau_hmt = f_tau(V, V_bd_hmt, V_0_hmt)

    # Load files.
    data = np.loadtxt(path, skiprows=0)
    # Create a PeakOTron Fit Object.
    f_data = PeakOTron(verbose=False)
    # Perform fit.
    f_data.Fit(data,
               tau=tau,  # SLOW PULSE COMPONENT TIME CONSTANT (ns)
               t_gate=t_gate,  # GATE LENGTH (ns)
               t_0=t_0,  # INTEGRATION TIME BEFORE GATE (ns)
               tau_R=f_tau_hmt*tau,
               bin_0=bin_0,
               truncate_nsigma0_up=truncate_nsigma0_up,
               truncate_nsigma0_do=truncate_nsigma0_do
               )

    f_data.PlotFit(plot_in_bins=True, display=False,
                   save_directory=f"{folder}/{file[:-4]}_fit.png")

    fit_out = {}
    prefit_val, prefit_err = f_data.GetPrefitResults(bin_units=False)
    print("\033[95m"+rf"Prefit: G = {prefit_val.get('G')} d_G = {prefit_err.get('G')}"+"\033[0m")
    for key, value in prefit_val.items():
        fit_out["prefit_{:s}".format(key)] = value
    for key, value in prefit_err.items():
        fit_out["prefit_d_{:s}".format(key)] = value

    #if not prefit_only:
    fit_val, fit_err = f_data.GetFitResults(bin_units=False)
    print("\033[95m"+rf"Fit: G = {fit_val.get('G')} d_G = {fit_err.get('G')}"+"\033[0m")
    for key, value in fit_val.items():
        fit_out["{:s}".format(key)] = value
    for key, value in fit_err.items():
        fit_out["d_{:s}".format(key)] = value

    df = pd.DataFrame.from_dict([fit_out])
    df.to_csv("{}/fit_results_{:s}.csv".format(folder, file[:-4]))

print("\033[95m=======================================\033[0m")