import argparse
import os
import joblib
import numpy as np
import pandas as pd
from sklearn import metrics
from sklearn.model_selection import LeaveOneOut
from sklearn.base import clone
from sklearn.preprocessing import StandardScaler

import config
import model_dispatcher

def run(fold, model):
    """
    Run a single fold of the model.
    """
    df = pd.read_pickle(config.TRAIN_DATA_PATH)
    df_train = df[df.kfold != fold].reset_index(drop=True)
    df_valid = df[df.kfold == fold].reset_index(drop=True)

    x_train = df_train.drop("target", axis=1).values
    y_train = df_train.target.values

    x_valid = df_valid.drop("target", axis=1).values
    y_valid = df_valid.target.values

    clf = model_dispatcher.models[model]
    clf.fit(x_train, y_train)
    preds = clf.predict(x_valid)

    rmse = np.sqrt(metrics.mean_squared_error(y_valid, preds))
    print(f"Fold: {fold}, RMSE: {rmse}")
    percentage_error = np.abs(y_valid - preds) / y_valid
    print(f"Fold: {fold}, Percentage Error: {percentage_error}")

    # save the model
    joblib.dump(clf, os.path.join(config.MODEL_PATH, f"{model}_{fold}.bin"))

def run_loocv(model):
    """
    Run Leave-One-Out Cross-Validation for the given model.
    This function is used to evaluate the performance of the model.
    
    This function performs LOOCV where each sample serves as validation set once,
    and trains on all remaining samples. After collecting all predictions,
    it trains a final model on the entire dataset using the same parameters.
    
    Args:
        model (str): Model name from model_dispatcher
    """
    # Load the dataset
    df = pd.read_pickle(config.TRAIN_DATA_PATH)
    
    # Remove kfold column if it exists (not needed for LOOCV)
    if 'kfold' in df.columns:
        df = df.drop('kfold', axis=1)
    
    # Prepare features and target
    X = df.drop("target", axis=1).values
    y = df.target.values

    if model == "huber":
        scaler = StandardScaler()
        X = scaler.fit_transform(X)
    
    print(f"Running LOOCV for {model} with {len(df)} samples...")
    
    # Initialize LOOCV
    loocv = LeaveOneOut()
    
    # Collect predictions for all folds
    all_predictions = []
    all_true_values = []
    fold_rmse_scores = []
    
    # Perform LOOCV
    for fold_idx, (train_idx, val_idx) in enumerate(loocv.split(X)):
        # Get training and validation data for this fold
        x_train_fold = X[train_idx]
        y_train_fold = y[train_idx]
        x_val_fold = X[val_idx]
        y_val_fold = y[val_idx]
        
        # Initialize and train model for this fold
        clf = clone(model_dispatcher.models[model])
        clf.fit(x_train_fold, y_train_fold)
        
        # Make prediction
        pred = clf.predict(x_val_fold)
        
        # Store predictions and true values
        all_predictions.extend(pred)
        all_true_values.extend(y_val_fold)
        
        # Calculate RMSE for this fold (single sample)
        fold_rmse = np.sqrt(metrics.mean_squared_error(y_val_fold, pred))
        fold_rmse_scores.append(fold_rmse)
        
        # Print progress every 10 folds for large datasets
        if (fold_idx + 1) % 10 == 0 or fold_idx == 0:
            print(f"Completed fold {fold_idx + 1}/{len(df)}, RMSE: {fold_rmse}")
    
    # Convert to numpy arrays for easier computation
    all_predictions = np.array(all_predictions)
    all_true_values = np.array(all_true_values)
    
    # Calculate overall metrics using all predictions
    overall_rmse = np.sqrt(metrics.mean_squared_error(all_true_values, all_predictions))
    overall_mae = metrics.mean_absolute_error(all_true_values, all_predictions)
    overall_r2 = metrics.r2_score(all_true_values, all_predictions)
    
    # Calculate percentage errors
    percentage_errors = np.abs(all_true_values - all_predictions) / all_true_values
    mean_percentage_error = np.mean(percentage_errors)
    
    # Print overall results
    print(f"\n=== LOOCV Results for {model} ===")
    print(f"Overall RMSE: {overall_rmse:.4f}")
    print(f"Overall MAE: {overall_mae:.4f}")
    print(f"Overall R²: {overall_r2:.4f}")
    print(f"Mean Percentage Error: {mean_percentage_error:.4f}")
    print(f"Mean of fold-wise RMSE: {np.mean(fold_rmse_scores):.4f}")
    print(f"Std of fold-wise RMSE: {np.std(fold_rmse_scores):.4f}")
    
    # Train final model on entire dataset
    print(f"\nTraining final model on entire dataset...")
    final_model = clone(model_dispatcher.models[model])
    final_model.fit(X, y)
    
    # Save the final model trained on all data
    model_filename = f"{model}_loocv_final.bin"
    model_path = os.path.join(config.MODEL_PATH, model_filename)
    joblib.dump(final_model, model_path)
    print(f"Final model saved to: {model_path}")
    
    # Save LOOCV results for later analysis
    results_dict = {
        'model': model,
        'overall_rmse': overall_rmse,
        'overall_mae': overall_mae,
        'overall_r2': overall_r2,
        'mean_percentage_error': mean_percentage_error,
        'all_predictions': all_predictions,
        'all_true_values': all_true_values,
        'fold_rmse_scores': fold_rmse_scores
    }
    
    results_filename = f"{model}_loocv_results.joblib"
    results_path = os.path.join(config.MODEL_PATH, results_filename)
    joblib.dump(results_dict, results_path)
    print(f"LOOCV results saved to: {results_path}")
    print(f"Training metrics: overall_rmse: {overall_rmse}, overall_mae: {overall_mae}, overall_r2: {overall_r2}, mean_percentage_error: {mean_percentage_error}")
    
    return results_dict

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--fold", type=int, default=0)
    parser.add_argument("--model", type=str, default="lr")
    parser.add_argument("--loocv", action="store_true", help="Run Leave-One-Out Cross-Validation")
    args = parser.parse_args()
    if args.loocv:
        # Run LOOCV
        run_loocv(model=args.model)
    else:
        # Run regular k-fold CV
        run(fold=args.fold, model=args.model)