import pandas as pd
import numpy as np
import os
import matplotlib.pyplot as plt
from . .regression import wls_regression
from . .queries import *
import statsmodels.api as sm
from statsmodels.sandbox.regression.predstd import wls_prediction_std
from . .figures import standard_check_figures


# VSMOW
def vsmow_correction(cfg, unknown, lin, drift, correction_log, folder_path,fig_path, log_file_path, isotope, vsmow, fig=True):
    print("Applying VSMOW correction")
    if isotope == "dD":
        label="dD"
    else:
        label = "dC"
    stds  = pd.concat([lin, drift])
    standards_iso = list(vsmow[vsmow['VSMOW accuracy check']==False]['ID'])
    standards_chain = list(vsmow[vsmow['VSMOW accuracy check']==False]['chain length'])
    vsmow_check = vsmow[vsmow['VSMOW accuracy check']==True]['chain length']
    vsmow_check = np.array(vsmow_check).flat[0]
    for i, chain in zip(standards_iso, standards_chain):
        # mask = stds.chain==i
        mask = (stds['Identifier 1'].str.contains(i))&(stds['chain']==chain)
        stds.loc[mask, 'VSMOW_dD_actual']    = vsmow.loc[vsmow['ID'] == i, 'isotope value'].iloc[0]
        stds.loc[mask, 'vsmow_error_actual'] = vsmow.loc[vsmow['ID'] == i, 'std'].iloc[0]
    dD_id = cfg.dD_col
    stds=stds[~stds[dD_id].isna()]
    mask = stds.chain.isin(standards_chain)
 
    # Create uncertainty for standards for weighting regression function
    stds['replicate_uncertainty'] = np.nan
    # for j in vsmow[vsmow['VSMOW accuracy check']==False]['chain length']:
    #     chain_mask = stds['chain']==j
    #     value = np.std(stds.loc[chain_mask, dD_id])
    #     stds.loc[chain_mask, 'replicate_uncertainty'] = value
    for chain_l, std_id in zip(vsmow[vsmow['VSMOW accuracy check']==False]['chain length'], vsmow[vsmow['VSMOW accuracy check']==False]['ID']):
        chain_mask = (stds['chain']==chain_l)&(stds['Identifier 1'].str.contains(std_id))
        value = np.std(stds.loc[chain_mask, dD_id])
        stds.loc[chain_mask, 'replicate_uncertainty'] = value
    # # Weighted least square regression
    x = stds.loc[mask, dD_id]
    y = stds.loc[mask,'VSMOW_dD_actual']
    sigma = np.sqrt(stds.loc[mask, 'vsmow_error_actual']**2+stds.loc[mask, 'replicate_uncertainty']**2)
    w = stds.loc[mask, 'area']#1/(sigma**2)
    w = 1/(sigma**2)

    m_d, b_d, r_squared, p_value, std_error, model = wls_regression(x,y, log_file_path,weights=w)

    # Predict VSMOW values - standards
    """
    Predicted error is: standard error of the prediction, from the scipy package wls_prediction_std.
    This is the correct prediction error associated with the prediction of weighted least squares (wls)
    regression.
    """
    new_x = sm.add_constant(np.array(stds[dD_id]))
    stds['VSMOW_dD'] = model.predict(new_x)
    prstd, std_pred_ci_lower, std_pred_ci_upper = wls_prediction_std(model, exog = new_x)
    stds['VSMOW_error'] = prstd

    # Predict VSMOW values - samples
    new_x = sm.add_constant(np.array(unknown[dD_id]))
    unknown['VSMOW_dD'] = model.predict(new_x)
    prstd, unknown_pred_ci_lower, unknown_pred_ci_upper = wls_prediction_std(model, exog = new_x)
    unknown['VSMOW_error'] = prstd
    relative_standard = vsmow[vsmow['VSMOW accuracy check']==True]['isotope value'].values
    rs_unc =  vsmow[vsmow['VSMOW accuracy check']==True]['std'].values
    
    # Figure
    fig = plt.figure()
    plt.scatter(stds.loc[mask,dD_id], stds.loc[mask,'VSMOW_dD_actual'], c='red', ec='k', label=f'VSMOW {standards_iso}', s=80, zorder=2)
    plt.plot(stds.loc[mask, dD_id], stds.loc[mask, dD_id]*m_d + b_d, c='k', linestyle='-', alpha=0.6, zorder=0)
    plt.scatter(unknown[dD_id], unknown['VSMOW_dD'], c='k', marker='x', alpha=0.5, label = 'corrected samples', zorder=1)
    if isotope=='dD':
        mask = stds.chain==vsmow_check
        plt.scatter(stds.loc[mask, dD_id], stds.loc[mask,'VSMOW_dD'], c='blue', ec='k', label=f'Predicted {vsmow_check}', zorder=99, s = 80)
        # plt.axhline(relative_standard,c='blue',linestyle='--',zorder=0,label=f"Actual VSMOW {vsmow_check}")
        ymin = float(np.squeeze(relative_standard - rs_unc))
        ymax = float(np.squeeze(relative_standard + rs_unc))
        plt.axhspan(ymin=ymin, ymax=ymax, color='blue', alpha=0.4, zorder=1, label=f"Actual VSMOW {vsmow_check}")
    plt.legend(loc='lower center', bbox_to_anchor=(0.5, -0.3), frameon=False, ncol=2)
    plt.xlabel("Measured "+str(label)+" (‰)")
    plt.ylabel("Predicted VSMOW "+str(label)+" (‰)")
    plt.savefig(os.path.join(fig_path, 'VSMOW correction.png'), dpi=300, bbox_inches='tight')
    plt.show()

    print("\nRegression statistics VSMOW 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}")
    print(f"\nAccuracy of {vsmow_check} prediction")
    mask = stds.chain==vsmow_check # "C24"
    test = stds.loc[mask,'VSMOW_dD']
    mae = (test - relative_standard).abs().mean()
    rmse = np.sqrt(((test - (relative_standard)) ** 2).mean())
    print(f"RMSE: {rmse:.2f}")

    # Uncertainty for samples
    #stds['vsmow_error'] = std_error
    if 'drift_error' not in stds.columns:
        stds['drift_error'] = np.nan
    if 'linearity_error' not in stds.columns:
        stds['linearity_error'] = np.nan

    append_to_log(log_file_path, 'VSMOW')
    append_to_log(log_file_path, f"- Linear equation: {label})  = ({m_d:.2f})(time) + {b_d:.2f}")
    append_to_log(log_file_path, f"- Adjusted R²: {r_squared:.2f}")
    append_to_log(log_file_path, f"- p-value: {p_value:.2f}")
    append_to_log(log_file_path, f"- Standard Error: {std_error:.2f}")
    append_to_log(log_file_path, f"- Accuracy of C24 prediction (RMSE): {rmse:.2f}")
    standard_check_figures(cfg, stds, fig_path, label, vsmow)
    return unknown, stds


