import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import ElasticNetCV, ElasticNet
from sklearn.ensemble import RandomForestRegressor
from xgboost import XGBRegressor
from sklearn.model_selection import GridSearchCV
import logging
from .utils import evaluate_and_plot

logger = logging.getLogger(__name__)

def run_ml_models(df_train, df_test, age_col, biom_cols, random_state=42, ml_params=None, out_dir=None):
    logger.info("Running Machine Learning Models...")
    scaler = StandardScaler()
    X_train_s = scaler.fit_transform(df_train[biom_cols])
    X_test_s = scaler.transform(df_test[biom_cols])
    y_train = df_train[age_col]
    
    results_df = pd.DataFrame({"Chronological_Age": df_test[age_col]}, index=df_test.index)
    ml_metrics = []
    
    params = ml_params or {}
    grid_search_enabled = params.get("grid_search", False)
    cv_folds = params.get("cv_folds", 5)

    models = {}
    
    # Elastic Net
    en_params = params.get("elastic_net", {})
    if en_params.get("run", True):
        if grid_search_enabled and "param_grid" in en_params:
            models["Elastic Net"] = (GridSearchCV(ElasticNet(random_state=random_state), en_params["param_grid"], cv=cv_folds, scoring='neg_mean_squared_error', n_jobs=-1), "#2ca02c")
        else:
            l1 = en_params.get("l1_ratio", 0.5)
            models["Elastic Net"] = (ElasticNetCV(cv=cv_folds, l1_ratio=l1, random_state=random_state, n_jobs=-1), "#2ca02c")

    # Random Forest
    rf_params = params.get("random_forest", {})
    if rf_params.get("run", True):
        if grid_search_enabled and "param_grid" in rf_params:
            models["Random Forest"] = (GridSearchCV(RandomForestRegressor(random_state=random_state), rf_params["param_grid"], cv=cv_folds, scoring='neg_mean_squared_error', n_jobs=-1), "#d62728")
        else:
            n_est = rf_params.get("n_estimators", 100)
            max_d = rf_params.get("max_depth", None)
            models["Random Forest"] = (RandomForestRegressor(n_estimators=n_est, max_depth=max_d, random_state=random_state, n_jobs=-1), "#d62728")

    # XGBoost
    xgb_params = params.get("xgboost", {})
    if xgb_params.get("run", True):
        if grid_search_enabled and "param_grid" in xgb_params:
            models["XGBoost"] = (GridSearchCV(XGBRegressor(random_state=random_state), xgb_params["param_grid"], cv=cv_folds, scoring='neg_mean_squared_error', n_jobs=-1), "#9467bd")
        else:
            n_est = xgb_params.get("n_estimators", 100)
            lr = xgb_params.get("learning_rate", 0.1)
            max_d = xgb_params.get("max_depth", 6)
            models["XGBoost"] = (XGBRegressor(n_estimators=n_est, learning_rate=lr, max_depth=max_d, random_state=random_state, n_jobs=-1), "#9467bd")

    for name, (model, color) in models.items():
        logger.info(f"Training {name}...")
        model.fit(X_train_s, y_train)
        
        if grid_search_enabled and hasattr(model, 'best_params_'):
            logger.info(f"[Tuned] Best params for {name}: {model.best_params_}")
            
        preds = model.predict(X_test_s)
        results_df[f"{name}_BA"] = preds
        
        plot_path = out_dir / f"plot_ml_{name.replace(' ', '_').lower()}.png" if out_dir else None
        metrics = evaluate_and_plot(results_df["Chronological_Age"], preds, f"{name} (Test Set)", plot_path, color)
        ml_metrics.append(metrics)

    return results_df, ml_metrics
