import numpy as np
import datetime
from scipy.optimize import curve_fit

# start = [11.75-0.25*i for i in range(8)]
# print(start)

# temps = []
# for s in start:
#     temps.append([s-i*0.0194 for i in range(150)])

# data = np.array(temps)
# times = [i for i in range(150)]


# Define linear fit function:
def linear_fit(x, a, b):
    '''
    Linear function for fitting of CRAFT data
    '''
    return a+b*x

# Define smooth_transition function for correcting not-linear section in the beginning of the cooldown
def smooth_transition(x, a, b, c, x0):
    """
    A function that is constant before t0 and transitions smoothly into a linear function after t0.
    Is needed for fitting of CRAFT data. In case the temperature does not increase linear at the 
    beginning of the cooldown.
    Parameters:
    t  : array-like, time values
    t0 : float, transition point
    a  : float, constant value before t0
    b  : float, slope of the linear function after t0
    c  : float, controls the smoothness of the transition
    
    Returns:
    array-like, function values
    """
    return a + (b * (x - x0)) / (1 + np.exp(-c * np.clip(x - x0, -100, 100)))
def mean_sigma(input): #calculate average an standart deviation
    '''
    Calculates the mean and standard deviation of a given list of numbers.
    Parameters
    ----------
    input : list of numbers
        List of numbers to be analyzed.
    Returns
    -------
    mean : float
        Mean of the input list.
    sigma : float
        Standard deviation of the input list.
    '''
    n = len(input)
    mean = np.nanmean([input])
    deviation = [(xi - mean)**2 for xi in input]
    sigma = np.sqrt(1/(n-1)*np.nansum(deviation))
    return(mean, sigma)

def find_index_of_first_equal_to(lst, value):
    '''Finds the index of the first element in a list that is equal to a given value. Returns -1 if no such element is found.'''
    return next((i for i, x in enumerate(lst) if x == value), np.nan)

def calc_gradient(input, times, sensors:list = [1,3,4,5,6,7,8], Tc = 9.2, return_indices: bool = False): 
    '''Analyses cooldown. "sensors" specifies all sensors that are being used (1 to 8)
    input is np.array of temperatures
    times is np.array of times
    sensors is list with sensors used for analysis
    Tc is critical temperature
    return_indices specifies if the indices of the sensors closest to the critical temperature should be returned'''
    #find indices of numbers closest to Tc
    positions = [9.5 - ((x-1)*9.5/7) for x in sensors] #get the sensor position on the sample
    indices = []#list of indices
    #convert input to np.array if necessary
    if type(input) == list:
        input = np.array(input)
    #convert datetime to timestamp, if necessary
    if type(times[0]) == datetime.datetime:
        times = [t.timestamp() for t in times]
    #transpose array if fist dimension if larger than second
    if input.shape[0] > input.shape[1]:
        input = input.T 

    #find indices at which critical temperature is reached
    input = input[[x-1 for x in sensors]]
    for i in input:
        index = (np.abs(i - Tc)).argmin()
        indices.append(index)


    #calculate local gradients and cooldown rates
    grad_loc = [np.NaN for x in range(8)] #local temperature gradients
    rate_loc = [np.NaN for x in range(8)] #local cooldown rates

    for i in range(len(sensors)):
        if i == 0: #for first sensor take temperature difference to the sensor below
            DT = input[0,indices[0]] - input[1,indices[0]] #temperature difference
            t = times[indices[0]] - times[indices[1]] #time it took the pahsefront to move from second sensor to first sensor
            Ds = positions[i] - positions[i+1] #distance between sensors
            grad_loc[sensors[i]-1] = (DT/Ds)
            try:
                rate_loc[sensors[i]-1] = (DT/t)
            except ValueError:
                rate_loc[sensors[i]-1] = 0
        elif i == len(sensors)-1: #for last sensor take temperature difference to the sensor above
            DT = input[i-1,indices[i]] - input[i,indices[i]] #temperature difference
            t = times[indices[i-1]] - times[indices[i]] #time it took the pahsefront to move from sensor 7 to sensor 8
            Ds = positions[i-1] - positions[i] #distance between sensors
            grad_loc[sensors[i]-1] = (DT/Ds)
            try:
                rate_loc[sensors[i]-1] = (DT/t)
            except ValueError:
                rate_loc[sensors[i]-1] = 0
        else: #for middle sensors take average temperature difference to sensor below and above
            Ds_top = positions[i-1] - positions[i] #distance to sensor above
            Ds_bot = positions[i] - positions[i+1] #distance to sensor above
            DT = np.mean([input[i-1,indices[i]] - input[i,indices[i]] , input[i,indices[i]] - input[i+1,indices[i]]]) #mean temperature difference to sensor above and below
            grad = np.mean([(input[i-1,indices[i]] - input[i,indices[i]]) / Ds_top , (input[i,indices[i]] - input[i+1,indices[i]]) / Ds_bot]) #mean temperature gradient to sensor above and below
            t = np.mean([times[indices[i-1]] - times[indices[i]], times[indices[i]] - times[indices[i+1]]]) #mean time to sensor above and from sensor below
            grad_loc[sensors[i]-1] = grad
            try:
                rate_loc[sensors[i]-1] = (DT/t)
            except ValueError:
                rate_loc[sensors[i]-1] = 0

    #calc global gradient and transition time
    grad_glob = (input[0,indices[-1]] - input[-1,indices[-1]]) / (positions[0] - positions[-1]) #gives gradient in K/cm. Its 9.5 cm instead of 10 cm since the cernox sensors are not at the very end of the sample. This is momentary global gradient when the lowest sensor passes Tc
    trans_time = times[indices[0]] - times[indices[-1]] 
    (av_grad_glob, e_av_grad_glob) = mean_sigma(grad_loc) #average gradient from local gradients. The error is also returned.
    (av_rate_glob, e_av_rate_glob) = mean_sigma(rate_loc) #average cooldown rate from local rates.

    if return_indices:
        return([grad_glob, av_grad_glob, e_av_grad_glob, trans_time, av_rate_glob, e_av_rate_glob, grad_loc, rate_loc, indices])
    else:
        return([grad_glob, av_grad_glob, e_av_grad_glob, trans_time, av_rate_glob, e_av_rate_glob, grad_loc, rate_loc])

def calc_gradient_CRAFT(input, times, sensors:list = [1,3,4,5,6,7,8], Tc_top = 9.2, Tc_bottom = 9.2, return_indices: bool = False): 
    '''Analyses cooldown. "sensors" specifies all sensors that are being used (1 to 8)
    input is np.array of temperatures
    times is np.array of times
    sensors is list with sensors used for analysis
    Tc is critical temperature
    return_indices specifies if the indices of the sensors closest to the critical temperature should be returned'''
    #find indices of numbers closest to Tc
    sensors = [1,8] #override sensor list that is passed to function for now because only 1 and 8 are used
    positions = [9.5 - ((x-1)*9.5/7) for x in sensors] #get the sensor position on the sample
    indices = []#list of indices
    #convert input to np.array if necessary
    if type(input) == list:
        input = np.array(input)
    #convert datetime to timestamp, if necessary and convert it to np.array
    if type(times[0]) == datetime.datetime:
        times = [t.timestamp() for t in times]
    times = np.array(times)
    #transpose array if fist dimension if larger than second
    if input.shape[0] > input.shape[1]:
        input = input.T 

    #find indices at which critical temperature is reached
    input = input[[x-1 for x in sensors]]
    #Calculate t_Tc_start: Lowest timepoint where Tc is reached by Cernox 1 OR 8. "OR" is important, because CRAFT allows for negative gradients during cooldown
    t_Tc_start = times[np.min(np.where((input[0] <= Tc_top) | (input[1] <= Tc_bottom)))]
    #Calculate t_Tc_stop: Highest timepoint where Tc is reached by Cernox 1 OR 8
    t_Tc_stop = times[np.max(np.where((input[0] >= Tc_top) | (input[1] >= Tc_bottom)))]

    # Smooth_transition fit to the temperature data of Cernox 1 and 8, from t_Tc_start - t_Tc_start_delta until t_Tc_stop + t_Tc_stop_delta
    # Often a linear fit is sufficient. But when the sc transitions starts immediately after cooldown start, the temperature data is not linear in the beginning.

    # Timepoints t_Tc_start / _stop are calculated automatically by checking time when Cernox 1 and 8 transition Tc_top and Tc_bottom. Use this delta to adjust the time interval.
    t_Tc_start_delta = 0
    t_Tc_stop_delta = 0

    fit_indices = np.where((times >= t_Tc_start - t_Tc_start_delta) & (times <= t_Tc_stop + t_Tc_stop_delta))   # Indices of the data array which are in the interesting time interval
    t0 = times[fit_indices[0][0]] #Set t0 to the first timepoint of the fit_indices. Reduces the number of digits in the fit parameters...
    t_end = times[fit_indices[0][-1]]
    t_T1 = times[fit_indices] - t0
    y_T1 = input[0][fit_indices]
    t_T8 = times[fit_indices] - t0
    y_T8 = input[1][fit_indices]
    trans_time = abs(t_T1[-1])
    # Fit the smooth_transition function to the data
    popt_T1, pcov_T1 = curve_fit(smooth_transition, t_T1, y_T1, p0=[y_T1[0], -0.03, 0.1, 0], maxfev=10000)
    popt_T8, pcov_T8 = curve_fit(smooth_transition, t_T8, y_T8, p0=[y_T8[0], -0.03, 0.1, 0], maxfev=10000)
    # popt_T1, pcov_T1_linear = curve_fit(linear_fit, t_T1, y_T1)
    # popt_T8, pcov_T8_linear = curve_fit(linear_fit, t_T8, y_T8)
    y_T1_fit = smooth_transition(t_T1, *popt_T1)
    y_T8_fit = smooth_transition(t_T8, *popt_T8)
    # y_T1_fit = linear_fit(t_T1, *popt_T1)
    # y_T8_fit = linear_fit(t_T8, *popt_T8)   

    # Linear fit for calculating cooldown speed
    popt_T1_linear, pcov_T1_linear = curve_fit(linear_fit, t_T1, y_T1)
    popt_T8_linear, pcov_T8_linear = curve_fit(linear_fit, t_T8, y_T8)
    y_T1_linearfit = linear_fit(t_T1, *popt_T1_linear)
    y_T8_linearfit = linear_fit(t_T8, *popt_T8_linear)


    #Calculate cooldown speed using the linear fit.
    cooldown_speed_T1 = popt_T1_linear[1]
    cooldown_speed_T1_std = np.sqrt(np.diag(pcov_T1_linear))[1]
    cooldown_speed_T8 = popt_T8_linear[1]
    cooldown_speed_T8_std = np.sqrt(np.diag(pcov_T8_linear))[1]
    #Calculate the mean cooldown speed
    cooldown_speed = np.mean([cooldown_speed_T1, cooldown_speed_T8])
    #Calculate the std of the mean cooldown speed
    cooldown_speed_std = np.sqrt(cooldown_speed_T1_std**2 + cooldown_speed_T8_std**2)

    #Calculate the mean and std temperature gradient by calc. difference of y_T1 and y_T8 for each x inbetween t_Tc_start and t_Tc_stop. Here, smooth_transition fit could be crucial!
    Distance = 9.5
    Gradients = (y_T1_fit - y_T8_fit) / Distance
    Gradients = Gradients[np.where((t_T1 + t0 >= t_Tc_start) & (t_T1 +t0 <= t_Tc_stop))]
    mean_gradient = np.mean(Gradients)
    std_gradient = np.std(Gradients)

    if return_indices:
        indices = [fit_indices[0][0], fit_indices[0][-1]]
        return([mean_gradient, mean_gradient, std_gradient, trans_time, cooldown_speed, cooldown_speed_std,
                0, 0, indices, np.append(popt_T1,np.array([t0, t_end])), 
                np.append(popt_T8,np.array([t0, t_end]))]) #pass times with fit parameters for plotting
    else:
        return([mean_gradient, mean_gradient, std_gradient, trans_time, cooldown_speed, cooldown_speed_std, 
                0, 0, np.append(popt_T1,np.array([t0, t_end])), 
                np.append(popt_T8,np.array([t0, t_end]))]) #pass times with fit parameters for plotting

def analyse_B_field(input, times, timestamps):
    '''
    input: np.array of full data set (temperaute,AMRs,fluxgates)
    times: list of times corresponding to the data points
    timestamps: list of timestamps: start_ramp, stop_ramp, save_point
    '''
    #timestamp "stop_ramp" must be adjusted by three seconds because main records the time when the PID controllers are turned off. After that it waits 3 seconds.
    # timestamps[1] = timestamps[1] +datetime.timedelta(0,3)
    
    #find indices corresponding to timestamps
    indices = [find_index_of_first_equal_to(times,timestamp) for timestamp in timestamps]
    print(indices)
    #calculate AMR magnitude               
    AMR_mag = np.zeros(shape = (len(input[:,0]),15))
    for j in range(15):
        for i in range(len(input[:,0])):
            AMR_mag[i,j] = np.sqrt(input[i,j+8]**2 + input[i,j+23]**2 + input[i,j+38]**2)

    #find B_start from fluxgates
    B_start = [input[indices[0],53],input[indices[0],54],input[indices[0],55]]
    B_start.append(np.sqrt(B_start[0]**2+B_start[1]**2+B_start[2]**2)) #magnitude of B_start

    #find expelled flux
    B_expelled =  np.hstack((input[indices[1],8:53], AMR_mag[indices[1],:]))

    #find trapped flux
    B_trapped = np.hstack((input[indices[2],8:53], AMR_mag[indices[2],:]))

    #calculate wave magnitude
    WM_x =  np.max(np.max(abs(input[indices[0]:indices[1],8:23]) - abs(input[indices[0],8:23])))
    WM_y =  np.max(np.max(abs(input[indices[0]:indices[1],23:38]) - abs(input[indices[0],23:38])))
    WM_z =  np.max(np.max(abs(input[indices[0]:indices[1],38:53]) - abs(input[indices[0],38:53])))

    return([B_start, B_expelled, B_trapped, [WM_x, WM_y, WM_z], indices])