import numpy as np
import pandas as pd
from scipy import stats
from sklearn.linear_model import LinearRegression
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error
import logging
from .utils import evaluate_and_plot

logger = logging.getLogger(__name__)

def run_kdm(df_train, df_test, age_col, biom_cols, s2_ba_floor=0.1, out_path=None):
    logger.info("Running Vectorized KDM...")
    k_vals, q_vals, s_vals, r_vals = [], [], [], []
    X_age_train = df_train[[age_col]].values

    for biom in biom_cols:
        y_train = df_train[[biom]].values
        if np.var(y_train) < 1e-10:
            k, q, s, r = 0.0, 0.0, 1.0, 0.0
        else:
            mdl = LinearRegression().fit(X_age_train, y_train)
            k, q = mdl.coef_[0][0], mdl.intercept_[0]
            s = max(np.sqrt(mean_squared_error(y_train, mdl.predict(X_age_train))), 1e-6)
            r = stats.pearsonr(df_train[age_col], df_train[biom])[0]
        k_vals.append(k); q_vals.append(q); s_vals.append(s); r_vals.append(r if not np.isnan(r) else 0.0)

    num = sum((r**2) / np.sqrt(1 - r**2) for r in r_vals if 1e-6 <= abs(r) < 1)
    den = sum(r / np.sqrt(1 - r**2) for r in r_vals if 1e-6 <= abs(r) < 1)
    r_char = np.clip(num / den if abs(den) > 1e-9 else 0.0, -0.9999, 0.9999)

    valid = np.abs(k_vals) > 1e-9
    k_arr, q_arr, s_arr = np.array(k_vals)[valid], np.array(q_vals)[valid], np.array(s_vals)[valid]
    w_arr = (k_arr / s_arr) ** 2

    X_train_bio = df_train[biom_cols].values[:, valid]
    num_train = ((X_train_bio - q_arr) / k_arr) @ w_arr
    den_train = w_arr.sum()
    train_diffs = (num_train / den_train) - df_train[age_col].values
        
    term1 = np.nanvar(train_diffs, ddof=0)
    ca_range = df_train[age_col].max() - df_train[age_col].min()
    term2 = ((1.0 - r_char**2) / (r_char**2)) * (((ca_range**2) / 12.0) / len(biom_cols)) if abs(r_char) > 1e-6 else 0
    s2_ba = max(s2_ba_floor, term1 - term2)

    X_test_bio = df_test[biom_cols].values[:, valid]
    num_test = ((X_test_bio - q_arr) / k_arr) @ w_arr
    ba_e = num_test / den_train
    ba_ec_vals = (ba_e * den_train + df_test[age_col].values / s2_ba) / (den_train + 1.0 / s2_ba)

    df_out = df_test[[age_col]].copy()
    df_out["KDM_BA"] = ba_ec_vals
    
    metrics = evaluate_and_plot(df_out[age_col], df_out["KDM_BA"], "KDM (Test Set)", out_path, "#1f77b4")
    return df_out, metrics

def run_pca_dubina(df_train, df_test, age_col, biom_cols, out_path=None):
    logger.info("Running PCA-Dubina...")
    scaler = StandardScaler()
    z_train = scaler.fit_transform(df_train[biom_cols])
    z_test = scaler.transform(df_test[biom_cols])
    
    ca_train = df_train[age_col].to_numpy(dtype=float)
    ca_test = df_test[age_col].to_numpy(dtype=float)

    pca = PCA(n_components=1)
    bas_train = pca.fit_transform(z_train).flatten()
    bas_test = pca.transform(z_test).flatten()

    if np.corrcoef(bas_train, ca_train)[0, 1] < 0:
        bas_train *= -1
        bas_test *= -1

    bas_train_mu, bas_train_sd = np.nanmean(bas_train), (np.nanstd(bas_train, ddof=0) or 1.0)
    ca_train_mu, ca_train_sd = np.nanmean(ca_train), np.nanstd(ca_train, ddof=0)
    
    ba_train = (bas_train - bas_train_mu) * (ca_train_sd / bas_train_sd) + ca_train_mu
    ba_test = (bas_test - bas_train_mu) * (ca_train_sd / bas_train_sd) + ca_train_mu
    
    b = np.nanmean((ba_train - np.nanmean(ba_train)) * (ca_train - ca_train_mu)) / (np.nanvar(ca_train, ddof=0) + 1e-12)
    bac_test = ba_test + (ca_test - ca_train_mu) * (1.0 - b)

    df_out = df_test[[age_col]].copy()
    df_out["PCA_BA"] = bac_test
    
    metrics = evaluate_and_plot(df_out[age_col], df_out["PCA_BA"], "PCA-Dubina (Test Set)", out_path, "#ff7f0e")
    return df_out, metrics
