from . .queries import *
from . .regression import *
from statsmodels.sandbox.regression.predstd import wls_prediction_std
import matplotlib.pyplot as plt
import pandas as pd
import time

def q_drift(log_file_path):
    """
    Ask the user for their choice of correction and keep asking until a valid response is given.
    """
    while True:  # Start an infinite loop
        choice = input("\nApply drift to correction? (Y/N)\n")

        if choice in ["yes", "y", "true", "t", "no", "n", "false", "f"]:
            return choice
            append_to_log(log_file_path, f"Drift application: {choice}")
        else:
            print("\nInvalid response\n")
            
def process_drift_correction(samp,lin_std, drift_std, correction_log, fig_path, log_file_path, isotope):
    user_choice = q_drift(log_file_path)
    chains_to_correct = []
    if neg_response(user_choice): dD_temp = "dD"
    else:
        result = drift_correction(samp, lin_std, drift_std, fig_path, isotope, log_file_path, fig=True)
        if result is not None:
            samp, lin_std, drift_std = result  # Apply the correction
            dD_temp = "drift_corrected_dD"
            correction_log.loc["Drift", "samples"] = 1
        else:
            print("Correction not applied. Reverting to original data.")
            dD_temp = "dD" 
            # Ensure drift_corrected_dD column is created
            samp['drift_corrected_dD'] = samp['dD']
            lin_std['drift_corrected_dD'] = ['dD']
            drift_std['drift_corrected_dD'] = drift_std['dD']
    time.sleep(0)  # Wait for 1 second
    #clear_output(wait=True)
    return samp, lin_std, drift_std, dD_temp, correction_log
    

def drift_correction(samp, lin_std, drift_std, fig_path, isotope, log_file_path, fig=True):
    lin_std = lin_std.copy()
    drift_std = drift_std.copy()
    fig, axs = plt.subplots(1, 2, figsize=(5, 3))
    # Normalize chains
    norm=pd.DataFrame()
    mean_isotope_dict={}
    for i in ["C18","C24"]:
        mask = drift_std.chain==i
        temp = drift_std.loc[mask, [isotope, 'time_rel','chain']].copy()
        mean_isotope_dict[i] = temp[isotope].mean()
        temp[isotope] -= temp[isotope].mean() # center the data at zero
        norm = pd.concat([norm, temp])
    # Perform robust linear regression
    m_d, b_d, r_squared, p_value, std_error, model = wls_regression(norm['time_rel'], norm[isotope], log_file_path)
    new_time_rel = sm.add_constant(np.array(drift_std.time_rel))
    results = model.predict(new_time_rel)
    prstd, pred_ci_lower, pred_ci_upper = wls_prediction_std(model)#, alpha=0.05)
    f3_drift_times = drift_std['time_rel'].unique()[:3]                                             # Find last three samples
    relative_mean = np.mean(model.predict(sm.add_constant(drift_std[drift_std["time_rel"].isin(f3_drift_times)].time_rel))) # estimated dD of last three samples
    drift_correction = results - relative_mean # difference between estimated drift times and lowest replicate mean 
    # Apply drift correction to the drift standards themselves
    drift_std["drift_corrected_dD"] = drift_std[isotope] - drift_correction
    drift_std["drift_error"]        = prstd

    # Plotting if required
    #fig, axs = plt.subplots(1, 2, figsize=(5, 3))
    x = 0
    for i in ["C18","C24"]:
        # Predict values for plotting the regression line
        temp = drift_std[drift_std.chain==i]
        axs[x].scatter(temp['time_rel'], temp[isotope], alpha=0.4, ec='k', s=80, c='blue', label="Drift std")
        axs[x].scatter(temp['time_rel'], temp['drift_corrected_dD'], alpha=0.7, marker='x', s=80, c='red', label="Corrected drift std")
        # Plot regression line
        xmin, xmax = temp['time_rel'].min(), temp['time_rel'].max()
        axs[x].plot([xmin, xmax], [((xmin * m_d + b_d)+temp[temp.time_rel==xmin][isotope].mean()), (xmax * m_d + b_d+temp[temp.time_rel==xmax][isotope].mean())], c='k', linestyle='--', alpha=0.6) # Generalized least-squares regression
        axs[x].set_xlabel("Time relative") 
        axs[x].text(0.9, 0.9, i,
                  horizontalalignment='center', verticalalignment='center',
                  transform=axs[x].transAxes)
        if isotope == "dD": label = "δD (‰)"
        else: label = "δC (‰)"
        axs[x].set_ylabel(label)
        if x == 1:
            axs[x].legend(loc='upper center', bbox_to_anchor=(0, -0.2), fancybox=False, shadow=True,frameon=False, ncol=1)
        x= x+1
        
    # Show the plot before asking for confirmation
    plt.tight_layout()
    #plt.savefig(os.path.join(fig_path, 'Drift correction.png'), dpi=300, bbox_inches='tight')
    plt.show()
    
    # Print regression statistics for each chain
    print("\nRegression statistics drift standards:")
    print("Linear equation: "+str(label)+f" = ({m_d:.2f})(time) + {b_d:.2f}")
    print(f"Adjusted R²: {r_squared:.2f}")
    print(f"P-value: {p_value:.2f}")
    print(f"Standard Error: {std_error:.2f}")

    # Ask user for confirmation
    user_confirmation = input("Apply the correction? (Y/N): \n")
    if pos_response(user_confirmation):
        append_to_log(log_file_path, "Drift correction equation: δD = (time)("+str(m_d)+"+"+str(b_d)+"\nr2 = "+str(r_squared)+"\np = "+str(p_value))
        append_to_log(log_file_path, "User applied drift correction.")
        # Correct linearity standards
        new_time_rel = sm.add_constant(np.array(lin_std.time_rel))
        results = model.predict(new_time_rel)
        prstd, pred_ci_lower, pred_ci_upper = wls_prediction_std(model, exog=new_time_rel)#, alpha=0.05)
        lin_est_corr = results - relative_mean # offset between estimates of linearity samples based on time and the time_zero estimate
        lin_std['drift_corrected_dD'] = lin_std[isotope]-lin_est_corr
        lin_std['drift_error']        = prstd
        
        # Correct samples 
        new_time_rel = sm.add_constant(np.array(samp.time_rel))
        results = model.predict(new_time_rel)
        prstd, pred_ci_lower, pred_ci_upper = wls_prediction_std(model, exog =new_time_rel)#, alpha=0.05)
        samp_est_corr = results - relative_mean # offset between estimates of linearity samples based on time and the time_zero estimate
        samp_est_corr = results-relative_mean
        samp['drift_corrected_dD'] = samp[isotope]-samp_est_corr
        samp['drift_error']        = prstd
        return samp, lin_std, drift_std
    else:
        return None
