import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from copy import deepcopy
import os
import sys
import argparse
import pickle
import time
import logging
# Define models - exactly as in original
from sklearn.linear_model import RidgeCV, LassoCV, ElasticNetCV
from xgboost import XGBRegressor
from joblib import Parallel, delayed
from quantile_forest import RandomForestQuantileRegressor, ExtraTreesQuantileRegressor
from pygam import ExpectileGAM
from utils import convert_to_serializable, get_intervals_manually, optimize_gam_smoothing, setup_logging, print_model_performance, WorkerLogger

# Update imports to use the new PCS implementation
# Try to import from the installed package first, fall back to local path if not available
try:
    from PCS.regression import PCS_UQ, PCS_OOB
except ImportError:
    # Add PCS_UQ to path if needed
    pcs_uq_path = os.path.join("..", "..", "PCS_UQ", "src")
    if pcs_uq_path not in sys.path:
        sys.path.append(pcs_uq_path)
    from PCS.regression import PCS_UQ, PCS_OOB

def run_single_experiment_pcs(dataset_name, seed_index, n_boot=50, base_data_path=None, alphas=[0.1, 0.05, 0.01], top_k=1, use_oob=False, calibration_method='multiplicative', conformalized=False, qxgb_only=False):
    """
    Run a single PCS ensemble experiment for a dataset using the specified seed index.
    
    Args:
        dataset_name: Name of the dataset to process (e.g., 'data_ailerons' or just 'ailerons')
        seed_index: The index to use for generating the seed (seed = 777 + seed_index)
        n_boot: Number of bootstrap samples
        base_data_path: Base directory path where data is located
        alphas: List of alpha values to use for calibration (default: [0.1, 0.05, 0.01])
        top_k: Number of top-performing models to use (default: 1)
        use_oob: Whether to use PCS_OOB instead of PCS_UQ (default: False)
        calibration_method: Method for calibration, 'multiplicative' or 'additive' (default: 'multiplicative')
        conformalized: Whether to use separate calibration set (conformalized) or reuse validation set (standard) (default: False)
        qxgb_only: If True, use only XGBoost quantile estimator (default: False)
    
    Returns:
        Dict containing model results and trained models for each alpha value
    """
    # Create a worker logger to capture logs
    worker_logger = WorkerLogger()
    
    # Normalize base path
    if base_data_path is None:
        base_data_path = os.getcwd()
    
    # Make sure dataset name has 'data_' prefix
    if not dataset_name.startswith("data_"):
        dataset_name = f"data_{dataset_name}"
    
    # List all possible paths where the data might be
    search_paths = [
        # Your current pwd shows you're in F:\research\clear\data
        os.path.join(base_data_path, dataset_name),
        os.path.join(base_data_path, dataset_name),
        # Try parent directories
        os.path.join(os.path.dirname(base_data_path), dataset_name),
        os.path.join(os.path.dirname(base_data_path), "data", dataset_name)
    ]
    
    # Find an existing path that contains the dataset
    data_dir = next((path for path in search_paths if os.path.exists(path)), None)
    
    if not data_dir:
        worker_logger.log(f"Searched for dataset in these directories: {search_paths}")
        raise FileNotFoundError(f"Could not find directory for {dataset_name}")
    
    # Now construct paths to the X and y files
    X_path = os.path.join(data_dir, "X.csv")
    y_path = os.path.join(data_dir, "y.csv")
    
    # Verify files exist
    if not os.path.exists(X_path) or not os.path.exists(y_path):
        worker_logger.log(f"Found directory {data_dir} but missing X.csv or y.csv files")
        raise FileNotFoundError(f"Missing data files in {data_dir}")
    
    worker_logger.log(f"Using X data from: {X_path}")
    worker_logger.log(f"Using y data from: {y_path}")
    
    # Read the data
    X = pd.read_csv(X_path)
    y = np.loadtxt(y_path, delimiter=',')
    
    # Calculate seed from index - exactly as in original
    seed = 777 + seed_index
    
    # Train, validation, and test split
    x_train_val, x_test, y_train_val, y_test = train_test_split(X, y, test_size=0.2, random_state=seed)
    
    if conformalized:
        # For conformalized: split train_val into train (75%), val (12.5%), calib (12.5%)
        # This gives us 60% train, 10% val, 10% calib, 20% test of original data
        x_train_temp, x_val_calib, y_train_temp, y_val_calib = train_test_split(
            x_train_val, y_train_val, test_size=0.25, random_state=seed)
        x_val, x_calib, y_val, y_calib = train_test_split(
            x_val_calib, y_val_calib, test_size=0.5, random_state=seed + 1)
        x_train, y_train = x_train_temp, y_train_temp
        
        # Create split_dict for conformalized version with separate calibration set
        split_dict = {
            "x_train": x_train.to_numpy(),
            "y_train": y_train,
            "x_val": x_val.to_numpy(),
            "y_val": y_val,
            "x_calib": x_calib.to_numpy(),
            "y_calib": y_calib,
            "x_test": x_test.to_numpy(),
            "y_test": y_test
        }
        worker_logger.log(f"[Seed {seed_index}] Conformalized mode: Train={len(y_train)}, Val={len(y_val)}, Calib={len(y_calib)}, Test={len(y_test)}")
    else:
        # Standard approach: reuse validation set as calibration set
        x_train, x_val, y_train, y_val = train_test_split(x_train_val, y_train_val, test_size=0.25, random_state=seed)
        
        # Create split_dict for standard version (val = calib)
        split_dict = {
            "x_train": x_train.to_numpy(),
            "y_train": y_train,
            "x_val": x_val.to_numpy(),
            "y_val": y_val,
            "x_calib": x_val.to_numpy(),  # Reuse validation set
            "y_calib": y_val,             # Reuse validation set
            "x_test": x_test.to_numpy(),
            "y_test": y_test
        }
        worker_logger.log(f"[Seed {seed_index}] Standard mode: Train={len(y_train)}, Val/Calib={len(y_val)}, Test={len(y_test)}")
    
    # Optimize GAM smoothing parameters using training data
    n_splines = 10
    
    # Define quantile estimators based on qxgb_only flag
    if qxgb_only:
        # Use only XGBoost for qxgb variant
        quantile_estimators = {
            "QXGB": XGBRegressor(random_state=777, objective='reg:quantileerror', n_estimators=100, tree_method='hist', quantile_alpha=0.5, min_child_weight = 10)
        }
        worker_logger.log(f"[Seed {seed_index}] QXGB-only mode: Using only XGBoost quantile estimator")
        # Set default values for GAM parameters when not using GAM
        optimal_lam = None
        gam_converged = False
    else:
        # Define basic quantile estimators first (will always have these)
        quantile_estimators = {
            "QRF": RandomForestQuantileRegressor(random_state=777, n_estimators=100, default_quantiles=0.5, min_samples_leaf = 10),
            "QXGB": XGBRegressor(random_state=777, objective='reg:quantileerror', n_estimators=100, tree_method='hist', quantile_alpha=0.5, min_child_weight = 10)
        }
        
        # Set timeout to 5 minutes (300 seconds)
        optimal_lam, gam_converged = optimize_gam_smoothing(split_dict["x_train"], split_dict["y_train"], n_splines=n_splines, timeout=300)
        
        if gam_converged:
            worker_logger.log(f"[Seed {seed_index}] ExpectileGAM converged with optimal smoothing parameter (lam): {optimal_lam:.4f}")
            # Add ExpectileGAM to the models if it converged successfully
            quantile_estimators["ExpectileGAM"] = ExpectileGAM(expectile=0.5, n_splines=n_splines, lam=optimal_lam, spline_order=3, basis='ps')
        else:
            worker_logger.log(f"[Seed {seed_index}] ExpectileGAM did not converge within timeout. Will use only {', '.join(quantile_estimators.keys())} models.")

    # Print the models we're using for this run
    worker_logger.log(f"[Seed {seed_index}] Using models: {', '.join(quantile_estimators.keys())}")
    
    # Create a dictionary to store results for each alpha
    all_results = {}
    
    # Process each alpha value
    for alpha in alphas:
        worker_logger.log(f"[Seed {seed_index}] Training and calibrating for alpha={alpha} (confidence level: {(1-alpha)*100}%)")
        
        # Create the appropriate PCS object
        if use_oob:
            worker_logger.log(f"[Seed {seed_index}] Using OOB method with top_k={top_k}")
            pcs_model = PCS_OOB(
                models=quantile_estimators,
                num_bootstraps=n_boot,
                alpha=alpha,
                seed=seed,
                top_k=top_k,
                calibration_method=calibration_method
            )
            # Fit the model on train+val combined data
            pcs_model.fit(
                X = np.concatenate([split_dict["x_train"], split_dict["x_val"]]), 
                y = np.concatenate([split_dict["y_train"], split_dict["y_val"]])
            )
        else:
            pcs_model = PCS_UQ(
                models=quantile_estimators,
                num_bootstraps=n_boot,
                alpha=alpha,
                seed=seed,
                top_k=top_k,
                calibration_method=calibration_method
            )
            pcs_model.fit(X = split_dict["x_train"], y = split_dict["y_train"], X_calib = split_dict["x_val"], y_calib = split_dict["y_val"])
        
        worker_logger.log(f"[Seed {seed_index}] Calibration complete. Gamma value: {pcs_model.gamma}")
        
        # Print model performance rankings to the worker logger
        print_model_performance(pcs_model, worker_logger)

        if use_oob:
            worker_logger.log(f"[Seed {seed_index}] Using OOB method with custom get_intervals")
            
            # Use get_intervals on train+val combined data
            train_intervals_raw = get_intervals_manually(
                pcs_model, 
                split_dict["x_train"],
                alpha=alpha
            )
            
            val_intervals_raw = get_intervals_manually(
                pcs_model, 
                split_dict["x_val"],
                alpha=alpha
            )
            
            calib_intervals_raw = get_intervals_manually(
                pcs_model, 
                split_dict["x_calib"],
                alpha=alpha
            )
            
            # Use the same method for test data directly
            test_intervals_raw = get_intervals_manually(pcs_model, split_dict['x_test'], alpha=alpha)
            
            # Get calibrated intervals
            train_intervals = pcs_model.predict(split_dict['x_train'])
            val_intervals = pcs_model.predict(split_dict['x_val'])
            calib_intervals = pcs_model.predict(split_dict['x_calib'])
            test_intervals = pcs_model.predict(split_dict['x_test'])
            
            worker_logger.log(f"[Seed {seed_index}] OOB intervals and medians calculated using custom method")
        else:
            # For standard PCS_UQ, use the built-in methods
            train_intervals_raw = pcs_model.get_intervals(split_dict['x_train'])
            val_intervals_raw = pcs_model.get_intervals(split_dict['x_val'])
            calib_intervals_raw = pcs_model.get_intervals(split_dict['x_calib'])
            test_intervals_raw = pcs_model.get_intervals(split_dict['x_test'])
            
            # Get calibrated intervals
            train_intervals = pcs_model.predict(split_dict['x_train'])
            val_intervals = pcs_model.predict(split_dict['x_val'])
            calib_intervals = pcs_model.predict(split_dict['x_calib'])
            test_intervals = pcs_model.predict(split_dict['x_test'])
            
            worker_logger.log(f"[Seed {seed_index}] Standard intervals and medians calculated")
        
        # Store results
        pcs_ensemble = {
            'x_train': split_dict['x_train'], 
            'y_train': split_dict['y_train'],
            'x_val': split_dict['x_val'], 
            'y_val': split_dict['y_val'],
            'x_calib': split_dict['x_calib'],
            'y_calib': split_dict['y_calib'],
            'x_test': split_dict['x_test'], 
            'y_test': split_dict['y_test'],
            # Raw intervals for all datasets
            'train_intervals_raw': train_intervals_raw,
            'val_intervals_raw': val_intervals_raw,
            'calib_intervals_raw': calib_intervals_raw,
            'test_intervals_raw': test_intervals_raw,
            # Calibrated intervals for all datasets
            'train_intervals': train_intervals,
            'val_intervals': val_intervals,
            'calib_intervals': calib_intervals,
            'test_intervals': test_intervals,
            # Median predictions for all datasets
            # 'train_median_preds': train_median_preds,
            # 'val_median_preds': val_median_preds,
            # 'test_median_preds': test_median_preds,
            'alpha': alpha,
            'confidence_level': f"{(1-alpha)*100:.0f}",
            'gamma': pcs_model.gamma,
            'top_k': top_k,
            'calibration_method': calibration_method,
            'use_oob': use_oob,
            'conformalized': conformalized,
            'top_model_names': list(pcs_model.top_k_models.keys()),
            'model_performances': {model: score for model, score in sorted(pcs_model.pred_scores.items(), key=lambda x: x[1], reverse=True)},
            'gam_parameters': {'lam': optimal_lam, 'n_splines': n_splines},
            'gam_converged': gam_converged
        }
        
        all_results[f"alpha_{alpha}"] = pcs_ensemble
    
    # Return both results and logs
    return {
        'results': all_results,
        'logs': worker_logger.get_logs(),
        'seed_index': seed_index
    }

def run_single_experiment_pcs_and_save(dataset_name, seed_index, n_boot=50, base_data_path=None, save_path_template=None, alphas=[0.1, 0.05, 0.01], top_k=1, use_oob=False, calibration_method='multiplicative', conformalized=False, qxgb_only=False):
    """Modified function that runs the experiment and saves result directly to disk"""
    full_result = run_single_experiment_pcs(dataset_name, seed_index, n_boot, base_data_path, alphas, top_k, use_oob, calibration_method, conformalized, qxgb_only)
    results = full_result['results']
    
    # Save results for each alpha to individual files
    if save_path_template:
        for alpha_key, result in results.items():
            alpha = result['alpha']
            confidence_level = result['confidence_level']
            save_path = save_path_template.format(f"{dataset_name}_run_{seed_index}_{confidence_level}")
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            
            # Convert data to a serializable format before saving
            serializable_result = convert_to_serializable(result)
            with open(save_path, "wb") as f:
                pickle.dump(serializable_result, f)
        
    return full_result

if __name__ == "__main__":
    # Default list of datasets
    default_datasets = [
        'data_computer',
        'data_ailerons',
        'data_airfoil',
        'data_ca_housing',
        'data_concrete',
        'data_elevator',
        'data_energy_efficiency',
        'data_insurance',
        'data_kin8nm',
        'data_miami_housing',
        'data_naval_propulsion',
        'data_parkinsons',
        'data_powerplant',
        'data_sulfur',
        'data_superconductor',
        'data_qsar',
        'data_allstate', # also takes long computationally, but we have maintained it
        # 'data_diamond', #disabled for computational reasons
    ]

    parser = argparse.ArgumentParser(description='Run PCS ensemble experiments')
    parser.add_argument('--datasets', type=str, nargs='?', default=','.join(default_datasets), 
                        help='Comma-separated list of dataset names (default: all datasets)')
    parser.add_argument('--seed_index', type=int, default=0, help='Seed index')
    parser.add_argument('--n_boot', type=int, default=100, help='Number of bootstrap samples')
    parser.add_argument('--n_jobs', type=int, default=5, help='Number of jobs for parallelizing the training')
    parser.add_argument('--project_root', type=str, default=os.path.join("..", ".."), help='Base path for the data')
    parser.add_argument('--alphas', type=str, default='0.05', 
                       help='Comma-separated list of alpha values for calibration (default: 0.1,0.05,0.01)')
    parser.add_argument('--top_k', type=int, default=1, help='Number of top-performing models to use (default: 1)')
    parser.add_argument('--use_oob', action='store_true', help='Use PCS_OOB instead of PCS_UQ')
    parser.add_argument('--calibration_method', type=str, choices=['multiplicative', 'additive'], 
                       default='multiplicative', help='Method for calibration (default: multiplicative)')
    parser.add_argument('--conformalized', action='store_true', 
                       help='Use separate calibration set (conformalized) instead of reusing validation set (standard)')
    parser.add_argument('--qxgb_only', action='store_true', 
                       help='Use only XGBoost quantile estimator (qxgb variant) instead of full PCS ensemble')
    parser.add_argument('--n_runs', type=int, default=10, help='Number of runs for each dataset (default: 5)')
    parser.add_argument('--log_file', type=str, default=None, 
                       help='Path to log file (default: logs/pcs_training_<timestamp>.log)')
    args = parser.parse_args()

    # Set up logging
    logger = setup_logging(args.log_file)
    logging.info("=" * 80)
    mode_type = "QXGB-only" if args.qxgb_only else "PCS ensemble"
    logging.info(f"Starting {mode_type} training with the following parameters:")
    logging.info(f"- Datasets: {args.datasets}")
    logging.info(f"- Seed index: {args.seed_index}")
    logging.info(f"- Bootstrap samples: {args.n_boot}")
    logging.info(f"- Jobs: {args.n_jobs}")
    logging.info(f"- Project root: {args.project_root}")
    logging.info(f"- Use OOB: {args.use_oob}")
    logging.info(f"- Calibration method: {args.calibration_method}")
    logging.info(f"- Conformalized: {args.conformalized}")
    logging.info(f"- QXGB only: {args.qxgb_only}")
    logging.info(f"- Top k: {args.top_k}")
    logging.info("=" * 80)

    # Parse the datasets from command-line input (or use the default list)
    datasets = args.datasets.split(',')
    
    # Parse alpha values
    alphas = [float(alpha) for alpha in args.alphas.split(',')]
    logging.info(f"Using alpha values: {alphas} (confidence levels: {[(1-alpha)*100 for alpha in alphas]}%)")

    data_base_path = os.path.join(args.project_root, "data")
    approach_suffix = "conformalized" if args.conformalized else "standard"
    # Naming convention:
    #  - QXGB-only:    models/pcs_top{top_k}_qxgb_{n_runs}_{approach}
    #  - PCS ensemble: models/pcs_top{top_k}_qpcs_{n_runs}_{approach}
    model_prefix = "qxgb" if args.qxgb_only else "qpcs"
    model_save_path = os.path.join(
        args.project_root,
        "models",
        f"pcs_top{args.top_k}_{model_prefix}_{args.n_runs}_{approach_suffix}",
        "{}_pcs_results_{}.pkl",
    )
    # Create results directory if it doesn't exist
    os.makedirs(os.path.dirname(model_save_path.format("", "")), exist_ok=True)

    logging.info("Starting to process datasets...")
    start_time = time.time()
    #---------------------------------------------------
    # Option 1: Process everything at once
    # Process each dataset
    for dataset_name in datasets:
        dataset_name = dataset_name.strip()
        logging.info(f"Processing dataset: {dataset_name}")

        # Special handling for Allstate due to convergence issues
        n_jobs_to_use = 1 if dataset_name == 'data_allstate' else args.n_jobs
        
        all_results = Parallel(n_jobs=n_jobs_to_use)(
            delayed(run_single_experiment_pcs)(
                dataset_name, i, 
                n_boot=args.n_boot, 
                base_data_path=data_base_path, 
                alphas=alphas,
                top_k=args.top_k,
                use_oob=args.use_oob,
                calibration_method=args.calibration_method,
                conformalized=args.conformalized,
                qxgb_only=args.qxgb_only
            )
            for i in range(args.n_runs)
        )
        
        # Print worker logs from each seed
        for result in all_results:
            seed_idx = result['seed_index']
            logging.info(f"\n----- Logs from worker process (seed {seed_idx}) -----")
            for log_line in result['logs']:
                logging.info(log_line)
        
        # For each alpha, create a separate results file
        for alpha in alphas:
            confidence_level = f"{(1-alpha)*100:.0f}"
            results_for_alpha = {f'run_{i}': all_results[i]['results'][f"alpha_{alpha}"] for i in range(args.n_runs)}
            
            # Convert to serializable before saving
            serializable_results = convert_to_serializable(results_for_alpha)
            
            # Save results
            with open(model_save_path.format(dataset_name, confidence_level), "wb") as f:
                pickle.dump(serializable_results, f)
            
            logging.info(f"Results saved for dataset {dataset_name} at confidence level {confidence_level}")

        # end time
        end_time = time.time()
        logging.info(f"Total time taken for {dataset_name}: {end_time - start_time} seconds")

    logging.info("All processing complete!")


# To run the experiments, use the following command (for our paper):
# python .\train_pcs_quantile.py --top_k 1
## Alternatively, for double-logging for each top_k, use the following command:
## python .\train_pcs_quantile.py --top_k 1 *> logs/pcs_training_top1.log ; python .\train_pcs_quantile.py --top_k 2 *> logs/pcs_training_top2.log