train_user_level.py

#!/usr/bin/env python3
"""
Question-Answer Pair Model Training and Output Saving with Question Strategies
Enhanced with proper hyperparameter tuning and user-level evaluation
"""

import os

# Torch 2.x optimizer initialization can import torch._dynamo, which imports
# optree. Preload optree's C-extension first to avoid late ABI surprises.
try:
    import optree._C  # noqa: F401
except Exception as e:
    msg = str(e)
    if "CXXABI_1.3.15" in msg or "libstdc++.so.6" in msg:
        raise RuntimeError(
            "Failed to load optree C-extension due to libstdc++ ABI mismatch. "
            "Activate the conda environment and ensure its runtime libraries are used. "
            "A common fix is: conda install -n maqua -c conda-forge libstdcxx-ng libgcc-ng"
        ) from e
    raise

import torch
import torch.nn as nn
import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler, StandardScaler, RobustScaler
from sklearn.model_selection import ParameterGrid
from sklearn.linear_model import LinearRegression, Ridge
from scipy.stats import pearsonr
from tqdm import tqdm
import json
import random
import argparse
import warnings
from maqua_paths import FORMATTED_DATA, MODEL_OUTPUTS, ensure_dir
warnings.filterwarnings('ignore')

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class Config:
    def __init__(self):
        self.random_seed = 42
        self.n_folds = 9
        
        # Hyperparameter tuning configuration
        self.hyperparam_tuning = True
        self.max_epochs = 1000
        self.patience = 20
        
        # Model configuration
        self.embedding_type = "full_embedding"
        self.num_outputs = 10
        self.aggregation_type = "all_question_aggregate"
        
        # Model type configuration
        self.model_type = "multitask"  # Options: 'multitask', 'singletask', 'both'
        
        # Question strategy configuration
        self.question_strategy = "all_questions"  # Options: 'specific_only', 'general_only', 'combined', 'all_questions'
        
        # Outcome mapping
        self.outcome_names = ["PHQ", "GAD", "MDQ", "RAADS", "DUDIT", "AUDIT", "BOCS", "ASRS", "NSE", "EDE_QS"]
        self.singletask_outcomes = list(self.outcome_names)
        
        # Hyperparameter search space (only Adam optimizer and MinMax scaler)
        self.param_grid = {
            'lr': [0.001, 0.01, 0.1],
            'weight_decay': [0.0, 0.001, 0.01]
        }
        
    def get_config_name(self):
        """Generate a descriptive name for this configuration"""
        if self.model_type == "multitask":
            return f"{self.embedding_type}_multitask_{self.num_outputs}outputs_{self.aggregation_type}_{self.question_strategy}"
        elif self.model_type == "singletask":
            return f"{self.embedding_type}_singletask_{self.aggregation_type}_{self.question_strategy}"
        else:  # both
            return f"{self.embedding_type}_both_{self.aggregation_type}_{self.question_strategy}"

class MultitaskLinearModel(nn.Module):
    def __init__(self, input_dim, num_outcomes=10):
        super().__init__()
        self.linear = nn.Linear(input_dim, num_outcomes)
        
    def forward(self, x):
        return self.linear(x)

class SingletaskLinearModel(nn.Module):
    def __init__(self, input_dim):
        super().__init__()
        self.linear = nn.Linear(input_dim, 1)  # Single output
        
    def forward(self, x):
        return self.linear(x)

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)

class DataProcessor:
    def __init__(self, config):
        self.config = config
        
    def load_data(self):
        # Load question embeddings
        question_data = pd.read_csv(FORMATTED_DATA / "question_embeddings.csv")
        question_data.set_index("code", inplace=True)
        
        # Load input data
        input_essay = pd.read_csv(FORMATTED_DATA / "input_dataset_essay.csv")
        input_words = pd.read_csv(FORMATTED_DATA / "input_dataset_words.csv")
        input_data = pd.concat([input_words, input_essay], ignore_index=True)
        
        # Load output data
        output_data = pd.read_csv(FORMATTED_DATA / "output_questionnaire_outcomes.csv")
        
        # Load question strategy mapping
        questions_mapping = self.load_question_strategy_mapping()
        
        return input_data, output_data, question_data, questions_mapping
    
    def load_question_strategy_mapping(self):
        """Load question mapping with strategy filtering"""
        try:
            questions_df = pd.read_csv(FORMATTED_DATA / "questions_symptom_outcomes_type.csv")
        except FileNotFoundError:
            questions_df = pd.read_csv(FORMATTED_DATA / "questions.csv")
        
        # Get all general questions (those with symptom == 'general')
        general_questions = questions_df[questions_df['symptom'] == 'general']['code'].tolist()
        
        # Get all questions regardless of symptom type
        all_questions = questions_df['code'].tolist()
        
        # Create mappings for specific symptoms for each outcome
        outcome_to_specific_symptoms = {
            'PHQ': ['depression'],
            'GAD': ['anxiety'], 
            'MDQ': ['bipolar'],
            'RAADS': ['autism'],
            'DUDIT': ['substance abuse'],
            'AUDIT': ['substance abuse'],
            'BOCS': ['OCD'],
            'ASRS': ['ADHD'],
            'NSE': ['PTSD'],
            'EDE_QS': ['eating']
        }
        
        # Create question sets for each outcome based on strategy
        outcome_questions = {}
        
        for outcome, specific_symptoms in outcome_to_specific_symptoms.items():
            # Get specific questions for this outcome
            specific_questions = questions_df[questions_df['symptom'].isin(specific_symptoms)]['code'].tolist()
            
            # Apply strategy with special handling for outcomes without specific questions
            if self.config.question_strategy == 'general_only':
                outcome_questions[outcome] = general_questions
            elif self.config.question_strategy == 'specific_only':
                # For outcomes without specific questions (like DUDIT, AUDIT), fall back to general
                if len(specific_questions) == 0:
                    print(f"WARNING: {outcome} has no specific questions, using general questions instead")
                    outcome_questions[outcome] = general_questions
                else:
                    outcome_questions[outcome] = specific_questions
            elif self.config.question_strategy == 'all_questions':
                outcome_questions[outcome] = all_questions
            else:  # combined (default)
                # For outcomes without specific questions, combined = general only
                if len(specific_questions) == 0:
                    outcome_questions[outcome] = general_questions
                else:
                    outcome_questions[outcome] = general_questions + specific_questions
        
        # For multitask: use same strategy
        if self.config.question_strategy == 'general_only':
            outcome_questions['MULTITASK'] = general_questions
        elif self.config.question_strategy == 'specific_only':
            # For specific_only in multitask, use all specific questions from all outcomes
            all_specific_questions = []
            for symptoms in outcome_to_specific_symptoms.values():
                specific_qs = questions_df[questions_df['symptom'].isin(symptoms)]['code'].tolist()
                all_specific_questions.extend(specific_qs)
            outcome_questions['MULTITASK'] = list(set(all_specific_questions))  # Remove duplicates
        elif self.config.question_strategy == 'all_questions':
            outcome_questions['MULTITASK'] = all_questions
        else:  # combined
            outcome_questions['MULTITASK'] = all_questions
        
        print(f"Question Strategy: {self.config.question_strategy}")
        print(f"Total questions available: {len(all_questions)}")
        print(f"General questions: {len(general_questions)}")
        
        # Show detailed breakdown for each outcome
        for outcome, questions in outcome_questions.items():
            if outcome != 'MULTITASK':
                specific_questions = questions_df[questions_df['symptom'].isin(outcome_to_specific_symptoms[outcome])]['code'].tolist()
                if len(specific_questions) == 0:
                    print(f"{outcome} questions: {len(questions)} (no specific questions available)")
                else:
                    print(f"{outcome} questions: {len(questions)} (has {len(specific_questions)} specific questions)")
        print(f"MULTITASK questions: {len(outcome_questions['MULTITASK'])}")
        
        return outcome_questions
    
    def prepare_data(self, input_data, output_data, question_data, questions_mapping, target_outcome=None):
        # Filter questions based on strategy
        if target_outcome and target_outcome in questions_mapping:
            # For singletask models, use questions specific to the target outcome
            relevant_questions = questions_mapping[target_outcome]
            input_data = input_data[input_data['question_code'].isin(relevant_questions)].copy()
            print(f"Filtered to {len(input_data)} records for {target_outcome} using {self.config.question_strategy} strategy")
        elif self.config.question_strategy != 'all_questions':
            # For multitask models, use MULTITASK questions
            relevant_questions = questions_mapping['MULTITASK']
            input_data = input_data[input_data['question_code'].isin(relevant_questions)].copy()
            print(f"Filtered to {len(input_data)} records using {self.config.question_strategy} strategy")
        
        # Use all input columns
        input_cols = [col for col in input_data.columns if col.startswith("input_")]
        
        # Process question embeddings based on embedding_type
        qu_embed_data = []
        if self.config.embedding_type == "full_embedding":
            for _, row in tqdm(input_data.iterrows(), total=len(input_data), desc="Processing question embeddings"):
                embed_str = question_data.loc[row["question_code"], "question_embeddings"]
                embed_values = [float(x) for x in embed_str.strip("[] ").split()][:64]
                qu_embed_data.append(embed_values)
        
        # Create final dataset
        processed_data = pd.DataFrame({
            'user_id': input_data['user_id'],
            'question_code': input_data['question_code']
        })
        
        # Add input features
        for col in input_cols:
            processed_data[col] = input_data[col].values
        
        # Add question embeddings if using full embedding
        if self.config.embedding_type == "full_embedding":
            qu_embed_df = pd.DataFrame(qu_embed_data, columns=[f"input_qu_embed_{i}" for i in range(64)])
            processed_data = pd.concat([processed_data, qu_embed_df], axis=1)
        
        # Process output data based on num_outputs
        output_cols = [f'output_{i}' for i in range(self.config.num_outputs)]
        combined_output_data = output_data[output_cols].copy()
        
        # Merge data
        merged_data = pd.merge(processed_data, combined_output_data, left_on="user_id", right_index=True)
        merged_data.reset_index(drop=True, inplace=True)
        
        return merged_data
    
    def create_cv_folds(self, data):
        # Aggregate by user, excluding user_id and question_code columns
        agg_cols = {
            col: (lambda x, c=col: np.mean(x) if c.startswith(('input_', 'output_')) else x.iloc[0])
            for col in data.columns if col not in ['user_id', 'question_code']
        }
        
        user_data = data.groupby("user_id").agg(agg_cols).reset_index()
        user_data["PHQ"] = user_data["output_0"]
        user_data_sorted = user_data.sort_values(by="PHQ", ascending=False)
        
        # Create folds
        user_ids = user_data_sorted['user_id'].tolist()
        bins = [[] for _ in range(self.config.n_folds)]
        
        for i, user_id in enumerate(user_ids):
            bins[i % self.config.n_folds].append(user_id)
        
        return bins

class HyperparameterTuner:
    def __init__(self, config):
        self.config = config
    
    def get_scaler(self, scaler_type='minmax'):
        """Always return MinMax scaler"""
        return MinMaxScaler()
    
    def train_single_model(self, train_data, val_data, params, target_outcome_idx=None):
        """Train a single model with given hyperparameters"""
        # Prepare features and labels
        input_cols = [col for col in train_data.columns if col.startswith("input_")]
        output_cols = [col for col in train_data.columns if col.startswith("output_")]
        
        # Aggregate by user
        train_agg = train_data.groupby("user_id").agg({
            **{col: 'mean' for col in input_cols},
            **{col: 'first' for col in output_cols}
        }).reset_index()
        
        val_agg = val_data.groupby("user_id").agg({
            **{col: 'mean' for col in input_cols},
            **{col: 'first' for col in output_cols}
        }).reset_index()
        
        # Convert to numpy arrays
        X_train = train_agg[input_cols].values
        y_train = train_agg[output_cols].values
        X_val = val_agg[input_cols].values
        y_val = val_agg[output_cols].values
        
        # Handle NaN/Inf
        X_train = np.nan_to_num(X_train, nan=0.0, posinf=1.0, neginf=-1.0)
        X_val = np.nan_to_num(X_val, nan=0.0, posinf=1.0, neginf=-1.0)
        y_train = np.nan_to_num(y_train, nan=0.0, posinf=1.0, neginf=-1.0)
        y_val = np.nan_to_num(y_val, nan=0.0, posinf=1.0, neginf=-1.0)
        
        # Normalize (always use MinMax)
        input_scaler = self.get_scaler()
        output_scaler = self.get_scaler()
        
        X_train_norm = input_scaler.fit_transform(X_train)
        X_val_norm = input_scaler.transform(X_val)
        
        # For singletask, select specific outcome
        if target_outcome_idx is not None:
            y_train = y_train[:, target_outcome_idx:target_outcome_idx+1]
            y_val = y_val[:, target_outcome_idx:target_outcome_idx+1]
        
        y_train_norm = output_scaler.fit_transform(y_train)
        
        # Create model
        train_inputs = torch.tensor(X_train_norm, dtype=torch.float32).to(device)
        train_true = torch.tensor(y_train_norm, dtype=torch.float32).to(device)
        val_inputs = torch.tensor(X_val_norm, dtype=torch.float32).to(device)
        
        # Choose model based on task type
        if target_outcome_idx is not None:
            # Singletask model
            model = SingletaskLinearModel(train_inputs.shape[1]).to(device)
        else:
            # Multitask model
            model = MultitaskLinearModel(train_inputs.shape[1], train_true.shape[1]).to(device)
        
        model.apply(init_weights)
        
        # Create optimizer (always use Adam)
        optimizer = torch.optim.Adam(
            model.parameters(), 
            lr=params['lr'], 
            weight_decay=params['weight_decay']
        )
        
        criterion = nn.MSELoss()
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, patience=10, factor=0.5, min_lr=1e-6
        )
        
        # Training loop
        best_val_loss = float('inf')
        patience_counter = 0
        
        for epoch in range(self.config.max_epochs):
            # Training
            model.train()
            optimizer.zero_grad()
            train_pred = model(train_inputs)
            train_loss = criterion(train_pred, train_true)
            
            if torch.isnan(train_loss) or torch.isinf(train_loss):
                break
                
            train_loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            
            # Validation check every 10 epochs.
            # Note: patience_counter increments only at these checkpoints,
            # so effective patience = self.config.patience × 10 epochs.
            if epoch % 10 == 0:
                model.eval()
                with torch.no_grad():
                    val_pred = model(val_inputs)
                    val_loss = criterion(val_pred, torch.tensor(
                        output_scaler.transform(y_val), dtype=torch.float32
                    ).to(device))
                
                scheduler.step(val_loss)
                
                if val_loss.item() < best_val_loss:
                    best_val_loss = val_loss.item()
                    patience_counter = 0
                else:
                    patience_counter += 1
                
                if patience_counter >= self.config.patience:
                    break
        
        return model, input_scaler, output_scaler, best_val_loss
    
    def tune_hyperparameters(self, train_folds_data, dev_folds_data, target_outcome_idx=None):
        """Perform hyperparameter tuning using cross-validation within train folds"""
        print("Starting hyperparameter tuning...")
        
        param_combinations = list(ParameterGrid(self.config.param_grid))
        best_params = None
        best_score = float('inf')
        error_messages = []
        
        results = []
        
        # Use the number of train folds for cross-validation
        num_train_folds = len(train_folds_data)
        
        for i, params in enumerate(tqdm(param_combinations, desc="Hyperparameter tuning")):
            print(f"\nTesting parameters {i+1}/{len(param_combinations)}: {params}")
            
            fold_scores = []
            
            # Cross-validation within training data
            for fold_idx in range(num_train_folds):
                # Use (num_train_folds-1) folds for training, 1 fold for validation
                train_fold_data = []
                val_fold_data = train_folds_data[fold_idx]
                
                for j in range(num_train_folds):
                    if j != fold_idx:
                        train_fold_data.append(train_folds_data[j])
                
                # Combine training folds
                train_combined = pd.concat(train_fold_data, ignore_index=True)
                
                try:
                    model, input_scaler, output_scaler, val_loss = self.train_single_model(
                        train_combined, val_fold_data, params, target_outcome_idx
                    )
                    fold_scores.append(val_loss)
                except Exception as e:
                    err = f"param_idx={i}, fold={fold_idx}, error={e}"
                    error_messages.append(err)
                    print(f"Error in fold {fold_idx}: {e}")
                    fold_scores.append(float('inf'))
            
            # Calculate average validation score
            avg_score = np.mean(fold_scores)
            std_score = np.std(fold_scores)
            
            results.append({
                'params': params,
                'mean_val_loss': avg_score,
                'std_val_loss': std_score,
                'fold_scores': fold_scores
            })
            
            print(f"Average validation loss: {avg_score:.6f} ± {std_score:.6f}")
            
            if avg_score < best_score:
                best_score = avg_score
                best_params = params
                print(f"New best parameters found! Score: {best_score:.6f}")
        
        print(f"\nBest hyperparameters: {best_params}")
        print(f"Best validation score: {best_score:.6f}")

        if best_params is None:
            details = "\n".join(error_messages[:5])
            if len(error_messages) > 5:
                details += f"\n... and {len(error_messages) - 5} more errors"
            raise RuntimeError(
                "Hyperparameter tuning failed: no valid parameter set produced a finite "
                "validation loss. This often indicates an environment/runtime failure."
                + (f"\nSample errors:\n{details}" if details else "")
            )
        
        return best_params, results

class Trainer:
    def __init__(self, config):
        self.config = config
        self.tuner = HyperparameterTuner(config)
    
    def train_model_with_hyperparams(self, train_folds_data, dev_folds_data, test_fold_data, target_outcome_idx=None):
        """Train model with hyperparameter tuning and evaluate at user level"""
        
        # Step 1: Hyperparameter tuning using train folds
        if self.config.hyperparam_tuning:
            best_params, tuning_results = self.tuner.tune_hyperparameters(
                train_folds_data, dev_folds_data, target_outcome_idx
            )
        else:
            # Use default parameters (Adam optimizer only)
            best_params = {
                'lr': 0.01,
                'weight_decay': 0.01
            }
            tuning_results = []
        
        # Step 2: Train final model on all train folds with best hyperparameters
        print("\nTraining final model with best hyperparameters...")
        train_combined = pd.concat(train_folds_data, ignore_index=True)
        dev_combined = pd.concat(dev_folds_data, ignore_index=True)
        
        final_model, input_scaler, output_scaler, _ = self.tuner.train_single_model(
            train_combined, dev_combined, best_params, target_outcome_idx
        )
        
        # Step 3: Evaluate at user level
        print("Evaluating at user level...")
        user_correlations, user_mse = self.evaluate_user_level(
            final_model, input_scaler, output_scaler, test_fold_data, target_outcome_idx
        )
        
        return final_model, input_scaler, output_scaler, best_params, user_correlations, user_mse, tuning_results
    
    def evaluate_user_level(self, model, input_scaler, output_scaler, test_data, target_outcome_idx=None):
        """Evaluate model at user level by aggregating predictions per user"""
        # Get QA-level predictions
        qa_predictions = self.predict_qa_pairs(model, input_scaler, output_scaler, test_data, target_outcome_idx)
        
        # Aggregate to user level
        if target_outcome_idx is not None:
            # For singletask model
            outcome_name = self.config.outcome_names[target_outcome_idx]
            user_agg = qa_predictions.groupby('user_id').agg({
                f'{outcome_name}_pred': 'mean',
                f'output_{target_outcome_idx}': 'first'
            }).reset_index()
            
            # Calculate user-level metrics for single outcome
            correlations = {}
            mse_results = {}
            
            try:
                true_col = f'output_{target_outcome_idx}'
                pred_col = f'{outcome_name}_pred'
                
                if true_col in user_agg.columns and pred_col in user_agg.columns:
                    true_vals = user_agg[true_col].values
                    pred_vals = user_agg[pred_col].values
                    
                    # Handle NaN/Inf
                    mask = np.isfinite(true_vals) & np.isfinite(pred_vals)
                    if np.sum(mask) > 1:
                        true_vals_clean = true_vals[mask]
                        pred_vals_clean = pred_vals[mask]
                        
                        if np.var(true_vals_clean) > 1e-10 and np.var(pred_vals_clean) > 1e-10:
                            corr, _ = pearsonr(true_vals_clean, pred_vals_clean)
                            mse = np.mean((true_vals_clean - pred_vals_clean) ** 2)
                            
                            correlations[outcome_name] = corr if np.isfinite(corr) else 0.0
                            mse_results[outcome_name] = mse if np.isfinite(mse) else float('inf')
                        else:
                            correlations[outcome_name] = 0.0
                            mse_results[outcome_name] = float('inf')
                    else:
                        correlations[outcome_name] = 0.0
                        mse_results[outcome_name] = float('inf')
                else:
                    correlations[outcome_name] = 0.0
                    mse_results[outcome_name] = float('inf')
                    
            except Exception as e:
                print(f"Error calculating metrics for {outcome_name}: {e}")
                correlations[outcome_name] = 0.0
                mse_results[outcome_name] = float('inf')
        else:
            # For multitask model (original behavior)
            user_agg = qa_predictions.groupby('user_id').agg({
                **{f'{outcome}_pred': 'mean' for outcome in self.config.outcome_names[:self.config.num_outputs]},
                **{f'output_{i}': 'first' for i in range(self.config.num_outputs)}
            }).reset_index()
            
            # Calculate user-level metrics
            correlations = {}
            mse_results = {}
            
            for i, outcome in enumerate(self.config.outcome_names[:self.config.num_outputs]):
                try:
                    true_col = f'output_{i}'
                    pred_col = f'{outcome}_pred'
                    
                    if true_col in user_agg.columns and pred_col in user_agg.columns:
                        true_vals = user_agg[true_col].values
                        pred_vals = user_agg[pred_col].values
                        
                        # Handle NaN/Inf
                        mask = np.isfinite(true_vals) & np.isfinite(pred_vals)
                        if np.sum(mask) > 1:
                            true_vals_clean = true_vals[mask]
                            pred_vals_clean = pred_vals[mask]
                            
                            if np.var(true_vals_clean) > 1e-10 and np.var(pred_vals_clean) > 1e-10:
                                corr, _ = pearsonr(true_vals_clean, pred_vals_clean)
                                mse = np.mean((true_vals_clean - pred_vals_clean) ** 2)
                                
                                correlations[outcome] = corr if np.isfinite(corr) else 0.0
                                mse_results[outcome] = mse if np.isfinite(mse) else float('inf')
                            else:
                                correlations[outcome] = 0.0
                                mse_results[outcome] = float('inf')
                        else:
                            correlations[outcome] = 0.0
                            mse_results[outcome] = float('inf')
                    else:
                        correlations[outcome] = 0.0
                        mse_results[outcome] = float('inf')
                        
                except Exception as e:
                    print(f"Error calculating metrics for {outcome}: {e}")
                    correlations[outcome] = 0.0
                    mse_results[outcome] = float('inf')
        
        return correlations, mse_results
    
    def predict_qa_pairs(self, model, input_scaler, output_scaler, data, target_outcome_idx=None):
        """Generate QA-level predictions"""
        input_cols = [col for col in data.columns if col.startswith("input_")]
        output_cols = [col for col in data.columns if col.startswith("output_")]
        
        X = data[input_cols].values
        X = np.nan_to_num(X, nan=0.0, posinf=1.0, neginf=-1.0)
        
        X_norm = input_scaler.transform(X)
        X_tensor = torch.tensor(X_norm, dtype=torch.float32).to(device)
        
        model.eval()
        with torch.no_grad():
            predictions_norm = model(X_tensor).cpu().numpy()
        
        # Handle different model types
        if target_outcome_idx is not None:
            # Singletask model - predictions_norm is shape (n_samples, 1)
            predictions = output_scaler.inverse_transform(predictions_norm)
            
            # Create QA predictions dataframe for single outcome - only include target output
            target_output_col = f'output_{target_outcome_idx}'
            qa_predictions = data[['user_id', 'question_code', target_output_col]].copy()
            outcome_name = self.config.outcome_names[target_outcome_idx]
            qa_predictions[f'{outcome_name}_pred'] = predictions[:, 0]
        else:
            # Multitask model - predictions_norm is shape (n_samples, n_outcomes)
            predictions = output_scaler.inverse_transform(predictions_norm)
            
            # Create QA predictions dataframe
            qa_predictions = data[['user_id', 'question_code'] + output_cols].copy()
            
            for i, outcome in enumerate(self.config.outcome_names[:self.config.num_outputs]):
                qa_predictions[f'{outcome}_pred'] = predictions[:, i]
        
        return qa_predictions

class Analysis:
    def __init__(self, question_strategy='all_questions', model_type='both'):
        self.config = Config()
        self.config.question_strategy = question_strategy
        self.config.model_type = model_type
        self.processor = DataProcessor(self.config)
        self.trainer = Trainer(self.config)
    
    def get_fold_split(self, n_folds):
        """
        Calculate fold distribution: train folds + dev folds + 1 test fold
        Ensure at least 1 fold for test, and train/dev are roughly equal
        """
        test_folds = 1
        remaining_folds = n_folds - test_folds
        
        # Split remaining folds roughly equally between train and dev
        train_folds = remaining_folds // 2
        dev_folds = remaining_folds - train_folds
        
        return train_folds, dev_folds, test_folds
    
    def run_multitask(self, input_data, output_data, question_data, questions_mapping):
        """Run multitask model training"""
        print("="*50)
        print("RUNNING MULTITASK MODEL")
        print("="*50)
        
        merged_data = self.processor.prepare_data(input_data, output_data, question_data, questions_mapping)
        cv_folds = self.processor.create_cv_folds(merged_data)
        
        # Calculate fold distribution
        train_folds_count, dev_folds_count, test_folds_count = self.get_fold_split(self.config.n_folds)
        
        # Store all results
        all_results = []
        all_correlations = []
        all_mse = []
        all_best_params = []
        all_tuning_results = []
        
        print(f"\nStarting {self.config.n_folds}-fold cross validation...")
        print(f"Model configuration: {self.config.get_config_name()}")
        print(f"Fold structure: {train_folds_count} train + {dev_folds_count} dev + {test_folds_count} test")
        
        # Run cross validation with proper fold structure
        for test_fold_idx in range(self.config.n_folds):
            print(f"\nFold {test_fold_idx + 1}/{self.config.n_folds}")
            print("=" * 40)
            
            # Test fold (1 fold)
            test_users = cv_folds[test_fold_idx]
            test_data = merged_data[merged_data['user_id'].isin(test_users)]
            
            # Split remaining folds into train + dev
            remaining_folds = [i for i in range(self.config.n_folds) if i != test_fold_idx]
            
            # First train_folds_count folds for training
            train_fold_indices = remaining_folds[:train_folds_count]
            train_folds_data = []
            train_users = []
            for idx in train_fold_indices:
                fold_users = cv_folds[idx]
                train_users.extend(fold_users)
                fold_data = merged_data[merged_data['user_id'].isin(fold_users)]
                train_folds_data.append(fold_data)
            
            # Remaining folds for dev
            dev_fold_indices = remaining_folds[train_folds_count:]
            dev_folds_data = []
            dev_users = []
            for idx in dev_fold_indices:
                fold_users = cv_folds[idx]
                dev_users.extend(fold_users)
                fold_data = merged_data[merged_data['user_id'].isin(fold_users)]
                dev_folds_data.append(fold_data)
            
            print(f"Train folds: {train_fold_indices} ({len(train_users)} users)")
            print(f"Dev folds: {dev_fold_indices} ({len(dev_users)} users)")
            print(f"Test fold: {test_fold_idx} ({len(test_users)} users)")
            
            # Train multitask model
            (model, input_scaler, output_scaler, best_params, 
             user_correlations, user_mse, tuning_results) = self.trainer.train_model_with_hyperparams(
                train_folds_data, dev_folds_data, test_data
            )
            
            all_correlations.append(user_correlations)
            all_mse.append(user_mse)
            all_best_params.append(best_params)
            all_tuning_results.append(tuning_results)
            
            # Generate predictions for all datasets
            # Train data
            train_combined = pd.concat(train_folds_data, ignore_index=True)
            train_qa_predictions = self.trainer.predict_qa_pairs(
                model, input_scaler, output_scaler, train_combined
            )
            train_qa_predictions['fold'] = test_fold_idx
            train_qa_predictions['dataset'] = 'train'
            all_results.append(train_qa_predictions)
            
            # Dev data
            dev_combined = pd.concat(dev_folds_data, ignore_index=True)
            dev_qa_predictions = self.trainer.predict_qa_pairs(
                model, input_scaler, output_scaler, dev_combined
            )
            dev_qa_predictions['fold'] = test_fold_idx
            dev_qa_predictions['dataset'] = 'dev'
            all_results.append(dev_qa_predictions)
            
            # Test data
            test_qa_predictions = self.trainer.predict_qa_pairs(
                model, input_scaler, output_scaler, test_data
            )
            test_qa_predictions['fold'] = test_fold_idx
            test_qa_predictions['dataset'] = 'test'
            all_results.append(test_qa_predictions)
            
            # Print fold results
            print(f"\nFold {test_fold_idx + 1} Results:")
            print(f"Best hyperparameters: {best_params}")
            print("User-level correlations:")
            for outcome, corr in user_correlations.items():
                print(f"  {outcome}: {corr:.4f}")
        
        return all_results, all_correlations, all_mse, all_best_params, all_tuning_results
    
    def run_singletask(self, input_data, output_data, question_data, questions_mapping):
        """Run singletask models for all outcomes"""
        print("="*50)
        print("RUNNING SINGLETASK MODELS")
        print("="*50)
        
        all_results = {}
        all_correlations = {}
        all_mse = {}
        all_best_params = {}
        all_tuning_results = {}
        
        # Train a separate model for each selected outcome
        selected_outcomes = self.config.singletask_outcomes
        outcome_to_idx = {name: i for i, name in enumerate(self.config.outcome_names)}
        for outcome_name in selected_outcomes:
            outcome_idx = outcome_to_idx[outcome_name]
            print(f"\n{'='*60}")
            print(f"TRAINING SINGLETASK MODEL FOR {outcome_name} (Index {outcome_idx})")
            print(f"{'='*60}")
            
            # Prepare data specific to this outcome
            merged_data = self.processor.prepare_data(
                input_data, output_data, question_data, questions_mapping, outcome_name
            )
            cv_folds = self.processor.create_cv_folds(merged_data)
            
            # Calculate fold distribution
            train_folds_count, dev_folds_count, test_folds_count = self.get_fold_split(self.config.n_folds)
            
            # Store results for this outcome
            outcome_results = []
            outcome_correlations = []
            outcome_mse = []
            outcome_best_params = []
            outcome_tuning_results = []
            
            print(f"Starting {self.config.n_folds}-fold cross validation for {outcome_name}...")
            print(f"Fold structure: {train_folds_count} train + {dev_folds_count} dev + {test_folds_count} test")
            
            # Run cross validation
            for test_fold_idx in range(self.config.n_folds):
                print(f"\nFold {test_fold_idx + 1}/{self.config.n_folds} - {outcome_name}")
                print("-" * 40)
                
                # Test fold (1 fold)
                test_users = cv_folds[test_fold_idx]
                test_data = merged_data[merged_data['user_id'].isin(test_users)]
                
                # Split remaining folds into train + dev
                remaining_folds = [i for i in range(self.config.n_folds) if i != test_fold_idx]
                
                # First train_folds_count folds for training
                train_fold_indices = remaining_folds[:train_folds_count]
                train_folds_data = []
                train_users = []
                for idx in train_fold_indices:
                    fold_users = cv_folds[idx]
                    train_users.extend(fold_users)
                    fold_data = merged_data[merged_data['user_id'].isin(fold_users)]
                    train_folds_data.append(fold_data)
                
                # Remaining folds for dev
                dev_fold_indices = remaining_folds[train_folds_count:]
                dev_folds_data = []
                dev_users = []
                for idx in dev_fold_indices:
                    fold_users = cv_folds[idx]
                    dev_users.extend(fold_users)
                    fold_data = merged_data[merged_data['user_id'].isin(fold_users)]
                    dev_folds_data.append(fold_data)
                
                print(f"Train folds: {train_fold_indices} ({len(train_users)} users)")
                print(f"Dev folds: {dev_fold_indices} ({len(dev_users)} users)")
                print(f"Test fold: {test_fold_idx} ({len(test_users)} users)")
                
                # Train singletask model for this outcome
                (model, input_scaler, output_scaler, best_params, 
                 user_correlations, user_mse, tuning_results) = self.trainer.train_model_with_hyperparams(
                    train_folds_data, dev_folds_data, test_data, outcome_idx
                )
                
                outcome_correlations.append(user_correlations)
                outcome_mse.append(user_mse)
                outcome_best_params.append(best_params)
                outcome_tuning_results.append(tuning_results)
                
                # Generate predictions for all datasets
                # Train data
                train_combined = pd.concat(train_folds_data, ignore_index=True)
                train_qa_predictions = self.trainer.predict_qa_pairs(
                    model, input_scaler, output_scaler, train_combined, outcome_idx
                )
                train_qa_predictions['fold'] = test_fold_idx
                train_qa_predictions['dataset'] = 'train'
                outcome_results.append(train_qa_predictions)
                
                # Dev data
                dev_combined = pd.concat(dev_folds_data, ignore_index=True)
                dev_qa_predictions = self.trainer.predict_qa_pairs(
                    model, input_scaler, output_scaler, dev_combined, outcome_idx
                )
                dev_qa_predictions['fold'] = test_fold_idx
                dev_qa_predictions['dataset'] = 'dev'
                outcome_results.append(dev_qa_predictions)
                
                # Test data
                test_qa_predictions = self.trainer.predict_qa_pairs(
                    model, input_scaler, output_scaler, test_data, outcome_idx
                )
                test_qa_predictions['fold'] = test_fold_idx
                test_qa_predictions['dataset'] = 'test'
                outcome_results.append(test_qa_predictions)
                
                # Print fold results
                print(f"\nFold {test_fold_idx + 1} Results for {outcome_name}:")
                print(f"Best hyperparameters: {best_params}")
                if outcome_name in user_correlations:
                    print(f"User-level correlation: {user_correlations[outcome_name]:.4f}")
            
            # Store results for this outcome
            all_results[outcome_name] = outcome_results
            all_correlations[outcome_name] = outcome_correlations
            all_mse[outcome_name] = outcome_mse
            all_best_params[outcome_name] = outcome_best_params
            all_tuning_results[outcome_name] = outcome_tuning_results
        
        return all_results, all_correlations, all_mse, all_best_params, all_tuning_results
    
    def run(self):
        print("Loading data...")
        torch.manual_seed(self.config.random_seed)
        np.random.seed(self.config.random_seed)
        random.seed(self.config.random_seed)
        
        input_data, output_data, question_data, questions_mapping = self.processor.load_data()
        
        all_results = []
        original_model_type = self.config.model_type
        
        if original_model_type in ['multitask', 'both']:
            print("\n" + "="*80)
            print("MULTITASK MODEL TRAINING")
            print("="*80)
            
            multitask_results, multitask_correlations, multitask_mse, multitask_best_params, multitask_tuning_results = self.run_multitask(
                input_data, output_data, question_data, questions_mapping
            )
            
            # Analyze and save multitask results
            self.config.model_type = 'multitask'  # Temporarily set for saving
            self.analyze_performance(multitask_correlations, multitask_mse, multitask_best_params)
            self.save_results(multitask_results, multitask_best_params, multitask_tuning_results, 
                             *self.get_fold_split(self.config.n_folds))
            
            all_results.extend(multitask_results)
        
        if original_model_type in ['singletask', 'both']:
            print("\n" + "="*80)
            print("SINGLETASK MODELS TRAINING")
            print("="*80)
            
            singletask_results, singletask_correlations, singletask_mse, singletask_best_params, singletask_tuning_results = self.run_singletask(
                input_data, output_data, question_data, questions_mapping
            )
            
            # Analyze and save singletask results for each outcome
            for outcome_name in self.config.singletask_outcomes:
                self.config.model_type = 'singletask'  # Temporarily set for saving
                self.analyze_performance_singletask(
                    singletask_correlations[outcome_name], 
                    singletask_mse[outcome_name], 
                    singletask_best_params[outcome_name],
                    outcome_name
                )
                self.save_results_singletask(
                    singletask_results[outcome_name], 
                    singletask_best_params[outcome_name], 
                    singletask_tuning_results[outcome_name],
                    outcome_name,
                    *self.get_fold_split(self.config.n_folds)
                )
        
        self.config.model_type = original_model_type  # restore original value
        return all_results
    
    def analyze_performance_singletask(self, fold_correlations, fold_mse, fold_params, outcome_name):
        """Analyze performance for singletask model"""
        print(f"\nPerformance Analysis for {outcome_name} (Singletask)")
        print("-" * 50)
        
        # Extract correlations and MSE values
        corr_values = [fold_corr[outcome_name] for fold_corr in fold_correlations if outcome_name in fold_corr]
        mse_values = [fm[outcome_name] for fm in fold_mse if outcome_name in fm]
        
        if corr_values:
            correlation_stats = {
                'mean': np.mean(corr_values),
                'std': np.std(corr_values),
                'values': corr_values
            }
            
            mse_stats = {
                'mean': np.mean(mse_values),
                'std': np.std(mse_values),
                'values': mse_values
            }
            
            print(f"Correlation: {correlation_stats['mean']:.4f} ± {correlation_stats['std']:.4f}")
            print(f"MSE: {mse_stats['mean']:.4f} ± {mse_stats['std']:.4f}")
            
            # Analyze hyperparameter choices
            hyperparameter_analysis = {}
            for param_key in fold_params[0].keys():
                param_values = [params[param_key] for params in fold_params]
                unique_values, counts = np.unique(param_values, return_counts=True)
                
                hyperparameter_analysis[param_key] = {}
                print(f"Hyperparameter {param_key}:")
                for val, count in zip(unique_values, counts):
                    hyperparameter_analysis[param_key][str(val)] = int(count)
                    print(f"  {val}: {count} times ({count/len(fold_params)*100:.1f}%)")
            
            # Store performance analysis
            self.performance_analysis = {
                'outcome': outcome_name,
                'correlation': correlation_stats,
                'mse': mse_stats,
                'hyperparameters': hyperparameter_analysis,
                'summary': {
                    'correlation_mean': float(correlation_stats['mean']),
                    'mse_mean': float(mse_stats['mean']),
                    'total_folds': len(fold_params)
                }
            }
        else:
            print(f"No valid results found for {outcome_name}")
            self.performance_analysis = {}
    
    def save_results_singletask(self, all_results, all_best_params, all_tuning_results, outcome_name, 
                               train_folds_count, dev_folds_count, test_folds_count):
        """Save results for singletask model"""
        # Create output directories based on model configuration
        config_name = f"{self.config.embedding_type}_singletask_{self.config.aggregation_type}_{self.config.question_strategy}_{outcome_name}"
        output_dir = ensure_dir(MODEL_OUTPUTS / config_name)
        
        # Merge all results
        all_data = pd.concat(all_results, ignore_index=True)
        
        print(f"\nSaving singletask results for {outcome_name} to: {output_dir}")
        
        saved_files = []
        
        # Create fold-specific subdirectories and save data
        for fold in range(self.config.n_folds):
            fold_dir = f"{output_dir}/fold_{fold}"
            os.makedirs(fold_dir, exist_ok=True)
            
            fold_data = all_data[all_data['fold'] == fold]
            
            for dataset in ['train', 'dev', 'test']:
                dataset_data = fold_data[fold_data['dataset'] == dataset]
                if len(dataset_data) > 0:
                    filename = f"{fold_dir}/{dataset}_{fold}.csv"
                    dataset_data.to_csv(filename, index=False)
                    saved_files.append(filename)
        
        # Save combined file for reference
        all_data.to_csv(f"{output_dir}/all_predictions.csv", index=False)
        
        # Save hyperparameter tuning results
        tuning_file = f"{output_dir}/hyperparameter_tuning_results.json"
        with open(tuning_file, 'w') as f:
            json.dump(all_tuning_results, f, indent=2, default=str)
        
        # Get performance analysis from stored data
        performance_analysis = getattr(self, 'performance_analysis', {})
        
        if performance_analysis:
            # Save performance analysis as JSON
            with open(f"{output_dir}/performance_analysis.json", 'w') as f:
                json.dump(performance_analysis, f, indent=2, default=str)
            
            # Save performance summary as CSV
            performance_csv = self.create_performance_csv_singletask(performance_analysis)
            performance_csv.to_csv(f"{output_dir}/performance_summary.csv", index=False)
        else:
            print("No performance analysis available for saving")
        
        # Save configuration
        config_info = {
            'model_type': 'singletask',
            'target_outcome': outcome_name,
            'question_strategy': self.config.question_strategy,
            'embedding_type': self.config.embedding_type,
            'n_folds': self.config.n_folds,
            'fold_structure': {
                'train_folds': train_folds_count,
                'dev_folds': dev_folds_count,
                'test_folds': test_folds_count
            },
            'random_seed': self.config.random_seed,
            'hyperparam_tuning': self.config.hyperparam_tuning,
            'param_grid': self.config.param_grid
        }
        
        with open(f"{output_dir}/config.json", 'w') as f:
            json.dump(config_info, f, indent=2)
        
        print(f"\nFiles saved for {outcome_name}:")
        print(f"  Output directory: {output_dir}")
        print(f"  Configuration: {config_name}")
        print(f"  Total files created: {len(saved_files) + 3 + (2 if performance_analysis else 0)}")
        print(f"  Fold structure: {train_folds_count} train + {dev_folds_count} dev + {test_folds_count} test")
        print(f"  Evaluation level: User-level")
    
    def create_performance_csv_singletask(self, performance_analysis):
        """Create a CSV summary of performance analysis for singletask"""
        rows = []
        
        # Correlation results
        corr_stats = performance_analysis['correlation']
        rows.append({
            'metric': 'correlation',
            'outcome': performance_analysis['outcome'],
            'mean': corr_stats['mean'],
            'std': corr_stats['std'],
            'values': str(corr_stats['values'])
        })
        
        # MSE results  
        mse_stats = performance_analysis['mse']
        rows.append({
            'metric': 'mse',
            'outcome': performance_analysis['outcome'],
            'mean': mse_stats['mean'],
            'std': mse_stats['std'],
            'values': str(mse_stats['values'])
        })
        
        # Hyperparameter analysis
        for param_key, param_stats in performance_analysis['hyperparameters'].items():
            rows.append({
                'metric': f'hyperparameter_{param_key}',
                'outcome': performance_analysis['outcome'],
                'mean': None,
                'std': None,
                'values': str(param_stats)
            })
        
        return pd.DataFrame(rows)
    
    def analyze_performance(self, fold_correlations, fold_mse, fold_params):
        print(f"\nOverall Performance Analysis (User-Level)")
        print("-" * 50)
        
        # Aggregate correlation results
        all_correlations = {}
        all_mse_results = {}
        
        for outcome in self.config.outcome_names[:self.config.num_outputs]:
            corr_values = [fold_corr[outcome] for fold_corr in fold_correlations]
            mse_values = [fm[outcome] for fm in fold_mse]
            
            all_correlations[outcome] = {
                'mean': np.mean(corr_values),
                'std': np.std(corr_values),
                'values': corr_values
            }
            
            all_mse_results[outcome] = {
                'mean': np.mean(mse_values),
                'std': np.std(mse_values),
                'values': mse_values
            }
        
        # Display results
        print("Correlation results (User-Level):")
        for outcome, stats in all_correlations.items():
            print(f"  {outcome}: {stats['mean']:.4f} ± {stats['std']:.4f}")
        
        print("\nMSE results (User-Level):")
        for outcome, stats in all_mse_results.items():
            print(f"  {outcome}: {stats['mean']:.4f} ± {stats['std']:.4f}")
        
        overall_corr = np.mean([stats['mean'] for stats in all_correlations.values()])
        overall_mse = np.mean([stats['mean'] for stats in all_mse_results.values()])
        print(f"\nOverall average correlation: {overall_corr:.4f}")
        print(f"Overall average MSE: {overall_mse:.4f}")
        
        # Analyze hyperparameter choices
        hyperparameter_analysis = {}
        print(f"\nHyperparameter Analysis:")
        for param_key in fold_params[0].keys():
            param_values = [params[param_key] for params in fold_params]
            unique_values, counts = np.unique(param_values, return_counts=True)
            
            hyperparameter_analysis[param_key] = {}
            print(f"  {param_key}:")
            for val, count in zip(unique_values, counts):
                percentage = count/len(fold_params)*100
                print(f"    {val}: {count}/{len(fold_params)} folds ({percentage:.1f}%)")
                hyperparameter_analysis[param_key][str(val)] = {
                    'count': int(count),
                    'percentage': float(percentage)
                }
        
        # Store performance analysis for saving
        self.performance_analysis = {
            'correlations': all_correlations,
            'mse': all_mse_results,
            'hyperparameters': hyperparameter_analysis,
            'summary': {
                'overall_correlation_mean': float(overall_corr),
                'overall_mse_mean': float(overall_mse),
                'total_folds': len(fold_params),
                'num_outcomes': self.config.num_outputs
            }
        }
        
        return self.performance_analysis
    
    def create_performance_csv(self, performance_analysis):
        """Create a CSV summary of performance analysis"""
        rows = []
        
        # Correlation results
        for outcome, stats in performance_analysis['correlations'].items():
            rows.append({
                'Metric': 'Correlation',
                'Outcome': outcome,
                'Mean': stats['mean'],
                'Std': stats['std'],
                'Values': ', '.join([f'{v:.4f}' for v in stats['values']])
            })
        
        # MSE results
        for outcome, stats in performance_analysis['mse'].items():
            rows.append({
                'Metric': 'MSE',
                'Outcome': outcome,
                'Mean': stats['mean'],
                'Std': stats['std'],
                'Values': ', '.join([f'{v:.4f}' for v in stats['values']])
            })
        
        # Hyperparameter analysis
        for param_key, param_stats in performance_analysis['hyperparameters'].items():
            for value, info in param_stats.items():
                rows.append({
                    'Metric': 'Hyperparameter',
                    'Outcome': f'{param_key}_{value}',
                    'Mean': info['count'],
                    'Std': info['percentage'],
                    'Values': f"{info['count']}/{performance_analysis['summary']['total_folds']} folds"
                })
        
        return pd.DataFrame(rows)

    def save_results(self, all_results, all_best_params, all_tuning_results, 
                    train_folds_count, dev_folds_count, test_folds_count):
        # Create output directories based on model configuration
        config_name = self.config.get_config_name()
        output_dir = ensure_dir(MODEL_OUTPUTS / config_name)
        
        # Merge all results
        all_data = pd.concat(all_results, ignore_index=True)
        
        print(f"\nSaving results to: {output_dir}")
        
        saved_files = []
        
        # Create fold-specific subdirectories and save data
        for fold in range(self.config.n_folds):
            fold_dir = f"{output_dir}/fold_{fold}"
            os.makedirs(fold_dir, exist_ok=True)
            
            # Save train data for this fold
            train_data = all_data[(all_data['fold'] == fold) & (all_data['dataset'] == 'train')]
            if len(train_data) > 0:
                train_file = f"{fold_dir}/train_{fold}.csv"
                train_data.to_csv(train_file, index=False)
                saved_files.append(train_file)
                print(f"  Saved train fold {fold}: {train_file} ({len(train_data)} records)")
            
            # Save dev data for this fold
            dev_data = all_data[(all_data['fold'] == fold) & (all_data['dataset'] == 'dev')]
            if len(dev_data) > 0:
                dev_file = f"{fold_dir}/dev_{fold}.csv"
                dev_data.to_csv(dev_file, index=False)
                saved_files.append(dev_file)
                print(f"  Saved dev fold {fold}: {dev_file} ({len(dev_data)} records)")
            
            # Save test data for this fold
            test_data = all_data[(all_data['fold'] == fold) & (all_data['dataset'] == 'test')]
            if len(test_data) > 0:
                test_file = f"{fold_dir}/test_{fold}.csv"
                test_data.to_csv(test_file, index=False)
                saved_files.append(test_file)
                print(f"  Saved test fold {fold}: {test_file} ({len(test_data)} records)")
            
            # Save best hyperparameters for this fold
            params_file = f"{fold_dir}/best_params_{fold}.json"
            with open(params_file, 'w') as f:
                json.dump(all_best_params[fold], f, indent=2)
        
        # Save combined file for reference
        all_data.to_csv(f"{output_dir}/all_predictions.csv", index=False)
        
        # Save hyperparameter tuning results
        tuning_file = f"{output_dir}/hyperparameter_tuning_results.json"
        with open(tuning_file, 'w') as f:
            json.dump(all_tuning_results, f, indent=2, default=str)
        
        # Get performance analysis from stored data
        performance_analysis = getattr(self, 'performance_analysis', {})
        
        if performance_analysis:
            # Save performance analysis as JSON
            performance_file = f"{output_dir}/performance_analysis.json"
            with open(performance_file, 'w') as f:
                json.dump(performance_analysis, f, indent=2)
            
            # Save performance analysis as CSV for easy viewing
            performance_csv = self.create_performance_csv(performance_analysis)
            performance_csv_file = f"{output_dir}/performance_summary.csv"
            performance_csv.to_csv(performance_csv_file, index=False)
            
            print(f"  Saved performance analysis: {performance_file}")
            print(f"  Saved performance summary CSV: {performance_csv_file}")
            
            # Save configuration info with performance summary
            config_info = {
                'embedding_type': self.config.embedding_type,
                'num_outputs': self.config.num_outputs,
                'aggregation_type': self.config.aggregation_type,
                'question_strategy': self.config.question_strategy,
                'n_folds': self.config.n_folds,
                'random_seed': self.config.random_seed,
                'max_epochs': self.config.max_epochs,
                'patience': self.config.patience,
                'outcome_names': self.config.outcome_names[:self.config.num_outputs],
                'fold_structure': f'{train_folds_count}_train_{dev_folds_count}_dev_{test_folds_count}_test',
                'train_folds_count': train_folds_count,
                'dev_folds_count': dev_folds_count,
                'test_folds_count': test_folds_count,
                'hyperparameter_tuning': self.config.hyperparam_tuning,
                'param_grid': self.config.param_grid,
                'evaluation_level': 'user_level',
                'overall_performance': performance_analysis.get('summary', {})
            }
        else:
            # Save configuration info without performance summary
            config_info = {
                'embedding_type': self.config.embedding_type,
                'num_outputs': self.config.num_outputs,
                'aggregation_type': self.config.aggregation_type,
                'question_strategy': self.config.question_strategy,
                'n_folds': self.config.n_folds,
                'random_seed': self.config.random_seed,
                'max_epochs': self.config.max_epochs,
                'patience': self.config.patience,
                'outcome_names': self.config.outcome_names[:self.config.num_outputs],
                'fold_structure': f'{train_folds_count}_train_{dev_folds_count}_dev_{test_folds_count}_test',
                'train_folds_count': train_folds_count,
                'dev_folds_count': dev_folds_count,
                'test_folds_count': test_folds_count,
                'hyperparameter_tuning': self.config.hyperparam_tuning,
                'param_grid': self.config.param_grid,
                'evaluation_level': 'user_level'
            }
        
        with open(f"{output_dir}/config.json", 'w') as f:
            json.dump(config_info, f, indent=2)
        
        print(f"\nFiles saved:")
        print(f"  Output directory: {output_dir}")
        print(f"  Configuration: {config_name}")
        print(f"  Total files created: {len(saved_files) + 3 + (2 if performance_analysis else 0)}")
        print(f"  Fold structure: {train_folds_count} train + {dev_folds_count} dev + {test_folds_count} test")
        print(f"  Evaluation level: User-level")
    
def run_all_question_strategies(model_type='both', strategies=None, singletask_outcomes=None):
    """Run analysis for selected question strategies and model types"""
    if strategies is None:
        strategies = ['all_questions', 'specific_only', 'general_only', 'combined']
    
    all_strategy_results = {}
    failed_strategies = []
    
    for strategy in strategies:
        print(f"\n{'='*80}")
        print(f"RUNNING ANALYSIS WITH QUESTION STRATEGY: {strategy.upper()}")
        print(f"MODEL TYPE: {model_type.upper()}")
        print(f"{'='*80}")
        
        try:
            analysis = Analysis(question_strategy=strategy, model_type=model_type)
            if singletask_outcomes is not None:
                analysis.config.singletask_outcomes = list(singletask_outcomes)
            results = analysis.run()
            all_strategy_results[strategy] = results
            print(f"✓ Completed {strategy} strategy with {model_type} model(s)")
        except Exception as e:
            print(f"✗ Error with {strategy} strategy: {e}")
            import traceback
            traceback.print_exc()
            failed_strategies.append((strategy, str(e)))
    
    return all_strategy_results, failed_strategies

def main():
    parser = argparse.ArgumentParser(
        description="Train user-level QA models across question strategies."
    )
    parser.add_argument(
        "--model-type",
        choices=["multitask", "singletask", "both"],
        default="multitask",
        help="Model family to train. Default: multitask",
    )
    parser.add_argument(
        "--question-strategy",
        choices=["all", "all_questions", "specific_only", "general_only", "combined"],
        default="all",
        help="Question strategy to run. Use 'all' to run all strategies. Default: all",
    )
    parser.add_argument(
        "--singletask-outcomes",
        default="all",
        help="Comma-separated singletask outcomes to run (e.g. PHQ,GAD) or 'all'.",
    )
    args = parser.parse_args()

    print("="*60)
    print("Enhanced Question-Answer Pair Model Training")
    print("With Hyperparameter Tuning and User-Level Evaluation")
    print("Support for Multitask and Singletask Models")
    print("="*60)
    print(f"Using device: {device}")
    print("Fold structure: 4 train + 4 dev + 1 test")
    print("Evaluation: User-level aggregation")
    print("Model types: multitask, singletask, both")
    print("="*60)
    
    # You can change this to run different model types:
    # 'multitask' - only multitask models
    # 'singletask' - only singletask models  
    # 'both' - both multitask and singletask models
    model_type = args.model_type
    if args.question_strategy == "all":
        strategies = ['all_questions', 'specific_only', 'general_only', 'combined']
    else:
        strategies = [args.question_strategy]

    singletask_outcomes = None
    if args.singletask_outcomes != "all":
        requested = [x.strip() for x in args.singletask_outcomes.split(",") if x.strip()]
        valid = ["PHQ", "GAD", "MDQ", "RAADS", "DUDIT", "AUDIT", "BOCS", "ASRS", "NSE", "EDE_QS"]
        unknown = [x for x in requested if x not in valid]
        if unknown:
            raise ValueError(f"Unknown outcomes in --singletask-outcomes: {unknown}")
        if not requested:
            raise ValueError("--singletask-outcomes specified but no valid outcomes provided")
        singletask_outcomes = requested
    
    print(f"\nRunning question strategies: {', '.join(strategies)} with model type: {model_type}")
    if singletask_outcomes is not None:
        print(f"Singletask outcomes subset: {', '.join(singletask_outcomes)}")
    try:
        all_results, failed_strategies = run_all_question_strategies(
            model_type=model_type,
            strategies=strategies,
            singletask_outcomes=singletask_outcomes,
        )
        if failed_strategies:
            details = "; ".join([f"{name}: {err}" for name, err in failed_strategies])
            raise RuntimeError(f"One or more strategies failed: {details}")

        print(f"\nAll strategies completed! Results saved with user-level evaluation.")
        print(f"Model type: {model_type}")
        return 0
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    raise SystemExit(main())