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

tau_list = []
def exponential_decay(t, y0, A, tau):
    return y0 + A * np.exp(-t / tau)

def fit_exponential(df, t_column, y_column, y_errors=None):
    t = df[t_column]
    y = df[y_column]

    popt, pcov = curve_fit(exponential_decay, t, y)

    t_fit = np.linspace(t.min(), t.max(), 1000)

    y_fit = exponential_decay(t_fit, *popt)


    # If errors are provided, use them as weights
    if y_errors is not None:
        weights = 1.0 / y_errors
    else:
        weights = None


    popt, pcov = curve_fit(exponential_decay, t, y, sigma=y_errors, absolute_sigma=True)

    y_fit = exponential_decay(t_fit, *popt)

    tau = popt[2]  # Extracting tau value

    plt.errorbar(t, y, yerr=y_errors, fmt='b.', label='Data')
    plt.plot(t_fit, y_fit, 'r--', label='Fit: y0={:.2f}, A={:.2f}, tau={:.2f}'.format(*popt))
    plt.xlabel('t')
    plt.ylabel('y')
    plt.title('Exponential Fit for {}'.format(df.name))
    plt.legend()
    plt.grid(True)
    plt.savefig(format(df.name)+'.png')
    plt.show()

#read in df
dataframe = pd.read_csv('datatau.csv', header = 0)
dataframe.dropna()
#names
well_names = ['e1c1c2','c1c2c3','c2c3c4','c3c4c5','c4c5c6','c5c6c7','c6c7c8','c7c8c9']
#well_names = ['c1c2c3']

for i in well_names:
    df = dataframe[dataframe['well name'] == i] 
    df.name = i
    df.sort_values('wait time [s]', inplace=True)

    # Calculate standard deviation for each row
    df['y_errors'] = df[['signal 1 [nVs]', 'signal 2 [nVs]', 'signal 3 [nVs]']].std(axis=1)

    # Fit exponential curve
    fit_exponential(df, 'wait time [s]', 'avrg signal [nVs]', y_errors=df['y_errors'])
print(tau_list)