qa_level_strategy_comparison.py

#!/usr/bin/env python3
"""
QA-Level Training Strategy Implementation - ULTIMATE FIXED VERSION
==============================================================================

This version specifically addresses the float() conversion error by:
1. Properly handling different data types in input columns
2. Ensuring consistent data flow between functions
3. Adding comprehensive debugging and error handling
"""

import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import json
import os
import sys
import random
import warnings
from sklearn.model_selection import ParameterGrid
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.linear_model import LinearRegression
from scipy.stats import pearsonr
from tqdm import tqdm
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.nn.utils.clip_grad import clip_grad_norm_
import ast
import traceback
import torch.utils.data

from maqua_data import build_qa_merged_dataset, load_question_embeddings
from maqua_mapping import get_questionnaire_mapping as _get_questionnaire_mapping

warnings.filterwarnings('ignore')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# ==================== CORE FUNCTIONS ====================

def create_stratified_cv_folds(user_data, n_folds=9):
    """Create stratified CV folds"""
    user_data_copy = user_data.copy()
    unique_users = user_data_copy["user_id"].nunique()
    if unique_users < n_folds:
        raise ValueError(
            f"Insufficient users for {n_folds}-fold CV: got {unique_users}. "
            "Small-sample debug mode is not supported."
        )
    user_data_copy["PHQ"] = user_data_copy["output_embeddings"].apply(lambda x: x[0] if isinstance(x, list) else x)
    
    user_data_sorted = user_data_copy.sort_values(by="PHQ", ascending=False)
    grouped = user_data_sorted.groupby("user_id", sort=False)
    bins = [list() for _ in range(n_folds)]
    
    for i, (user_id, group) in enumerate(grouped):
        bins[i % n_folds].append(user_id)
    
    print(f"Created {n_folds} stratified CV folds:")
    for i, bin_users in enumerate(bins):
        phq_scores = []
        for user_id in bin_users:
            user_row = user_data_copy[user_data_copy['user_id'] == user_id].iloc[0]
            phq_scores.append(user_row['PHQ'])
        
        phq_array = np.array(phq_scores)
        print(f"  Fold {i+1}: {len(bin_users)} users, PHQ range: {phq_array.min():.3f}-{phq_array.max():.3f}, mean: {phq_array.mean():.3f}")
    
    return bins

def get_questionnaire_mapping(filepath="formatted_data/questions_symptom_outcomes_type.csv"):
    """Load questionnaire mapping (delegates to maqua_mapping)."""
    return _get_questionnaire_mapping(filepath, verbose=True)

class LinearRegressionMultitaskQA(nn.Module):
    """Multitask linear regression for QA-level data"""
    def __init__(self, input_dim, output_dim):
        super(LinearRegressionMultitaskQA, self).__init__()
        self.linear = nn.Linear(input_dim, output_dim)
        
    def forward(self, x):
        return self.linear(x)

def smart_float_conversion(val, col_name="unknown"):
    """
    Smart conversion of various data types to float values
    Handles: scalars, strings, lists, arrays, etc.
    """
    try:
        # Case 1: Already a number
        if isinstance(val, (int, float, np.integer, np.floating)):
            return [float(val)]
        
        # Case 2: String representation
        elif isinstance(val, str):
            val = val.strip()
            
            # String representation of a list/array
            if val.startswith('[') and val.endswith(']'):
                try:
                    # Parse list string safely
                    parsed_list = ast.literal_eval(val)
                    if isinstance(parsed_list, (list, tuple)):
                        return [float(item) for item in parsed_list]
                    else:
                        return [float(parsed_list)]
                except (ValueError, SyntaxError, TypeError):
                    # If eval fails, try manual parsing
                    val_clean = val.strip('[]')
                    numbers = [float(x.strip()) for x in val_clean.split(',') if x.strip()]
                    return numbers
            
            # Single string number
            else:
                return [float(val)]
        
        # Case 3: Already a list/tuple/array
        elif isinstance(val, (list, tuple, np.ndarray)):
            return [float(item) for item in val]
        
        # Case 4: Other types - attempt direct conversion
        else:
            return [float(val)]
            
    except Exception as e:
        raise ValueError(
            f"Cannot convert value {val!r} (type: {type(val).__name__}) "
            f"to float in column {col_name}"
        ) from e

def prepare_qa_level_data_ultimate(data, embedding_strategy='full', separate_dudit_audit=True):
    """
    Ultimate QA-level data preparation with robust type handling
    """
    print(f"DEBUG: Starting QA-level data preparation with strategy '{embedding_strategy}'")
    print(f"DEBUG: Input data shape: {data.shape}")
    print(f"DEBUG: Input data columns sample: {list(data.columns)[:10]}")
    
    qa_level_data = []
    
    # Identify column types
    all_cols = list(data.columns)
    generic_input_cols = [col for col in all_cols if col.startswith("input_") and not col.startswith("input_qu_embed_")]
    qu_embed_cols = [col for col in all_cols if col.startswith("input_qu_embed_")]
    output_cols = sorted([col for col in all_cols if col.startswith("output_")])
    
    print(f"DEBUG: Found {len(generic_input_cols)} generic input columns")
    print(f"DEBUG: Found {len(qu_embed_cols)} question embedding columns")
    print(f"DEBUG: Found {len(output_cols)} output columns")
    
    # Select columns based on strategy
    if embedding_strategy == 'full':
        selected_input_cols = generic_input_cols + qu_embed_cols
    elif embedding_strategy == 'question_only':
        selected_input_cols = qu_embed_cols
    else:
        selected_input_cols = generic_input_cols
    
    print(f"DEBUG: Selected {len(selected_input_cols)} input columns for strategy '{embedding_strategy}'")
    
    # Handle DUDIT/AUDIT combination
    if not separate_dudit_audit:
        dudit_col, audit_col = 'output_4', 'output_5'
        if dudit_col in output_cols and audit_col in output_cols:
            data = data.copy()
            data["output_combined_dudit_audit"] = data[dudit_col] + data[audit_col]
            output_cols = [col for col in output_cols if col not in [dudit_col, audit_col]]
            output_cols.insert(4, "output_combined_dudit_audit")
            output_cols = sorted(output_cols)
    
    print(f"DEBUG: Final output columns: {output_cols}")
    
    # Process each row
    for idx, row in data.iterrows():
        user_id = row['user_id']
        question_code = row['question_code']
        
        # Process input embeddings with smart conversion
        input_embeddings = []
        for col in selected_input_cols:
            val = row[col]
            converted_vals = smart_float_conversion(val, col)
            input_embeddings.extend(converted_vals)
        
        # Process output embeddings
        output_embeddings = []
        for col in output_cols:
            val = row[col]
            converted_vals = smart_float_conversion(val, col)
            # Outputs should be single values, so take the first
            output_embeddings.append(converted_vals[0])
        
        qa_level_data.append({
            'user_id': user_id,
            'question_code': question_code,
            'input_embeddings': input_embeddings,
            'output_embeddings': output_embeddings
        })
        
        # Debug first few rows
        if idx < 3:
            print(f"DEBUG: Row {idx} - User: {user_id}, Question: {question_code}")
            print(f"DEBUG: Input length: {len(input_embeddings)}, Output length: {len(output_embeddings)}")
            print(f"DEBUG: First few inputs: {input_embeddings[:5]}")
            print(f"DEBUG: Outputs: {output_embeddings}")
    
    result_df = pd.DataFrame(qa_level_data)
    
    if len(result_df) > 0:
        print(f"DEBUG: Final QA-level data shape: {result_df.shape}")
        print(f"DEBUG: Sample input embedding length: {len(result_df['input_embeddings'].iloc[0])}")
        print(f"DEBUG: Sample output embedding length: {len(result_df['output_embeddings'].iloc[0])}")
        print(f"DEBUG: Data types - inputs: {type(result_df['input_embeddings'].iloc[0])}, outputs: {type(result_df['output_embeddings'].iloc[0])}")
    
    return result_df

def filter_qa_pairs_by_strategy(qa_data, question_strategy, all_outcome_mappings):
    """
    Filter QA pairs by strategy - CORRECTED LOGIC FOR MULTITASK
    
    For multitask models:
    - all_questions: Use ALL questions (specific + general)
    - specific_only: Use ONLY questions that belong to specific questionnaires (remove general)
    - general_only: Use ONLY general questions (questions not in any questionnaire)
    - combined: Same as all_questions for multitask
    """
    print(f"DEBUG: Filtering {len(qa_data)} QA pairs using strategy '{question_strategy}'")
    
    # Get all questions that belong to specific questionnaires
    all_specific_questions = set()
    for outcome, questions in all_outcome_mappings['questionnaire_specific'].items():
        all_specific_questions.update(questions)
    print(all_specific_questions)
    print(f"DEBUG: Found {len(all_specific_questions)} questions across all specific questionnaires")
    
    # Check what questions we have in data
    unique_questions_in_data = set(qa_data['question_code'].unique())
    print(f"DEBUG: Data contains {len(unique_questions_in_data)} unique question codes")
    
    # Find which questions in our data are specific vs general
    specific_questions_in_data = all_specific_questions.intersection(unique_questions_in_data)
    general_questions_in_data = unique_questions_in_data - all_specific_questions
    
    print(f"DEBUG: Questions in data that are questionnaire-specific: {len(specific_questions_in_data)}")
    print(f"DEBUG: Questions in data that are general: {len(general_questions_in_data)}")
    
    if len(general_questions_in_data) > 0:
        print(f"DEBUG: Sample general questions: {list(general_questions_in_data)[:5]}")
    if len(specific_questions_in_data) > 0:
        print(f"DEBUG: Sample specific questions: {list(specific_questions_in_data)[:5]}")
    
    # Apply filtering based on strategy
    if question_strategy == 'all_questions':
        result = qa_data.copy()
        print(f"DEBUG: 'all_questions' strategy - keeping all {len(result)} pairs")
        
    elif question_strategy == 'specific_only':
        # CORRECTED: Only keep questions that belong to specific questionnaires
        result = qa_data[qa_data['question_code'].isin(specific_questions_in_data)]
        print(f"DEBUG: 'specific_only' strategy - keeping only questionnaire-specific questions: {len(result)} pairs")
        
    elif question_strategy == 'general_only':
        # Only keep questions that don't belong to any questionnaire
        result = qa_data[qa_data['question_code'].isin(general_questions_in_data)]
        print(f"DEBUG: 'general_only' strategy - keeping only general questions: {len(result)} pairs")
        
    elif question_strategy == 'combined':
        # For multitask, combined is same as all_questions
        result = qa_data.copy()
        print(f"DEBUG: 'combined' strategy - keeping all questions (same as all_questions): {len(result)} pairs")
        
    else:
        raise ValueError(f"Unknown question strategy: {question_strategy}")
    
    # Verify the filtering worked correctly
    if len(result) > 0:
        result_questions = set(result['question_code'].unique())
        print(f"DEBUG: Result contains {len(result_questions)} unique question codes")
        
        # Check the composition
        result_specific = result_questions.intersection(specific_questions_in_data)
        result_general = result_questions.intersection(general_questions_in_data)
        
        print(f"DEBUG: Result composition - Specific: {len(result_specific)}, General: {len(result_general)}")
        
        # Verify strategy was applied correctly
        if question_strategy == 'specific_only' and len(result_general) > 0:
            print(f"WARNING: specific_only strategy still contains {len(result_general)} general questions!")
            
        elif question_strategy == 'general_only' and len(result_specific) > 0:
            print(f"WARNING: general_only strategy still contains {len(result_specific)} specific questions!")
    
    return result

def aggregate_predictions_by_questionnaire(test_df, outcome_names, questionnaire_mapping):
    """Aggregate predictions by questionnaire"""
    print(f"DEBUG: Starting questionnaire aggregation for {len(test_df)} QA pairs")
    print(f"DEBUG: Unique users in test set: {test_df['user_id'].nunique()}")
    print(f"DEBUG: Questionnaire mapping keys: {list(questionnaire_mapping.keys())}")
    
    user_level_results = []
    
    # 处理每个用户
    for user_id in test_df['user_id'].unique():
        user_data = test_df[test_df['user_id'] == user_id]
        
        # 获取该用户的真实问卷分数 (应该所有QA对的真实分数都相同)
        true_labels = user_data['output_embeddings'].iloc[0]  # 取第一个，因为同一用户的真实分数应该相同
        
        print(f"DEBUG: Processing user {user_id} with {len(user_data)} QA pairs")
        print(f"DEBUG: True labels shape: {len(true_labels) if isinstance(true_labels, (list, np.ndarray)) else 'scalar'}")
        
        # 处理每个问卷
        for outcome_idx, outcome_name in enumerate(outcome_names):
            true_score = true_labels[outcome_idx]
            
            # 获取该问卷对应的问题列表
            if outcome_name in questionnaire_mapping:
                questionnaire_questions = set(questionnaire_mapping[outcome_name])
                
                # 找到属于该问卷的QA对
                relevant_qa = user_data[user_data['question_code'].isin(questionnaire_questions)]
                
                if len(relevant_qa) > 0:
                    # 从每个QA对的预测中提取该问卷对应的输出
                    questionnaire_predictions = []
                    for _, qa_row in relevant_qa.iterrows():
                        # qa_row['predictions'] 应该是长度为10的列表 (multitask输出)
                        pred_vector = qa_row['predictions']
                        if isinstance(pred_vector, (list, np.ndarray)) and len(pred_vector) > outcome_idx:
                            questionnaire_predictions.append(pred_vector[outcome_idx])
                        else:
                            print(f"WARNING: Invalid prediction format for user {user_id}, question {qa_row['question_code']}")
                    
                    # 对该问卷的所有QA对预测求平均值
                    if len(questionnaire_predictions) > 0:
                        final_pred_score = np.mean(questionnaire_predictions)
                        print(f"DEBUG: {outcome_name} for user {user_id}: {len(questionnaire_predictions)} predictions, avg = {final_pred_score:.4f}, true = {true_score:.4f}")
                    else:
                        # 如果没有找到相关问题，使用所有问题的平均值作为fallback
                        all_predictions = [pred[outcome_idx] for pred in user_data['predictions'] 
                                         if isinstance(pred, (list, np.ndarray)) and len(pred) > outcome_idx]
                        final_pred_score = np.mean(all_predictions) if all_predictions else 0.0
                        print(f"DEBUG: {outcome_name} for user {user_id}: No specific questions found, using fallback avg = {final_pred_score:.4f}")
                else:
                    # 如果没有找到该问卷的问题，使用所有问题的平均值
                    all_predictions = [pred[outcome_idx] for pred in user_data['predictions'] 
                                     if isinstance(pred, (list, np.ndarray)) and len(pred) > outcome_idx]
                    final_pred_score = np.mean(all_predictions) if all_predictions else 0.0
                    print(f"DEBUG: {outcome_name} for user {user_id}: No relevant QA pairs, using all questions avg = {final_pred_score:.4f}")
            else:
                # 问卷不在映射中，使用所有问题的平均值
                all_predictions = [pred[outcome_idx] for pred in user_data['predictions'] 
                                 if isinstance(pred, (list, np.ndarray)) and len(pred) > outcome_idx]
                final_pred_score = np.mean(all_predictions) if all_predictions else 0.0
                print(f"DEBUG: {outcome_name} for user {user_id}: Not in mapping, using all questions avg = {final_pred_score:.4f}")
            
            # 存储结果
            user_level_results.append({
                'user_id': user_id,
                'questionnaire': outcome_name,
                'true_score': true_score,
                'pred_score': final_pred_score
            })
    
    result_df = pd.DataFrame(user_level_results)
    
    print(f"DEBUG: Final user-level results shape: {result_df.shape}")
    print(f"DEBUG: Sample results:")
    for outcome in outcome_names[:3]:  # 显示前3个问卷的样本
        outcome_data = result_df[result_df['questionnaire'] == outcome]
        if len(outcome_data) > 0:
            print(f"  {outcome}: {len(outcome_data)} users, true range: {outcome_data['true_score'].min():.3f}-{outcome_data['true_score'].max():.3f}, pred range: {outcome_data['pred_score'].min():.3f}-{outcome_data['pred_score'].max():.3f}")
    
    return result_df

def qa_level_linear_regression_training_with_hyperparameter_search(train_inputs, train_true, dev_inputs, dev_true, 
                                                                  test_inputs, test_true, task_type, long_training=False):
    """Train with hyperparameter search"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Simplified hyperparameter grid
    param_grid = {
        'lr': [0.01, 0.001],
        'weight_decay': [1e-5, 1e-4],
        'batch_size': [64],
        'epochs': [500],
        'patience': [50]
    }
    
    best_dev_loss = float('inf')
    best_model = None
    best_params = None
    best_test_pred = None
    
    print(f"    Starting hyperparameter search with {len(list(ParameterGrid(param_grid)))} configurations...")
    
    for params in ParameterGrid(param_grid):
        try:
            model = LinearRegressionMultitaskQA(train_inputs.shape[1], train_true.shape[1]).to(device)
            
            train_inputs_gpu = train_inputs.to(device)
            train_true_gpu = train_true.to(device)
            dev_inputs_gpu = dev_inputs.to(device)
            dev_true_gpu = dev_true.to(device)
            test_inputs_gpu = test_inputs.to(device)
            
            optimizer = torch.optim.Adam(model.parameters(), lr=params['lr'], weight_decay=params['weight_decay'])
            scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=params['patience']//4)
            criterion = nn.MSELoss()
            
            model_best_loss = float('inf')
            patience_counter = 0
            
            dataset = torch.utils.data.TensorDataset(train_inputs_gpu, train_true_gpu)
            dataloader = torch.utils.data.DataLoader(dataset, batch_size=params['batch_size'], shuffle=True)
            
            for epoch in range(params['epochs']):
                model.train()
                for batch_inputs, batch_targets in dataloader:
                    optimizer.zero_grad()
                    predictions = model(batch_inputs)
                    loss = criterion(predictions, batch_targets)
                    loss.backward()
                    clip_grad_norm_(model.parameters(), max_norm=1.0)
                    optimizer.step()
                
                # Validation check every 20 epochs.
                # Note: patience_counter increments only at these checkpoints,
                # so effective patience = params['patience'] × 20 epochs.
                if epoch % 20 == 0:
                    model.eval()
                    with torch.no_grad():
                        dev_pred = model(dev_inputs_gpu)
                        dev_loss = criterion(dev_pred, dev_true_gpu)
                    
                    scheduler.step(dev_loss)
                    
                    if dev_loss < model_best_loss:
                        model_best_loss = dev_loss
                        patience_counter = 0
                    else:
                        patience_counter += 1
                        
                    if patience_counter >= params['patience']:
                        break
            
            model.eval()
            with torch.no_grad():
                final_dev_pred = model(dev_inputs_gpu)
                final_dev_loss = criterion(final_dev_pred, dev_true_gpu).item()
            
            if final_dev_loss < best_dev_loss:
                best_dev_loss = final_dev_loss
                best_params = params.copy()
                best_params['final_epoch'] = epoch
                
                with torch.no_grad():
                    best_test_pred = model(test_inputs_gpu)
                
                best_model = type(model)(train_inputs.shape[1], train_true.shape[1]).to(device)
                best_model.load_state_dict(model.state_dict())
        
        except Exception as e:
            print(f"    Error with params {params}: {e}")
            continue
    
    if best_params:
        print(f"    Best hyperparameters: lr={best_params['lr']}, wd={best_params['weight_decay']}")
        print(f"    Best dev loss: {best_dev_loss:.6f}")
    else:
        print(f"    No valid hyperparameters found, using defaults")
        best_params = {'lr': 0.01, 'weight_decay': 1e-4, 'batch_size': 64, 'epochs': 500, 'patience': 50}
        
        model = LinearRegressionMultitaskQA(train_inputs.shape[1], train_true.shape[1]).to(device)
        optimizer = torch.optim.Adam(model.parameters(), lr=best_params['lr'], weight_decay=best_params['weight_decay'])
        criterion = nn.MSELoss()
        
        train_inputs_gpu = train_inputs.to(device)
        train_true_gpu = train_true.to(device)
        dev_inputs_gpu = dev_inputs.to(device)
        dev_true_gpu = dev_true.to(device)
        test_inputs_gpu = test_inputs.to(device)
        
        fallback_best_loss = float('inf')
        fallback_patience = 0
        
        for epoch in range(best_params['epochs']):
            model.train()
            optimizer.zero_grad()
            train_pred = model(train_inputs_gpu)
            loss = criterion(train_pred, train_true_gpu)
            loss.backward()
            clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            
            # Early stopping based on dev loss (check every 20 epochs)
            if epoch % 20 == 0:
                model.eval()
                with torch.no_grad():
                    dev_pred = model(dev_inputs_gpu)
                    dev_loss = criterion(dev_pred, dev_true_gpu)
                if dev_loss < fallback_best_loss:
                    fallback_best_loss = dev_loss
                    fallback_patience = 0
                else:
                    fallback_patience += 1
                if fallback_patience >= best_params['patience']:
                    break
        
        model.eval()
        with torch.no_grad():
            best_test_pred = model(test_inputs_gpu)
        
        best_model = model
    
    return best_model, best_test_pred, best_params


def calculate_questionnaire_metrics(user_level_results, outcome_names):
    """
    计算每个问卷的 correlation 和 MSE
    
    Args:
        user_level_results: DataFrame with columns [user_id, questionnaire, true_score, pred_score]
        outcome_names: 问卷名称列表
    
    Returns:
        correlations: dict {questionnaire_name: correlation_value}
        mse_values: dict {questionnaire_name: mse_value}
    """
    from scipy.stats import pearsonr
    
    correlations = {}
    mse_values = {}
    
    print(f"DEBUG: Calculating metrics for {len(outcome_names)} questionnaires")
    
    for outcome in outcome_names:
        outcome_data = user_level_results[user_level_results['questionnaire'] == outcome]
        
        if len(outcome_data) > 1:  # 需要至少2个点来计算correlation
            true_scores = outcome_data['true_score'].values
            pred_scores = outcome_data['pred_score'].values
            
            # 计算correlation
            corr, p_value = pearsonr(true_scores, pred_scores)
            correlations[outcome] = corr if not np.isnan(corr) else 0.0
            
            # 计算MSE
            mse = np.mean((true_scores - pred_scores) ** 2)
            mse_values[outcome] = mse
            
            print(f"DEBUG: {outcome}: n={len(outcome_data)}, corr={correlations[outcome]:.4f}, mse={mse_values[outcome]:.4f}")
        else:
            print(f"WARNING: {outcome} has insufficient data points ({len(outcome_data)}) for correlation calculation")
            correlations[outcome] = 0.0
            mse_values[outcome] = float('inf')
    
    return correlations, mse_values


def train_qa_level_model_ultimate(qa_level_data, n_folds=9, normalization='minmax', long_training=False):
    """
    Ultimate training function that takes pre-prepared QA-level data
    """
    
    outcome_names = ["PHQ", "GAD", "MDQ", "RAADS", "DUDIT", "AUDIT", "BOCS", "ASRS", "NSE", "EDE_QS"]
    
    print(f"\nTraining QA-Level Questionnaire-Specific Multitask (FIXED VERSION)")
    print(f"Normalization: {normalization}")
    print(f"Using {n_folds}-fold cross-validation with questionnaire-specific aggregation")
    
    # 验证数据结构
    if len(qa_level_data) == 0:
        print("ERROR: No QA pairs in the dataset!")
        return {outcome: [] for outcome in outcome_names}, {outcome: [] for outcome in outcome_names}, []
    
    print(f"QA-level dataset shape: {qa_level_data.shape}")
    print(f"Number of unique users: {qa_level_data['user_id'].nunique()}")
    print(f"Number of QA pairs: {len(qa_level_data)}")
    print(f"Input dimension: {len(qa_level_data['input_embeddings'].iloc[0])}")
    print(f"Output dimension: {len(qa_level_data['output_embeddings'].iloc[0])}")
    
    # 获取问卷映射
    questionnaire_mapping_data = get_questionnaire_mapping()
    questionnaire_mapping = questionnaire_mapping_data['questionnaire_specific']
    
    print(f"DEBUG: Using questionnaire mapping with {len(questionnaire_mapping)} questionnaires")
    
    bins = create_stratified_cv_folds(qa_level_data, n_folds)
    
    all_test_correlations = {outcome: [] for outcome in outcome_names}
    all_test_mse = {outcome: [] for outcome in outcome_names}
    fold_details = []
    
    for i in range(n_folds):
        print(f"  Processing fold {i+1}/{n_folds}...")
        
        # 分割用户
        test_users = bins[i]
        train_dev_bins = list(set(range(n_folds)) - {i})
        random.shuffle(train_dev_bins)
        
        n_train_bins = max(1, (len(train_dev_bins) * 2) // 3)
        train_users = [user for j in train_dev_bins[:n_train_bins] for user in bins[j]]
        dev_users = [user for j in train_dev_bins[n_train_bins:] for user in bins[j]]

        train_set = set(train_users)
        dev_set = set(dev_users)
        test_set = set(test_users)
        if not train_set.isdisjoint(dev_set):
            raise ValueError(f"User leakage between train/dev in fold {i}")
        if not train_set.isdisjoint(test_set):
            raise ValueError(f"User leakage between train/test in fold {i}")
        if not dev_set.isdisjoint(test_set):
            raise ValueError(f"User leakage between dev/test in fold {i}")
        
        # 获取QA对
        train = qa_level_data[qa_level_data.user_id.isin(train_users)]
        dev = qa_level_data[qa_level_data.user_id.isin(dev_users)]
        test = qa_level_data[qa_level_data.user_id.isin(test_users)]
        
        print(f"    Train QA pairs: {len(train)} ({len(train_users)} users)")
        print(f"    Dev QA pairs: {len(dev)} ({len(dev_users)} users)")
        print(f"    Test QA pairs: {len(test)} ({len(test_users)} users)")
        
        # 数据转换和标准化 (代码与原版相同)
        train_inputs_raw = np.array(train["input_embeddings"].tolist())
        train_outputs_raw = np.array(train["output_embeddings"].tolist())
        dev_inputs_raw = np.array(dev["input_embeddings"].tolist())
        dev_outputs_raw = np.array(dev["output_embeddings"].tolist())
        test_inputs_raw = np.array(test["input_embeddings"].tolist())
        test_outputs_raw = np.array(test["output_embeddings"].tolist())
        
        # 应用标准化
        if normalization == 'standard':
            scaler = StandardScaler()
            train_inputs_norm = scaler.fit_transform(train_inputs_raw)
            dev_inputs_norm = scaler.transform(dev_inputs_raw)
            test_inputs_norm = scaler.transform(test_inputs_raw)
        elif normalization == 'minmax':
            scaler = MinMaxScaler()
            train_inputs_norm = scaler.fit_transform(train_inputs_raw)
            dev_inputs_norm = scaler.transform(dev_inputs_raw)
            test_inputs_norm = scaler.transform(test_inputs_raw)
        else:
            train_inputs_norm = train_inputs_raw
            dev_inputs_norm = dev_inputs_raw
            test_inputs_norm = test_inputs_raw
        
        # 标准化输出
        output_scaler = StandardScaler()
        train_outputs_norm = output_scaler.fit_transform(train_outputs_raw)
        dev_outputs_norm = output_scaler.transform(dev_outputs_raw)
        
        # 准备张量
        train_inputs = torch.tensor(train_inputs_norm, dtype=torch.float32)
        train_true = torch.tensor(train_outputs_norm, dtype=torch.float32)
        dev_inputs = torch.tensor(dev_inputs_norm, dtype=torch.float32)
        dev_true = torch.tensor(dev_outputs_norm, dtype=torch.float32)
        test_inputs = torch.tensor(test_inputs_norm, dtype=torch.float32)
        test_true = torch.tensor(test_outputs_raw, dtype=torch.float32)
        
        # 训练模型
        model, test_pred, best_params = qa_level_linear_regression_training_with_hyperparameter_search(
            train_inputs, train_true, dev_inputs, dev_true, test_inputs, test_true, 'multitask', long_training
        )
        
        # 反标准化预测
        test_pred_denorm = output_scaler.inverse_transform(test_pred.cpu().numpy())
        
        # 准备测试数据DataFrame
        test_df = test.copy()
        test_df['predictions'] = [pred.tolist() for pred in test_pred_denorm]
        
        print(f"    Applying FIXED questionnaire-specific aggregation...")
        
        # 使用修正的聚合函数
        user_level_results = aggregate_predictions_by_questionnaire(test_df, outcome_names, questionnaire_mapping)
        
        # 计算指标
        fold_correlations, fold_mse = calculate_questionnaire_metrics(user_level_results, outcome_names)
        
        # 存储结果
        for outcome in outcome_names:
            if outcome in fold_correlations:
                all_test_correlations[outcome].append(fold_correlations[outcome])
                all_test_mse[outcome].append(fold_mse[outcome])
        
        fold_details.append({
            'fold': i+1,
            'train_qa_pairs': len(train),
            'test_qa_pairs': len(test),
            'train_users': len(train_users),
            'test_users': len(test_users),
            'correlations': fold_correlations,
            'mse': fold_mse,
            'best_hyperparameters': best_params
        })
        
        print(f"    Fold {i+1} Results (FIXED):")
        for outcome in outcome_names[:3]:
            if outcome in fold_correlations:
                print(f"      {outcome}: Corr={fold_correlations[outcome]:.4f}, MSE={fold_mse[outcome]:.4f}")
    
    # 打印总体结果
    print(f"\nOverall QA-Level Results (FIXED):")
    for outcome in outcome_names:
        if outcome in all_test_correlations and all_test_correlations[outcome]:
            corr_mean = np.mean(all_test_correlations[outcome])
            corr_std = np.std(all_test_correlations[outcome])
            mse_mean = np.mean(all_test_mse[outcome])
            mse_std = np.std(all_test_mse[outcome])
            print(f"{outcome}: Corr={corr_mean:.4f}±{corr_std:.4f}, MSE={mse_mean:.4f}±{mse_std:.4f}")
        else:
            print(f"{outcome}: No valid results")
    
    return all_test_correlations, all_test_mse, fold_details

def convert_numpy_to_python(obj):
    """Convert numpy types to Python native types"""
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, dict):
        return {key: convert_numpy_to_python(value) for key, value in obj.items()}
    elif isinstance(obj, list):
        return [convert_numpy_to_python(item) for item in obj]
    else:
        return obj

def fetch_question_embeddings(question_file):
    """Load question embeddings"""
    return load_question_embeddings(question_file)

def run_ultimate_questionnaire_specific_comparison():
    """Ultimate comparison function with robust error handling"""
    print("="*120)
    print("QA-LEVEL QUESTIONNAIRE-SPECIFIC COMPARISON - ULTIMATE FIXED VERSION")
    print("="*120)
    
    # Set random seed
    random.seed(42)
    torch.manual_seed(42)
    np.random.seed(42)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(42)
    
    # Load data
    print("Loading data...")
    merged_data = build_qa_merged_dataset()
    print(f"Final dataset shape: {merged_data.shape}")
    
    # Test strategies
    question_strategies = ['specific_only', 'all_questions', 'general_only', 'combined']
    embedding_strategy = 'full'
    normalization = 'minmax'
    outcome_names = ["PHQ", "GAD", "MDQ", "RAADS", "DUDIT", "AUDIT", "BOCS", "ASRS", "NSE", "EDE_QS"]
    
    all_results = {}
    summary_results = []
    
    for question_strategy in question_strategies:
        print(f"\n{'='*80}")
        print(f"TESTING STRATEGY: {question_strategy.upper()}")
        print(f"{'='*80}")
        
        config_name = f"qa_level_questionnaire_specific_{question_strategy}_{embedding_strategy}_{normalization}_multitask"
        
        try:
            # Get questionnaire mapping with proper general/specific distinction
            questionnaire_mapping_data = get_questionnaire_mapping()

            # Prepare QA-level data ONCE with ultimate function
            qa_level_data = prepare_qa_level_data_ultimate(
                merged_data, embedding_strategy=embedding_strategy, separate_dudit_audit=True
            )
            
            # Filter by strategy using the updated mapping
            strategy_data = filter_qa_pairs_by_strategy(
                qa_level_data, question_strategy, questionnaire_mapping_data
            )
            
            print(f"Data for strategy '{question_strategy}': {len(strategy_data)} pairs from {strategy_data['user_id'].nunique()} users.")
            
            # Skip if no data
            if len(strategy_data) == 0:
                print(f"SKIPPING {question_strategy}: No QA pairs found for this strategy")
                continue
            
            # Train the model
            correlations, mse, fold_details = train_qa_level_model_ultimate(
                strategy_data.copy(),
                n_folds=9,
                normalization=normalization,
                long_training=False
            )
            
            # Calculate summary statistics
            avg_correlations = {}
            avg_mse = {}
            for outcome in outcome_names:
                if outcome in correlations and correlations[outcome]:
                    corr_vals = correlations[outcome]
                    mse_vals = mse[outcome]
                    avg_correlations[outcome] = {
                        'mean': np.mean(corr_vals),
                        'std': np.std(corr_vals),
                        'values': corr_vals
                    }
                    avg_mse[outcome] = {
                        'mean': np.mean(mse_vals),
                        'std': np.std(mse_vals),
                        'values': mse_vals
                    }
            
            overall_avg_corr = np.mean([v['mean'] for v in avg_correlations.values()]) if avg_correlations else 0.0
            overall_avg_mse = np.mean([v['mean'] for v in avg_mse.values()]) if avg_mse else float('inf')
            
            # Store results
            all_results[config_name] = {
                'correlations': correlations,
                'mse': mse,
                'fold_details': fold_details,
                'question_strategy': question_strategy,
                'embedding_strategy': embedding_strategy,
                'normalization': normalization,
                'task_type': 'qa_level_questionnaire_specific',
                'avg_correlations': avg_correlations,
                'avg_mse': avg_mse
            }
            
            summary_results.append({
                'config': config_name,
                'question_strategy': question_strategy,
                'embedding_strategy': embedding_strategy,
                'normalization': normalization,
                'task_type': 'qa_level_questionnaire_specific',
                'overall_avg_corr': overall_avg_corr,
                'overall_avg_mse': overall_avg_mse,
                'outcome_correlations': avg_correlations,
                'outcome_mse': avg_mse
            })
            
            print(f"✓ Completed {config_name}: Avg Corr = {overall_avg_corr:.4f}, Avg MSE = {overall_avg_mse:.4f}")
            
        except Exception as e:
            print(f"✗ Error in {config_name}: {e}")
            traceback.print_exc()
    
    # Results Analysis
    print(f"\n{'='*120}")
    print("QUESTIONNAIRE-SPECIFIC RESULTS ANALYSIS")
    print(f"{'='*120}")
    
    if summary_results:
        # Sort results by correlation
        summary_results.sort(key=lambda x: x['overall_avg_corr'], reverse=True)
        
        print(f"\n{'QUESTIONNAIRE-SPECIFIC CONFIGURATIONS (by average correlation)'}")
        print(f"{'='*120}")
        print(f"{'Rank':<5} {'Question Strategy':<15} {'Embedding':<12} {'Norm':<8} {'Avg Corr':<12} {'Avg MSE':<12}")
        print("-"*85)
        
        for i, result in enumerate(summary_results, 1):
            print(f"{i:<5} {result['question_strategy']:<15} {result['embedding_strategy']:<12} "
                  f"{result['normalization']:<8} {result['overall_avg_corr']:.4f}      {result['overall_avg_mse']:.4f}")
        
        # Strategy impact analysis
        print(f"\n{'QUESTION STRATEGY IMPACT ANALYSIS'}")
        print(f"{'='*80}")
        
        for strategy in question_strategies:
            strategy_results = [r for r in summary_results if r['question_strategy'] == strategy]
            if strategy_results:
                result = strategy_results[0]
                print(f"{strategy:<15}: Corr={result['overall_avg_corr']:.4f}, MSE={result['overall_avg_mse']:.4f}")
            else:
                print(f"{strategy:<15}: No results (likely skipped)")
        
        # Best configuration per outcome
        print(f"\n{'BEST CONFIGURATION FOR EACH OUTCOME'}")
        print(f"{'='*120}")
        print(f"{'Outcome':<12} {'Question Strategy':<15} {'Correlation':<12} {'MSE':<12}")
        print("-"*60)
        
        for outcome in outcome_names:
            best_corr = -1
            best_config = None
            
            for result in summary_results:
                if outcome in result['outcome_correlations']:
                    corr = result['outcome_correlations'][outcome]['mean']
                    if corr > best_corr:
                        best_corr = corr
                        best_config = result
            
            if best_config:
                mse_val = best_config['outcome_mse'][outcome]['mean']
                print(f"{outcome:<12} {best_config['question_strategy']:<15} {best_corr:.4f}      {mse_val:.4f}")
            else:
                print(f"{outcome:<12} {'No results':<15} {'N/A':<12} {'N/A':<12}")
        
        # Save results
        print(f"\nSaving results...")
        
        output_dir = "model_outputs/qa_level_questionnaire_specific_comparison_ultimate"
        os.makedirs(output_dir, exist_ok=True)
        
        # Save detailed results
        json_safe_results = convert_numpy_to_python(all_results)
        with open(f"{output_dir}/detailed_results.json", 'w') as f:
            json.dump(json_safe_results, f, indent=2)
        
        # Save summary results  
        json_safe_summary = convert_numpy_to_python(summary_results)
        with open(f"{output_dir}/summary_results.json", 'w') as f:
            json.dump(json_safe_summary, f, indent=2)
        
        # Create comparison CSV
        comparison_rows = []
        for result in summary_results:
            comparison_rows.append({
                'Question_Strategy': result['question_strategy'],
                'Overall_Correlation': result['overall_avg_corr'],
                'Overall_MSE': result['overall_avg_mse'],
                'Embedding_Strategy': result['embedding_strategy'],
                'Normalization': result['normalization']
            })
        
        comparison_df = pd.DataFrame(comparison_rows)
        comparison_df.to_csv(f"{output_dir}/strategy_comparison_summary.csv", index=False)
        
        # Save configuration for downstream consumers (e.g. compute_tables.py)
        config_info = {
            'outcome_names': ["PHQ", "GAD", "MDQ", "RAADS", "DUDIT", "AUDIT", "BOCS", "ASRS", "NSE", "EDE_QS"],
            'model_type': 'multitask_qa_level',
            'question_strategies': question_strategies,
            'embedding_strategy': embedding_strategy,
            'normalization': normalization,
            'n_folds': 9
        }
        with open(f"{output_dir}/config.json", 'w') as f:
            json.dump(config_info, f, indent=2)
        
        print("Ultimate questionnaire-specific QA-level comparison completed!")
        print(f"Results saved to: {output_dir}")
        
        # Print final summary
        best_strategy = summary_results[0]
        print(f"\n🏆 BEST PERFORMING STRATEGY:")
        print(f"   Strategy: {best_strategy['question_strategy']}")
        print(f"   Overall Correlation: {best_strategy['overall_avg_corr']:.4f}")
        print(f"   Overall MSE: {best_strategy['overall_avg_mse']:.4f}")
        
    else:
        print("No successful results to analyze!")
    
    return all_results, summary_results

if __name__ == "__main__":
    print("="*80)
    print("QA-LEVEL QUESTIONNAIRE-SPECIFIC COMPARISON - ULTIMATE FIXED VERSION")
    print("="*80)
    print("This script implements questionnaire-specific EVALUATION where:")
    print("- Training: Use ALL QA pairs (multitask learns from all questionnaires)")
    print("- Prediction: Use ALL QA pairs (model predicts all outcomes)")
    print("- Evaluation: For each outcome, aggregate only that questionnaire's QA pairs")
    print()
    print("QUESTION STRATEGIES FOR MULTITASK:")
    print("📊 all_questions: Use ALL questions (questionnaire-specific + general)")
    print("🎯 specific_only: Use ONLY questionnaire-specific questions (remove general)")  
    print("🌐 general_only: Use ONLY general questions (not in any questionnaire)")
    print("🔄 combined: Same as all_questions for multitask models")
    print()
    print("ULTIMATE FIXES INCLUDED:")
    print("✅ Smart float conversion handling different data types")
    print("✅ Robust input/output data processing") 
    print("✅ Proper data flow between functions")
    print("✅ Comprehensive error handling and debugging")
    print("✅ Single-pass data preparation to avoid conflicts")
    print("✅ Corrected filtering logic for multitask strategies")
    print("="*80)
    
    if not sys.stdin.isatty():
        print("Non-interactive mode detected. Skipping interactive run. Use run_qa_level_save_outputs.py for automation.")
        sys.exit(0)

    choice = input("Run ultimate questionnaire-specific QA-level comparison? (y/n): ").strip().lower()
    
    if choice == 'y':
        print("\nRunning ultimate questionnaire-specific QA-level comparison...")
        try:
            results, summary = run_ultimate_questionnaire_specific_comparison()
            print("\n🎉 Comparison completed successfully!")
            
            if summary:
                print(f"📊 Results summary:")
                print(f"   - Total configurations tested: {len(summary)}")
                print(f"   - Best overall correlation: {max(s['overall_avg_corr'] for s in summary):.4f}")
                print(f"   - Results saved to: model_outputs/qa_level_questionnaire_specific_comparison_ultimate/")
            else:
                print("⚠️  No successful configurations - check the logs above")
                
        except Exception as e:
            print(f"\n❌ Fatal error during comparison: {e}")
            traceback.print_exc()
    else:
        print("Comparison cancelled.")
        
    print("\n" + "="*80)
    print("QA-LEVEL QUESTIONNAIRE-SPECIFIC COMPARISON - END")
    print("="*80)