path: root/collaborativeagents/scripts/run_experiments.py

#!/usr/bin/env python3
"""
Main experiment orchestrator for personalization benchmark.

This script runs all baselines and the proposed methods with PROPER multi-turn
conversation simulation, user preference enforcement, and LLM-based evaluation.

Usage:
    python run_experiments.py --config config.yaml
    python run_experiments.py --methods vanilla,rag,rag_vector --datasets gpqa,aime
"""

import argparse
import json
import yaml
import os
import sys
import numpy as np
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Any, Optional
from dataclasses import dataclass, asdict
import logging
import re
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading
import time

# Add paths
sys.path.insert(0, str(Path(__file__).parent.parent))
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from datasets_extended import get_dataset, get_all_datasets, get_challenging_datasets
from evaluation.llm_judge import LLMJudge, BatchEvaluator, ConversationMetrics
from conflict_scenario_generator import ConflictScenarioGenerator
from adapters.personalized_llm_adapter import PersonalizedLLMAdapter, create_baseline_adapter
from agents.local_user_agent import LocalUserAgent, SharedLocalUserAgent, TERMINATION_SIGNAL
from agents.vllm_user_agent import VLLMUserAgent, VLLMAgentClient
from agents.openai_user_agent import OpenAIUserAgent
from agents.batch_vllm_agent import BatchConversationGenerator, BatchVLLMClient


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


@dataclass
class ExperimentConfig:
    """Configuration for an experiment run."""
    # Methods to compare
    methods: List[str]

    # Datasets to use
    datasets: List[str]

    # User profiles
    n_profiles: int = 200
    profile_path: Optional[str] = None

    # Profile range (for splitting jobs)
    start_profile: int = 0  # Inclusive, 0-indexed
    end_profile: Optional[int] = None  # Exclusive, None means all

    # Session settings
    n_sessions_per_profile: int = 30
    max_turns_per_session: int = 15  # Increased for harder tasks

    # Model settings
    user_model: str = "meta-llama/Llama-3.3-70B-Instruct"
    agent_model: str = "meta-llama/Llama-3.1-8B-Instruct"
    judge_model: str = "meta-llama/Llama-3.3-70B-Instruct"

    # Output settings
    output_dir: str = "results"
    save_conversations: bool = True

    # Conflict testing
    conflict_ratio: float = 0.3  # proportion of queries that trigger conflicts

    # Compute settings
    batch_size: int = 4
    n_gpus: int = 4

    # vLLM settings (for high-performance inference)
    use_vllm: bool = False
    vllm_user_url: str = "http://localhost:8004/v1"  # 70B user simulator
    vllm_agent_url: str = "http://localhost:8003/v1"  # 8B agent

    # OpenAI user simulator (alternative to vLLM user agent)
    use_openai_user: bool = False
    openai_user_model: str = "gpt-5"  # Model name for OpenAI user agent

    # Reward mode: "keyword" (implicit user signals), "llm" (GPT-4o-mini), or "llm_local" (local vLLM)
    # This is a global option applied to ALL methods that use RL updates
    reward_mode: str = "keyword"

    # vLLM URL for local reward model (only used when reward_mode="llm_local")
    reward_vllm_url: str = "http://localhost:8005/v1"

    # Parallel/Batch processing
    parallel_profiles: int = 50  # Number of profiles to process in parallel
    use_batch_processing: bool = True  # Use turn-synchronous batch processing for vanilla/all_memory
    batch_size_conversations: int = 50  # Number of conversations to batch together

    # Continue from existing experiment (for extending sessions)
    continue_from: Optional[str] = None  # Path to existing output directory to continue from


# Available methods
AVAILABLE_METHODS = {
    "vanilla": "No memory, no personalization",
    "contextual": "Full history in context, summarize when overflow",
    "reflection": "CollaborativeAgents' agent_notes approach",
    "reflection_grpo": "Reflection + GRPO training",
    "all_memory": "All extracted memories in context (no retrieval)",
    "rag": "Extractor + RAG (no user vector)",
    "rag_dynamic": "Extractor + RAG with dynamic topk (min=3, max=8, ratio=0.5)",
    "rag_rewrite": "Extractor + RAG with LLM preference rewrite/merge",
    "rag_rewrite_vector": "Extractor + RAG + user vector + LLM preference rewrite",
    "rag_vector": "Extractor + RAG + user vector (proposed method)",
    "rag_vector_fast": "Extractor + RAG + user vector with 10x learning rate",
    "rag_vector_consolidate": "Extractor + RAG + user vector with session-level preference consolidation",
    "rag_vector_balanced": "Extractor + RAG + user vector with balanced rewards (10x LR + positive signal for good turns)",
    "rag_bge": "Extractor + RAG with BGE reranker (278M)",
    "rag_vector_bge": "Extractor + RAG + user vector with BGE reranker (278M)",
}


class ExperimentRunner:
    """Main experiment runner."""

    def __init__(self, config: ExperimentConfig):
        self.config = config

        # Use existing directory if continuing, otherwise create new timestamped one
        if config.continue_from:
            self.output_dir = Path(config.continue_from)
            if not self.output_dir.exists():
                raise ValueError(f"Continue-from directory does not exist: {config.continue_from}")
            logger.info(f"Continuing from existing experiment: {self.output_dir}")
        else:
            self.output_dir = Path(config.output_dir) / datetime.now().strftime("%Y%m%d_%H%M%S")
            self.output_dir.mkdir(parents=True, exist_ok=True)

        # Save/update config
        with open(self.output_dir / "config.yaml", "w") as f:
            yaml.dump(asdict(config), f)

        # Initialize components
        self.judge = LLMJudge(model_name=config.judge_model)
        self.batch_evaluator = BatchEvaluator(self.judge)
        self.conflict_generator = ConflictScenarioGenerator()

        # Load datasets
        self.datasets = {}
        for ds_name in config.datasets:
            try:
                self.datasets[ds_name] = get_dataset(ds_name)
                logger.info(f"Loaded dataset: {ds_name}")
            except Exception as e:
                logger.warning(f"Failed to load dataset {ds_name}: {e}")

        # Load or generate profiles
        self.profiles = self._load_profiles()

    def _load_profiles(self) -> List[Dict]:
        """Load user profiles from file or generate."""
        logger.info(f"Profile path configured: {self.config.profile_path}")

        if self.config.profile_path:
            profile_path = Path(self.config.profile_path)
            if profile_path.exists():
                profiles = []
                with open(profile_path) as f:
                    for line in f:
                        line = line.strip()
                        if line:
                            profiles.append(json.loads(line))
                logger.info(f"Loaded {len(profiles)} profiles from {self.config.profile_path}")
                return profiles[:self.config.n_profiles]
            else:
                logger.warning(f"Profile path does not exist: {self.config.profile_path}")

        # Generate simple placeholder profiles if no file provided
        logger.info(f"Generating {self.config.n_profiles} placeholder profiles...")
        profiles = []
        for i in range(self.config.n_profiles):
            profiles.append({
                "id": i,
                "persona": f"User {i+1} is a curious individual seeking help with problem solving.",
                "preferences": [
                    "Provide clear, step-by-step explanations",
                    "Use simple language when possible",
                    "Give examples to illustrate concepts",
                    "Be concise but thorough",
                    "Acknowledge when something is uncertain"
                ]
            })

        # Save generated profiles
        profile_path = self.output_dir / "generated_profiles.json"
        with open(profile_path, "w") as f:
            json.dump(profiles, f, indent=2)

        logger.info(f"Generated and saved {len(profiles)} placeholder profiles")
        return profiles

    def _create_method_adapter(self, method: str, profile: Dict, use_shared_models: bool = False) -> Any:
        """Create adapter for a specific method.

        Args:
            method: One of the baseline method names
            profile: User profile dict (used later in start_session, not constructor)
            use_shared_models: If True, share embedding/reranker models across parallel
                workers. ESSENTIAL for parallel profile processing to avoid OOM.

        Returns:
            Configured adapter instance
        """
        # Auto-detect available GPUs and set device assignment accordingly
        # Layout with local 70B user (4 GPUs):
        #   GPU 0-1: 70B user simulator (TP=2)
        #   GPU 2: 8B agent vLLM server
        #   GPU 3: Embedding + Reranker + Extractor
        # Layout with OpenAI user (2 GPUs):
        #   GPU 0: 8B agent vLLM server
        #   GPU 1: Embedding + Reranker + Extractor
        device_assignment = None
        try:
            import torch
            n_gpus = torch.cuda.device_count()
            if n_gpus >= 4:
                # 4 GPU layout: 70B user on 0-1, agent on 2, adapters on 3
                device_assignment = {
                    "embed": "cuda:3",
                    "reranker": "cuda:3",
                    "extractor": "cuda:3",
                }
            elif n_gpus >= 2:
                # 2 GPU layout: agent on 0, adapters on 1
                device_assignment = {
                    "embed": "cuda:1",
                    "reranker": "cuda:1",
                    "extractor": "cuda:1",
                }
            elif n_gpus == 1:
                device_assignment = {
                    "embed": "cuda:0",
                    "reranker": "cuda:0",
                    "extractor": "cuda:0",
                }
        except ImportError:
            pass

        adapter = create_baseline_adapter(
            method,
            device_assignment=device_assignment,
            use_vllm=self.config.use_vllm,
            use_shared_models=use_shared_models,
            reward_mode=self.config.reward_mode,
            reward_vllm_url=self.config.reward_vllm_url,
        )
        # Profile will be passed to start_session() when the conversation begins
        return adapter

    def _export_user_vectors(self, method: str, adapters: Dict[int, Any]) -> None:
        """
        Export user vectors from all adapters to disk for later analysis.

        Saves both .npz (efficient numpy format) and .json (human-readable).

        Args:
            method: Method name for the output directory
            adapters: Dict mapping profile_idx to adapter instances
        """
        method_dir = self.output_dir / method

        # Collect all user vectors from adapters
        all_vectors = {}
        for profile_idx, adapter in adapters.items():
            if hasattr(adapter, 'export_all_user_vectors'):
                vectors = adapter.export_all_user_vectors()
                all_vectors.update(vectors)

        if not all_vectors:
            logger.info(f"  No user vectors to export for {method}")
            return

        # Save as .npz for efficient analysis
        npz_path = method_dir / "user_vectors.npz"
        user_ids = list(all_vectors.keys())
        k = len(all_vectors[user_ids[0]]["z_long"])
        z_long = np.zeros((len(user_ids), k), dtype=np.float32)
        z_short = np.zeros((len(user_ids), k), dtype=np.float32)
        reward_ma = np.zeros(len(user_ids), dtype=np.float32)

        for i, uid in enumerate(user_ids):
            z_long[i] = all_vectors[uid]["z_long"]
            z_short[i] = all_vectors[uid]["z_short"]
            reward_ma[i] = all_vectors[uid]["reward_ma"]

        np.savez(
            npz_path,
            user_ids=np.array(user_ids),
            z_long=z_long,
            z_short=z_short,
            reward_ma=reward_ma,
        )

        # Also save summary stats as JSON
        summary = {
            "n_users": len(user_ids),
            "vector_dim": k,
            "z_long_norms": {uid: all_vectors[uid]["z_long_norm"] for uid in user_ids},
            "z_short_norms": {uid: all_vectors[uid]["z_short_norm"] for uid in user_ids},
            "reward_mas": {uid: all_vectors[uid]["reward_ma"] for uid in user_ids},
            "stats": {
                "z_long_norm_mean": float(np.mean([all_vectors[uid]["z_long_norm"] for uid in user_ids])),
                "z_long_norm_max": float(np.max([all_vectors[uid]["z_long_norm"] for uid in user_ids])),
                "z_long_norm_std": float(np.std([all_vectors[uid]["z_long_norm"] for uid in user_ids])),
            }
        }
        with open(method_dir / "user_vectors_summary.json", "w") as f:
            json.dump(summary, f, indent=2)

        logger.info(f"  Exported {len(user_ids)} user vectors to {npz_path}")

    def run_single_session(
        self,
        method: str,
        profile: Dict,
        problem: Dict,
        is_conflict_query: bool = False,
        adapter: Any = None,
        user_agent: Any = None
    ) -> Dict:
        """Run a single session with PROPER multi-turn conversation and user simulation.

        This implements:
        1. User simulator that role-plays with preferences
        2. Multi-turn conversation (up to max_turns)
        3. Preference enforcement by simulated user
        4. Proper metrics extraction from conversation
        """
        # Use provided adapter (reused across sessions) or create new one
        agent_adapter = adapter if adapter else self._create_method_adapter(method, profile)

        # Prepare conflict scenario if needed
        conflict_scenario = None
        original_problem = problem.get("problem", problem.get("question", ""))
        if is_conflict_query:
            conflict_scenario = self.conflict_generator.generate_for_profile(
                profile["preferences"],
                problem.get("domain", "general")
            )
            if conflict_scenario:
                problem = dict(problem)
                problem["problem"] = conflict_scenario["query"]

        query = problem.get("problem", problem.get("question", ""))

        # Extract user preferences as formatted string
        user_prefs = profile.get("preferences", [])
        if isinstance(user_prefs, list) and len(user_prefs) > 0:
            if isinstance(user_prefs[0], dict):
                # Structured preferences with condition/action
                pref_str = "\n".join([
                    f"- When {p.get('condition', '')}, {p.get('action', '')}"
                    for p in user_prefs
                ])
            else:
                # Simple string preferences
                pref_str = "\n".join([f"- {p}" for p in user_prefs])
        else:
            pref_str = str(user_prefs)

        user_persona = profile.get("persona", "A user seeking help with problem solving.")

        # Create user agent for this session (or reuse provided one)
        if user_agent is None:
            if self.config.use_openai_user:
                user_agent = OpenAIUserAgent(
                    user_task_description="Help the user solve their problem.",
                    problem=query,
                    user_persona=user_persona,
                    user_preferences=pref_str,
                    model=self.config.openai_user_model,
                )
            elif self.config.use_vllm:
                user_agent = VLLMUserAgent(
                    user_task_description="Help the user solve their problem.",
                    problem=query,
                    user_persona=user_persona,
                    user_preferences=pref_str,
                    vllm_url=self.config.vllm_user_url,
                )
            else:
                user_agent = SharedLocalUserAgent(
                    user_task_description="Help the user solve their problem.",
                    problem=query,
                    user_persona=user_persona,
                    user_preferences=pref_str,
                )

        # Initialize conversation
        turns = []
        full_user_log = []  # Detailed user agent outputs

        # Metrics tracking
        enforcement_count = 0
        disappointment_count = 0
        user_token_count = 0
        agent_token_count = 0
        preference_compliance_scores = []

        try:
            # Initialize adapter for this user
            if hasattr(agent_adapter, 'initialize'):
                agent_adapter.initialize()
            if hasattr(agent_adapter, 'start_session'):
                agent_adapter.start_session(
                    user_id=profile.get("user_id", "test_user"),
                    user_profile={"preferences": user_prefs, "persona": user_persona}
                )

            # Start with agent greeting
            conversation = [{"role": "assistant", "content": "How can I help you today?"}]

            # Multi-turn conversation loop
            for turn_num in range(self.config.max_turns_per_session):
                # === User Turn ===
                user_response = user_agent.generate_user_response(conversation)

                if user_response is None:
                    logger.warning(f"User agent failed to respond at turn {turn_num}")
                    break

                user_message = str(user_response.get("response", ""))
                user_token_count += len(user_message.split())

                # Add to conversation
                conversation.append({"role": "user", "content": user_message})
                turns.append({"role": "user", "content": user_message})
                full_user_log.append(user_response)

                # Check for termination
                if user_response.get("should_terminate", False) or TERMINATION_SIGNAL in user_message:
                    break

                # Detect preference enforcement (user correcting agent)
                enforcement_keywords = ["please", "i asked", "i said", "i prefer", "can you", "could you", "instead"]
                if any(kw in user_message.lower() for kw in enforcement_keywords):
                    enforcement_count += 1

                # === Agent Turn ===
                if hasattr(agent_adapter, 'generate_response'):
                    response = agent_adapter.generate_response(user_message, conversation[:-1])
                    agent_content = response.get("response", str(response)) if isinstance(response, dict) else str(response)
                elif callable(agent_adapter):
                    agent_content = agent_adapter(conversation)
                else:
                    agent_content = "[Error: Adapter not properly configured]"

                agent_token_count += len(agent_content.split())

                # Add to conversation
                conversation.append({"role": "assistant", "content": agent_content})
                turns.append({"role": "assistant", "content": agent_content})

                # Estimate preference compliance for this turn (heuristic based on user satisfaction)
                # If user doesn't enforce in next turn, assume compliance
                # This is a simplified heuristic - LLM judge would be more accurate
                compliance_score = 0.8 if enforcement_count == 0 else max(0.2, 1.0 - 0.2 * enforcement_count)
                preference_compliance_scores.append(compliance_score)

            # End session
            if hasattr(agent_adapter, 'end_session'):
                adapter_metrics = agent_adapter.end_session(task_success=True)
            else:
                adapter_metrics = {}

        except Exception as e:
            import traceback
            logger.error(f"Error in session: {e}")
            logger.error(f"Full traceback:\n{traceback.format_exc()}")
            turns.append({"role": "assistant", "content": f"[Error: {e}]"})

        # Compute metrics
        total_turns = len(turns)
        total_token_count = user_token_count + agent_token_count

        # Check if user reached a satisfactory answer (from last user response)
        task_success = False
        if full_user_log:
            last_user = full_user_log[-1]
            if last_user.get("should_terminate", False):
                draft = last_user.get("draft_answer", "")
                # Consider success if draft answer is not empty/"I don't know"
                task_success = bool(draft) and draft.lower() != "i don't know"

        # Compute average compliance
        avg_compliance = sum(preference_compliance_scores) / len(preference_compliance_scores) if preference_compliance_scores else 0.5

        # Conflict resolution (if this was a conflict test)
        conflict_accuracy = 0.0
        if is_conflict_query and conflict_scenario:
            # Check if the correct preference was applied
            expected_pref = conflict_scenario.get("expected_preference", "")
            # Simple heuristic: check if expected preference keywords appear in agent responses
            agent_texts = " ".join([t["content"] for t in turns if t["role"] == "assistant"])
            if expected_pref and any(kw in agent_texts.lower() for kw in expected_pref.lower().split()[:3]):
                conflict_accuracy = 1.0

        # Over-personalization detection (heuristic: if agent mentions preferences not in profile)
        over_personalization = 0.0

        metrics = ConversationMetrics(
            task_success=task_success,
            turns_to_success=total_turns if task_success else -1,
            total_turns=total_turns,
            user_token_count=user_token_count,
            enforcement_count=enforcement_count,
            disappointment_count=disappointment_count,
            total_token_count=total_token_count,
            agent_token_count=agent_token_count,
            preference_compliance_scores=preference_compliance_scores,
            conflict_resolution_accuracy=conflict_accuracy,
            over_personalization_rate=over_personalization,
        )

        return {
            "method": method,
            "profile_id": profile.get("user_id", "unknown"),
            "problem_id": problem.get("problem_id", str(hash(query))[:8]),
            "problem": original_problem,
            "ground_truth_solution": problem.get("solution", problem.get("answer", "")),
            "is_conflict_test": is_conflict_query,
            "conflict_scenario": conflict_scenario,
            "conversation": {"turns": turns} if self.config.save_conversations else None,
            "full_user_log": full_user_log if self.config.save_conversations else None,
            "metrics": asdict(metrics),
            "adapter_metrics": adapter_metrics if 'adapter_metrics' in dir() else {},
        }

    def _run_profile_sessions(
        self,
        method: str,
        profile_idx: int,
        profile: Dict,
        adapter: Any = None
    ) -> List[Dict]:
        """Run all sessions for a single profile. Thread-safe for parallel execution."""
        profile_results = []

        # Create vLLM-based agent client if using vLLM (for methods that need it)
        vllm_agent = None
        if self.config.use_vllm and method == "vanilla":
            vllm_agent = VLLMAgentClient(
                vllm_url=self.config.vllm_agent_url,
                system_prompt="You are a helpful AI assistant for problem-solving tasks."
            )

        # Run sessions across datasets
        session_idx = 0
        for ds_name, dataset in self.datasets.items():
            samples = dataset.get_testset()

            for sample in samples:
                if session_idx >= self.config.n_sessions_per_profile:
                    break

                # Decide if this is a conflict query
                is_conflict = (session_idx % int(1 / self.config.conflict_ratio)) == 0

                problem = {
                    "problem": sample.problem,
                    "solution": sample.solution,
                    "problem_id": sample.problem_id,
                    "domain": sample.domain,
                }

                try:
                    result = self.run_single_session(
                        method=method,
                        profile=profile,
                        problem=problem,
                        is_conflict_query=is_conflict,
                        adapter=vllm_agent if vllm_agent else adapter
                    )
                    profile_results.append(result)
                except Exception as e:
                    logger.error(f"Error in session for profile {profile_idx}: {e}")

                session_idx += 1

        return profile_results

    def run_method(self, method: str) -> List[Dict]:
        """Run all sessions for a single method with checkpointing and parallel processing."""
        logger.info(f"Running method: {method}")

        # Setup method directory and checkpoint
        method_dir = self.output_dir / method
        method_dir.mkdir(exist_ok=True)
        checkpoint_file = method_dir / "checkpoint.json"
        results_file = method_dir / "results.json"

        # Load existing results and checkpoint
        results = []
        completed_profiles = set()
        sessions_per_profile = {}  # Track session count per profile for continue functionality
        if checkpoint_file.exists():
            with open(checkpoint_file, "r") as f:
                checkpoint = json.load(f)
                completed_profiles = set(checkpoint.get("completed_profiles", []))
                sessions_per_profile = checkpoint.get("sessions_per_profile", {})
                logger.info(f"  Resuming from checkpoint: {len(completed_profiles)} profiles completed")
                if sessions_per_profile:
                    total_sessions = sum(sessions_per_profile.values())
                    logger.info(f"  Session-level tracking: {total_sessions} sessions across {len(sessions_per_profile)} profiles")
        if results_file.exists():
            with open(results_file, "r") as f:
                results = json.load(f)

        # Determine profile range
        start_idx = self.config.start_profile
        end_idx = self.config.end_profile if self.config.end_profile else len(self.profiles)

        # Build list of profiles that need more sessions
        profiles_to_run = []
        for idx in range(start_idx, min(end_idx, len(self.profiles))):
            existing_sessions = sessions_per_profile.get(str(idx), 0)
            if existing_sessions < self.config.n_sessions_per_profile:
                profiles_to_run.append(idx)

        # Log what we're running
        if sessions_per_profile:
            total_existing = sum(sessions_per_profile.get(str(idx), 0) for idx in profiles_to_run)
            total_needed = len(profiles_to_run) * self.config.n_sessions_per_profile
            logger.info(f"  Running profiles {start_idx} to {end_idx-1}: {len(profiles_to_run)} profiles need sessions")
            logger.info(f"  Sessions: {total_existing} existing, {total_needed - total_existing} remaining")
        else:
            logger.info(f"  Running profiles {start_idx} to {end_idx-1} ({len(profiles_to_run)} remaining)")

        # When using batch processing with vLLM or OpenAI user: use turn-synchronous batch mode
        # This batches both user and agent calls for maximum throughput
        if self.config.use_batch_processing and self.config.use_vllm:
            user_type = "OpenAI" if self.config.use_openai_user else "local vLLM"
            logger.info(f"  Using BATCH processing ({user_type} user) for {method}")
            return self._run_method_batch(
                method, profiles_to_run, results, completed_profiles,
                sessions_per_profile, checkpoint_file, results_file
            )

        # Decide on parallelization for sequential methods
        n_parallel = self.config.parallel_profiles if (self.config.use_vllm or self.config.use_openai_user) else 1

        if n_parallel > 1:
            logger.info(f"  Using parallel processing with {n_parallel} workers")
            self._run_method_parallel(
                method, profiles_to_run, results, completed_profiles,
                sessions_per_profile, checkpoint_file, results_file
            )
        else:
            # Sequential execution (original behavior)
            # Create ONE adapter per method and reuse it (avoids GPU OOM from repeated model loading)
            adapter = self._create_method_adapter(method, None)
            adapter.initialize()

            for profile_idx in profiles_to_run:
                profile = self.profiles[profile_idx]
                logger.info(f"  Profile {profile_idx + 1}/{len(self.profiles)}")

                profile_results = self._run_profile_sessions(method, profile_idx, profile, adapter)

                # Add profile results to overall results
                results.extend(profile_results)
                completed_profiles.add(profile_idx)
                sessions_per_profile[str(profile_idx)] = self.config.n_sessions_per_profile

                # Save checkpoint and results after each profile
                with open(checkpoint_file, "w") as f:
                    json.dump({
                        "completed_profiles": sorted(list(completed_profiles)),
                        "sessions_per_profile": sessions_per_profile
                    }, f)
                with open(results_file, "w") as f:
                    json.dump(results, f, indent=2)
                logger.info(f"  Profile {profile_idx + 1} completed and checkpointed")

            # Export user vectors at the end of sequential processing
            self._export_user_vectors(method, {0: adapter})

        return results

    def _run_method_parallel(
        self,
        method: str,
        profiles_to_run: List[int],
        results: List[Dict],
        completed_profiles: set,
        sessions_per_profile: Dict[str, int],
        checkpoint_file: Path,
        results_file: Path
    ):
        """Run profiles in parallel using ThreadPoolExecutor.

        Uses shared model singletons for embedding/reranker to avoid OOM
        when multiple workers try to load their own copies.
        """
        n_parallel = self.config.parallel_profiles
        results_lock = threading.Lock()
        start_time = time.time()
        profiles_completed = 0

        def process_profile(profile_idx: int) -> tuple:
            """Process a single profile and return (profile_idx, results)."""
            profile = self.profiles[profile_idx]
            # Create adapter with shared models to avoid OOM from duplicate model loading
            adapter = self._create_method_adapter(method, profile, use_shared_models=True)
            profile_results = self._run_profile_sessions(method, profile_idx, profile, adapter)
            return profile_idx, profile_results

        with ThreadPoolExecutor(max_workers=n_parallel) as executor:
            # Submit all profile jobs
            future_to_profile = {
                executor.submit(process_profile, idx): idx
                for idx in profiles_to_run
            }

            # Process completed profiles
            for future in as_completed(future_to_profile):
                profile_idx = future_to_profile[future]
                try:
                    idx, profile_results = future.result()

                    with results_lock:
                        results.extend(profile_results)
                        completed_profiles.add(idx)
                        sessions_per_profile[str(idx)] = self.config.n_sessions_per_profile
                        profiles_completed += 1

                        # Save checkpoint with session-level tracking
                        with open(checkpoint_file, "w") as f:
                            json.dump({
                                "completed_profiles": sorted(list(completed_profiles)),
                                "sessions_per_profile": sessions_per_profile
                            }, f)
                        with open(results_file, "w") as f:
                            json.dump(results, f, indent=2)

                        # Log progress with throughput estimate
                        elapsed = time.time() - start_time
                        profiles_per_hour = profiles_completed / elapsed * 3600 if elapsed > 0 else 0
                        sessions_per_hour = len(results) / elapsed * 3600 if elapsed > 0 else 0
                        logger.info(
                            f"  Profile {idx + 1} completed "
                            f"({profiles_completed}/{len(profiles_to_run)}) - "
                            f"{profiles_per_hour:.1f} profiles/hr, {sessions_per_hour:.1f} sessions/hr"
                        )

                except Exception as e:
                    logger.error(f"  Profile {profile_idx} failed: {e}")

        # Note: Parallel mode doesn't export user vectors because adapters are
        # created/destroyed per profile. Use batch mode for vector export.
        logger.info(f"  Parallel mode: user vectors not exported (use batch mode)")

    def _run_method_batch(
        self,
        method: str,
        profiles_to_run: List[int],
        results: List[Dict],
        completed_profiles: set,
        sessions_per_profile: Dict[str, int],
        checkpoint_file: Path,
        results_file: Path
    ) -> List[Dict]:
        """
        Turn-synchronous batch processing for ALL methods.

        At each turn, user calls are batched concurrently via AsyncOpenAI,
        then agent responses go through personalization adapters.
        Sessions within a profile run sequentially (for stateful memory).
        """
        from agents.batch_vllm_agent import BatchOpenAIClient, BatchVLLMClient, TERMINATION_SIGNAL
        from json_repair import repair_json

        start_time = time.time()

        # Create user client (OpenAI API or local vLLM)
        if self.config.use_openai_user:
            user_client = BatchOpenAIClient(
                model=self.config.openai_user_model,
                max_tokens=4096,
                max_concurrent=32,
                api_key=os.environ.get("OPENAI_API_KEY"),
            )
            logger.info(f"  Using OpenAI user simulator: {self.config.openai_user_model}")
        else:
            user_client = BatchVLLMClient(
                vllm_url=self.config.vllm_user_url,
                max_tokens=1024,  # User responses typically short, but allow for edge cases
                temperature=1.0,
                timeout=None,
                max_concurrent=100,
                json_mode=True,  # User simulator needs JSON output
            )
            logger.info(f"  Using local vLLM user simulator: {self.config.vllm_user_url}")

        # Create async agent client for batched vLLM calls
        agent_client = BatchVLLMClient(
            vllm_url=self.config.vllm_agent_url,
            max_tokens=2048,
            temperature=0.7,
            timeout=None,  # Infinite timeout for long generations
            max_concurrent=100,
        )

        USER_PROMPT_TEMPLATE = (
            "You are a user simulator collaborating with an agent to solve a problem. "
            "You will be provided with a problem description, and you must get the agent to help you solve it. "
            "You will also be provided with user preferences, which you must follow and actively enforce throughout the conversation.\n\n"
            "# Problem Description\n{problem}\nNote: the agent cannot see this problem description.\n\n"
            "# User Persona\n{user_persona}\n\n"
            "# User Preferences\n{user_preferences}\n"
            "These preferences are NON-NEGOTIABLE that define how you prefer the agent to behave. They must be strictly enforced:\n"
            "   - **Answer clarifying questions**: The agent may ask clarifying questions before attempting an answer. "
            "Answer such questions, and do not enforce preferences about answer format or content while the agent is clarifying.\n"
            "   - **Enforce immediately**: Every agent response must satisfy your preferences before you can proceed. "
            "Explicitly ask the agent to adjust their response until it complies.\n"
            "   - **Never proceed without compliance**: Do NOT update your draft answer, do NOT consider terminating, "
            "and do NOT move forward until the agent follows your preferences.\n\n"
            "# Draft Answer Management\n"
            "- **Maintain a working draft**: Start with \"I don't know\". Update your draft answer based on what you learn from agent responses.\n"
            "- **Don't update when enforcing preferences**: If the agent response does not follow your preferences, "
            "do NOT update your draft answer, regardless of whether the agent provides helpful information.\n\n"
            "# Conversation Termination\n"
            "Before generating your response, determine if you should terminate:\n"
            "   - Do you feel like your draft answer is a good answer to the problem?\n"
            "   - Do you feel like the agent cannot help further?\n"
            "If the agent response does not follow your preferences, you must NOT terminate - instead, enforce the preferences.\n"
            "When ready to terminate, respond with \"TERMINATE\".\n\n"
            "# Output Format (respond in JSON):\n"
            "{{\n"
            "   \"preferences_check\": \"For EACH relevant preference, evaluate: is it satisfied?\",\n"
            "   \"enforce_preferences\": true/false,\n"
            "   \"reasoning\": \"Brief reasoning (2-3 sentences). Does agent follow preferences? If no, enforce. If yes, update draft.\",\n"
            "   \"draft_answer\": \"Your current working draft answer\",\n"
            "   \"should_terminate\": true/false,\n"
            "   \"response\": \"Your response to the agent\"\n"
            "}}"
        )

        def parse_user_response(content):
            if not content:
                return None
            try:
                parsed = repair_json(content, return_objects=True)
                if isinstance(parsed, dict) and "response" in parsed:
                    return parsed
            except:
                pass
            if TERMINATION_SIGNAL in (content or ""):
                return {"reasoning": "", "draft_answer": "", "should_terminate": True, "response": TERMINATION_SIGNAL}
            return {"reasoning": "", "draft_answer": "", "should_terminate": False, "response": content or ""}

        def reverse_roles(conversation):
            return [
                {"role": "user" if m["role"] == "assistant" else "assistant", "content": m["content"]}
                for m in conversation
            ]

        # Create per-profile adapters
        adapters = {}
        profile_sessions = {}

        # Build session problem list ONCE (shared across all profiles for controlled comparison)
        # Each dataset contributes exactly n_per_dataset problems (front 10), no repeats
        shared_sessions = []
        dataset_names = list(self.datasets.keys())
        n_per_dataset = self.config.n_sessions_per_profile // len(dataset_names)
        remainder = self.config.n_sessions_per_profile % len(dataset_names)

        for i, ds_name in enumerate(dataset_names):
            ds_obj = self.datasets[ds_name]
            items = ds_obj.get_testset()
            n_take = n_per_dataset + (1 if i < remainder else 0)
            if n_take > len(items):
                logger.warning(f"  Dataset {ds_name} has only {len(items)} problems, need {n_take}")
            for j in range(n_take):
                item = items[j % len(items)]
                shared_sessions.append({"problem": item.problem, "solution": item.solution, "domain": ds_obj.domain})

        n_conflict = int(len(shared_sessions) * self.config.conflict_ratio)
        shared_session_list = [(s, idx < n_conflict) for idx, s in enumerate(shared_sessions)]
        logger.info(f"  Built shared session list: {len(shared_sessions)} problems from {len(dataset_names)} datasets ({n_per_dataset} each, same for all profiles)")

        for profile_idx in profiles_to_run:
            profile = self.profiles[profile_idx]
            adapter = self._create_method_adapter(method, profile, use_shared_models=True)
            if hasattr(adapter, 'initialize'):
                adapter.initialize()
            adapters[profile_idx] = adapter
            profile_sessions[profile_idx] = shared_session_list

        n_sessions = self.config.n_sessions_per_profile

        # Calculate sessions to run per profile (accounting for existing sessions)
        sessions_to_run_per_profile = {}
        for profile_idx in profiles_to_run:
            existing = sessions_per_profile.get(str(profile_idx), 0)
            remaining = n_sessions - existing
            if remaining > 0:
                sessions_to_run_per_profile[profile_idx] = (existing, remaining)  # (start_session, count)

        if sessions_to_run_per_profile:
            total_remaining = sum(v[1] for v in sessions_to_run_per_profile.values())
            logger.info(f"  Batch: {len(sessions_to_run_per_profile)} profiles, {total_remaining} sessions remaining")
        else:
            logger.info(f"  Batch: All sessions already completed")
            return results

        # Process sessions in rounds
        for session_idx in range(n_sessions):
            # Initialize all conversations for this round
            all_states = {}  # profile_idx -> state dict
            active_set = set()

            for profile_idx in profiles_to_run:
                # Skip if this profile doesn't need this session
                if profile_idx not in sessions_to_run_per_profile:
                    continue
                start_session, _ = sessions_to_run_per_profile[profile_idx]
                if session_idx < start_session:
                    continue  # Already completed this session
                if session_idx >= len(profile_sessions[profile_idx]):
                    continue
                problem_dict, is_conflict = profile_sessions[profile_idx][session_idx]
                profile = self.profiles[profile_idx]
                query = problem_dict["problem"]

                if is_conflict:
                    cs = self.conflict_generator.generate_for_profile(
                        profile.get("preferences", []), problem_dict.get("domain", "general"))
                    if cs:
                        query = cs["query"]

                user_prefs = profile.get("preferences", [])
                if isinstance(user_prefs, list) and user_prefs:
                    if isinstance(user_prefs[0], dict):
                        pref_str = "\n".join([f"- When {p.get('condition','')}, {p.get('action','')}" for p in user_prefs])
                    else:
                        pref_str = "\n".join([f"- {p}" for p in user_prefs])
                else:
                    pref_str = str(user_prefs)

                user_persona = profile.get("persona", "A user seeking help with problem solving.")
                adapter = adapters[profile_idx]
                if hasattr(adapter, 'start_session'):
                    adapter.start_session(
                        user_id=profile.get("user_id", f"user_{profile_idx}"),
                        user_profile={"preferences": user_prefs, "persona": user_persona}
                    )

                all_states[profile_idx] = {
                    "conversation": [{"role": "assistant", "content": "How can I help you today?"}],
                    "full_log": [],
                    "system_prompt": USER_PROMPT_TEMPLATE.format(
                        problem=query, user_persona=user_persona, user_preferences=pref_str),
                    "problem_dict": problem_dict,
                    "is_conflict": is_conflict,
                    "enforcement_count": 0,
                }
                active_set.add(profile_idx)

            # Turn-synchronous loop
            for turn in range(self.config.max_turns_per_session):
                if not active_set:
                    break

                # Batch user calls
                active_list = sorted(active_set)
                user_msgs_batch = []
                for pidx in active_list:
                    state = all_states[pidx]
                    msgs = [{"role": "system", "content": state["system_prompt"]}]
                    msgs.extend(reverse_roles(state["conversation"]))
                    user_msgs_batch.append(msgs)

                user_responses = user_client.batch_completion(user_msgs_batch)

                # Process user responses and prepare agent prompts for batching
                to_remove = []
                agent_prompts_batch = []  # List of (pidx, messages, context)
                for i, pidx in enumerate(active_list):
                    state = all_states[pidx]
                    parsed = parse_user_response(user_responses[i])

                    if parsed is None:
                        to_remove.append(pidx)
                        continue

                    user_msg = str(parsed.get("response", ""))
                    state["conversation"].append({"role": "user", "content": user_msg})
                    state["full_log"].append(parsed)

                    enforce = parsed.get("enforce_preferences", False)
                    if isinstance(enforce, str):
                        enforce = enforce.lower() == "true"
                    if enforce:
                        state["enforcement_count"] += 1

                    # Detect disappointment and satisfaction from user message
                    # Disappointment indicators (not quite right, could be better, etc.)
                    user_msg_lower = user_msg.lower()
                    disappointment = any(phrase in user_msg_lower for phrase in [
                        "not quite", "not what i", "that's not", "incorrect",
                        "wrong", "mistake", "error", "confused", "doesn't make sense",
                        "try again", "not helpful", "not useful"
                    ])
                    # Satisfaction indicators (explicit positive feedback)
                    satisfaction = parsed.get("should_terminate", False) or any(phrase in user_msg_lower for phrase in [
                        "perfect", "exactly", "great", "thanks", "helpful",
                        "that's right", "correct", "good job", "well done",
                        "makes sense", "understand now", "got it"
                    ])

                    # Store parsed feedback for REINFORCE (applied AFTER prepare_prompt sets pending_rl_update)
                    state["_pending_feedback"] = {
                        "user_msg": user_msg,
                        "enforce": bool(enforce),
                        "disappointment": disappointment and not enforce,  # Don't double-count
                        "satisfaction": satisfaction and not enforce,  # Don't count if also enforcing
                        "draft_answer": bool(parsed.get("draft_answer")),
                    }

                    if parsed.get("should_terminate", False) or TERMINATION_SIGNAL in user_msg:
                        to_remove.append(pidx)
                        continue

                # Batch preference extraction for PersonalizedLLM adapters
                extraction_batch = []  # (pidx, query)
                remaining_active = [pidx for pidx in active_list if pidx not in to_remove]
                for pidx in remaining_active:
                    adapter = adapters.get(pidx)
                    if adapter and hasattr(adapter, '_llm') and hasattr(adapter._llm, 'enable_preference_extraction'):
                        if adapter._llm.enable_preference_extraction and adapter._llm._extractor is not None:
                            query = adapter._llm.get_last_user_query(adapter._current_user_id) if hasattr(adapter._llm, 'get_last_user_query') else None
                            if not query:
                                state = all_states[pidx]
                                query = state["conversation"][-1]["content"] if state["conversation"] else ""
                            if query:
                                extraction_batch.append((pidx, query))

                if extraction_batch:
                    extractor = extraction_batch[0][1]  # just need any adapter to get the extractor
                    adapter0 = adapters[extraction_batch[0][0]]
                    shared_extractor = adapter0._llm._extractor
                    if hasattr(shared_extractor, 'batch_extract_preferences'):
                        queries = [q for _, q in extraction_batch]
                        batch_results = shared_extractor.batch_extract_preferences(queries)
                        for (pidx, _), pref_dict in zip(extraction_batch, batch_results):
                            adapter = adapters[pidx]
                            adapter._llm.apply_extracted_preferences(adapter._current_user_id, pref_dict)
                    else:
                        # Fallback: sequential
                        for pidx, query in extraction_batch:
                            adapter = adapters[pidx]
                            adapter._llm._extractor.extract_turn(adapter._llm._sessions[adapter._current_user_id].session_state.history)

                # Batch scaffolding for reflection adapters before prepare_prompt
                scaffolding_batch = []  # (pidx, prompt)
                remaining_active = [pidx for pidx in active_list if pidx not in to_remove]
                for pidx in remaining_active:
                    adapter = adapters.get(pidx)
                    if adapter and hasattr(adapter, 'get_scaffolding_prompt'):
                        state = all_states[pidx]
                        # Temporarily add user msg to history for scaffolding
                        agent_notes = adapter._user_notes.get(adapter._current_user_id, "No notes yet about this user.")
                        if adapter.with_scaffolding and agent_notes != "No notes yet about this user.":
                            prompt = adapter.get_scaffolding_prompt(
                                state["conversation"], agent_notes)
                            if prompt is not None:
                                scaffolding_batch.append((pidx, prompt))

                if scaffolding_batch:
                    scaff_messages = [[{"role": "user", "content": p}] for _, p in scaffolding_batch]
                    scaff_responses = agent_client.batch_completion(scaff_messages)
                    for (pidx, _), resp in zip(scaffolding_batch, scaff_responses):
                        adapter = adapters[pidx]
                        adapter._scaffolding_result = resp if resp else None

                # Prepare agent prompts for batching
                # NOTE: prepare_prompt calls chat_prepare which sets pending_rl_update
                # from the previous turn's data. REINFORCE feedback must be applied
                # AFTER this call so that pending_rl_update is available.
                for pidx in remaining_active:
                    state = all_states[pidx]
                    try:
                        adapter = adapters[pidx]
                        user_msg = state["conversation"][-1]["content"]
                        if hasattr(adapter, 'prepare_prompt'):
                            messages, context = adapter.prepare_prompt(user_msg, state["conversation"][:-1])
                            agent_prompts_batch.append((pidx, messages, context))
                        elif hasattr(adapter, 'generate_response'):
                            agent_prompts_batch.append((pidx, None, None))
                        else:
                            state["conversation"].append({"role": "assistant", "content": "[Error: Adapter not configured]"})
                    except Exception as e:
                        logger.error(f"  Agent prepare error p{pidx} t{turn}: {e}")
                        state["conversation"].append({"role": "assistant", "content": "I apologize, I encountered an error. Could you rephrase?"})

                # Apply REINFORCE feedback NOW (after prepare_prompt set pending_rl_update)
                for pidx in remaining_active:
                    state = all_states[pidx]
                    fb = state.pop("_pending_feedback", None)
                    if fb:
                        adapter = adapters.get(pidx)
                        if adapter and hasattr(adapter, 'process_user_turn'):
                            adapter.process_user_turn(
                                user_response=fb["user_msg"],
                                enforce_preferences=fb["enforce"],
                                express_disappointment=fb.get("disappointment", False),
                                express_satisfaction=fb["satisfaction"],
                                draft_answer_updated=fb["draft_answer"],
                            )

                # Also apply feedback for terminated sessions (they skipped prepare_prompt
                # but still need the reward signal from their last turn)
                for pidx in to_remove:
                    state = all_states.get(pidx)
                    if not state:
                        continue
                    fb = state.pop("_pending_feedback", None)
                    if fb:
                        adapter = adapters.get(pidx)
                        if adapter and hasattr(adapter, 'process_user_turn'):
                            # For terminated sessions, we can't call prepare_prompt
                            # (no next turn), but we still want the reward applied.
                            # Call chat_prepare with a dummy to set pending_rl_update,
                            # then apply feedback.
                            try:
                                if hasattr(adapter, '_llm') and hasattr(adapter._llm, 'chat_prepare'):
                                    adapter._llm.chat_prepare(
                                        adapter._current_user_id,
                                        fb["user_msg"],
                                        skip_extraction=True,
                                        skip_auto_reward=True,
                                    )
                                adapter.process_user_turn(
                                    user_response=fb["user_msg"],
                                    enforce_preferences=fb["enforce"],
                                    express_disappointment=fb.get("disappointment", False),
                                    express_satisfaction=fb["satisfaction"],
                                    draft_answer_updated=fb["draft_answer"],
                                )
                            except Exception:
                                pass  # Best effort for terminated sessions

                # Batch vLLM call for all agent prompts
                if agent_prompts_batch:
                    # Separate prompts that can be batched from fallback
                    batchable = [(pidx, msgs, ctx) for pidx, msgs, ctx in agent_prompts_batch if msgs is not None]
                    fallback = [(pidx, msgs, ctx) for pidx, msgs, ctx in agent_prompts_batch if msgs is None]

                    # Batch call for batchable prompts
                    if batchable:
                        batch_messages = [msgs for _, msgs, _ in batchable]
                        batch_responses = agent_client.batch_completion(batch_messages)

                        # Process batched responses
                        for (pidx, _, context), response in zip(batchable, batch_responses):
                            try:
                                adapter = adapters[pidx]
                                state = all_states[pidx]
                                if response is not None:
                                    result = adapter.process_response(response, context)
                                    agent_content = result.get("response", str(result)) if isinstance(result, dict) else str(result)
                                else:
                                    agent_content = "I apologize, I encountered an error. Could you rephrase?"
                                state["conversation"].append({"role": "assistant", "content": agent_content})
                            except Exception as e:
                                logger.error(f"  Agent process error p{pidx} t{turn}: {e}")
                                all_states[pidx]["conversation"].append({"role": "assistant", "content": "I apologize, I encountered an error. Could you rephrase?"})

                    # Handle fallback (adapters without prepare_prompt - sequential calls)
                    for pidx, _, _ in fallback:
                        try:
                            adapter = adapters[pidx]
                            state = all_states[pidx]
                            user_msg = state["conversation"][-1]["content"]
                            resp = adapter.generate_response(user_msg, state["conversation"][:-1])
                            agent_content = resp.get("response", str(resp)) if isinstance(resp, dict) else str(resp)
                            state["conversation"].append({"role": "assistant", "content": agent_content})
                        except Exception as e:
                            logger.error(f"  Agent fallback error p{pidx} t{turn}: {e}")
                            all_states[pidx]["conversation"].append({"role": "assistant", "content": "I apologize, I encountered an error. Could you rephrase?"})

                active_set -= set(to_remove)

            # Batch note-update for reflection adapters before end_session
            note_update_batch = []  # (profile_idx, messages)
            for profile_idx in profiles_to_run:
                if profile_idx not in all_states:
                    continue
                adapter = adapters.get(profile_idx)
                if adapter and hasattr(adapter, 'get_note_update_prompt'):
                    prompt_msgs = adapter.get_note_update_prompt()
                    if prompt_msgs is not None:
                        note_update_batch.append((profile_idx, prompt_msgs))

            if note_update_batch:
                note_messages = [msgs for _, msgs in note_update_batch]
                note_responses = agent_client.batch_completion(note_messages)
                for (profile_idx, _), resp in zip(note_update_batch, note_responses):
                    if resp:
                        adapter = adapters[profile_idx]
                        adapter.apply_note_update_response(resp)

            # Save results for this session round
            for profile_idx in profiles_to_run:
                if profile_idx not in all_states:
                    continue
                state = all_states[profile_idx]
                problem_dict = state["problem_dict"]
                conversation = state["conversation"]
                full_log = state["full_log"]

                user_tokens = sum(len(m["content"].split()) for m in conversation if m["role"] == "user")
                agent_tokens = sum(len(m["content"].split()) for m in conversation if m["role"] == "assistant")

                enforcement_count = state["enforcement_count"]
                task_success = 0
                for entry in full_log:
                    if entry.get("should_terminate", False):
                        draft = str(entry.get("draft_answer", ""))
                        if draft and "don't know" not in draft.lower():
                            task_success = 1

                # End session on adapter (applies task completion reward for REINFORCE)
                adapter = adapters.get(profile_idx)
                if adapter and hasattr(adapter, 'end_session'):
                    # Skip note update if batch already handled it
                    skip_notes = hasattr(adapter, 'get_note_update_prompt')
                    try:
                        adapter.end_session(task_success=bool(task_success), skip_note_update=skip_notes)
                    except TypeError:
                        adapter.end_session(task_success=bool(task_success))

                results.append({
                    "method": method,
                    "profile_id": self.profiles[profile_idx].get("user_id", f"user_{profile_idx}"),
                    "problem_id": f"session_{session_idx}",
                    "problem": problem_dict.get("problem", ""),
                    "ground_truth_solution": problem_dict.get("solution", ""),
                    "is_conflict_test": state["is_conflict"],
                    "conversation": {"turns": conversation},
                    "full_user_log": full_log,
                    "metrics": {
                        "task_success": bool(task_success),
                        "total_turns": len(conversation),
                        "user_token_count": user_tokens,
                        "agent_token_count": agent_tokens,
                        "total_token_count": user_tokens + agent_tokens,
                        "enforcement_count": enforcement_count,
                        "disappointment_count": 0,
                        "preference_compliance_scores": [],
                        "conflict_resolution_accuracy": 0,
                        "over_personalization_rate": 0,
                    },
                    "adapter_metrics": {},
                })

                # Collect adapter metrics (e.g. user_vector_norm for rag_vector)
                adapter = adapters.get(profile_idx)
                if adapter and hasattr(adapter, 'get_user_vector'):
                    user_id = self.profiles[profile_idx].get("user_id", f"user_{profile_idx}")
                    vec = adapter.get_user_vector(user_id)
                    if vec is not None:
                        results[-1]["adapter_metrics"] = {
                            "user_vector_norm": float(np.linalg.norm(vec)),
                        }

            # Save user vector snapshots every 10 sessions
            if (session_idx + 1) % 10 == 0:
                vectors_dir = checkpoint_file.parent / "vectors"
                vectors_dir.mkdir(parents=True, exist_ok=True)
                user_vectors = {}
                for profile_idx in profiles_to_run:
                    adapter = adapters.get(profile_idx)
                    if adapter and hasattr(adapter, 'get_user_vector'):
                        user_id = self.profiles[profile_idx].get("user_id", f"user_{profile_idx}")
                        vec = adapter.get_user_vector(user_id)
                        if vec is not None:
                            user_vectors[user_id] = vec
                if user_vectors:
                    snapshot_path = vectors_dir / f"vectors_session_{session_idx+1}.npy"
                    np.save(snapshot_path, user_vectors)
                    logger.info(f"  Saved {len(user_vectors)} user vectors to {snapshot_path}")

            # Checkpoint after each session round with session-level tracking
            # Only increment for profiles that actually ran in this round (those in all_states)
            for profile_idx in all_states.keys():
                sessions_per_profile[str(profile_idx)] = sessions_per_profile.get(str(profile_idx), 0) + 1
                if sessions_per_profile[str(profile_idx)] >= self.config.n_sessions_per_profile:
                    completed_profiles.add(profile_idx)

            with open(checkpoint_file, "w") as f:
                json.dump({
                    "completed_profiles": sorted(list(completed_profiles)),
                    "sessions_per_profile": sessions_per_profile
                }, f)
            with open(results_file, "w") as f:
                json.dump(results, f, indent=2)

            elapsed = time.time() - start_time
            sessions_done = len(results)
            rate = sessions_done / elapsed * 3600 if elapsed > 0 else 0
            logger.info(f"  Session round {session_idx+1}/{n_sessions}: {sessions_done} total, {rate:.0f} sessions/hr")

        # Export user vectors before cleanup (for RAG methods with user vectors)
        self._export_user_vectors(method, adapters)

        # Explicitly free adapter models to prevent GPU OOM across methods
        for pidx, adapter in adapters.items():
            if hasattr(adapter, 'cleanup'):
                adapter.cleanup()
        del adapters

        return results

    def run_all(self) -> Dict[str, Any]:
        """Run all methods and generate comparative analysis."""
        all_results = {}

        for method in self.config.methods:
            if method not in AVAILABLE_METHODS:
                logger.warning(f"Unknown method: {method}, skipping")
                continue

            results = self.run_method(method)
            all_results[method] = results

            # Free GPU memory between methods to prevent OOM on later adapters
            try:
                from personalization.serving.personalized_llm import clear_shared_models
                clear_shared_models()
            except ImportError:
                pass
            try:
                import gc
                import torch
                gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                    logger.info(f"  GPU memory freed after {method}: {torch.cuda.memory_allocated()/1e9:.1f}GB allocated")
            except ImportError:
                pass

        # Comparative analysis
        analysis = self._analyze_results(all_results)

        # Save analysis
        with open(self.output_dir / "analysis.json", "w") as f:
            json.dump(analysis, f, indent=2)

        # Generate report
        self._generate_report(analysis)

        return analysis

    def _analyze_results(self, all_results: Dict[str, List[Dict]]) -> Dict:
        """Analyze results across all methods."""
        analysis = {
            "per_method": {},
            "comparison": {},
        }

        for method, results in all_results.items():
            n = len(results)
            if n == 0:
                continue

            # Aggregate metrics
            task_success = sum(r["metrics"]["task_success"] for r in results) / n
            avg_user_tokens = sum(r["metrics"]["user_token_count"] for r in results) / n
            avg_total_tokens = sum(r["metrics"]["total_token_count"] for r in results) / n
            avg_enforcement = sum(r["metrics"]["enforcement_count"] for r in results) / n
            avg_turns = sum(r["metrics"]["total_turns"] for r in results) / n

            # Compliance and conflict metrics
            compliance_scores = [
                sum(r["metrics"]["preference_compliance_scores"]) / len(r["metrics"]["preference_compliance_scores"])
                if r["metrics"]["preference_compliance_scores"] else 0.5
                for r in results
            ]
            avg_compliance = sum(compliance_scores) / len(compliance_scores)

            conflict_results = [r for r in results if r["is_conflict_test"]]
            conflict_accuracy = sum(
                r["metrics"]["conflict_resolution_accuracy"] for r in conflict_results
            ) / len(conflict_results) if conflict_results else 0

            over_personalization = sum(
                r["metrics"]["over_personalization_rate"] for r in results
            ) / n

            analysis["per_method"][method] = {
                "n_sessions": n,
                "task_success_rate": task_success,
                "avg_user_tokens": avg_user_tokens,
                "avg_total_tokens": avg_total_tokens,
                "avg_enforcement_count": avg_enforcement,
                "avg_turns": avg_turns,
                "avg_preference_compliance": avg_compliance,
                "conflict_resolution_accuracy": conflict_accuracy,
                "over_personalization_rate": over_personalization,
            }

        # Comparison
        metrics_to_compare = [
            ("task_success_rate", True),  # higher is better
            ("avg_user_tokens", False),  # lower is better
            ("avg_total_tokens", False),  # lower is better
            ("avg_enforcement_count", False),  # lower is better
            ("avg_preference_compliance", True),  # higher is better
            ("conflict_resolution_accuracy", True),  # higher is better
            ("over_personalization_rate", False),  # lower is better
        ]

        for metric, higher_better in metrics_to_compare:
            values = {m: analysis["per_method"][m][metric] for m in analysis["per_method"]}
            if not values:
                logger.warning(f"No values for metric {metric}, skipping comparison")
                continue
            if higher_better:
                best = max(values, key=values.get)
            else:
                best = min(values, key=values.get)

            analysis["comparison"][metric] = {
                "values": values,
                "best_method": best,
                "best_value": values[best],
            }

        return analysis

    def _generate_report(self, analysis: Dict) -> None:
        """Generate a human-readable report."""
        report_lines = [
            "# Personalization Experiment Report",
            f"\nGenerated: {datetime.now().isoformat()}",
            f"\nConfig: {self.config.n_profiles} profiles, {self.config.n_sessions_per_profile} sessions each",
            "\n## Method Comparison\n",
        ]

        # Create comparison table
        metrics_display = [
            ("Task Success", "task_success_rate", "{:.1%}"),
            ("User Effort (tokens)", "avg_user_tokens", "{:.0f}"),
            ("Total Tokens", "avg_total_tokens", "{:.0f}"),
            ("Enforcement Count", "avg_enforcement_count", "{:.2f}"),
            ("Preference Compliance", "avg_preference_compliance", "{:.1%}"),
            ("Conflict Resolution", "conflict_resolution_accuracy", "{:.1%}"),
            ("Over-personalization", "over_personalization_rate", "{:.1%}"),
        ]

        methods = list(analysis["per_method"].keys())

        # Header
        header = "| Metric |" + "|".join(f" {m} " for m in methods) + "| Best |"
        separator = "|" + "|".join(["-" * (len(m) + 2) for m in ["Metric"] + methods + ["Best"]]) + "|"

        report_lines.extend([header, separator])

        for display_name, metric_key, fmt in metrics_display:
            row = f"| {display_name} |"
            for m in methods:
                val = analysis["per_method"].get(m, {}).get(metric_key, 0)
                row += f" {fmt.format(val)} |"

            if metric_key in analysis.get("comparison", {}):
                best = analysis["comparison"][metric_key]["best_method"]
            else:
                best = "N/A"
            row += f" {best} |"
            report_lines.append(row)

        # Key findings
        report_lines.extend([
            "\n## Key Findings\n",
        ])

        # Find advantages of proposed methods
        rag_vector = analysis["per_method"].get("rag_vector", {})
        rag = analysis["per_method"].get("rag", {})
        contextual = analysis["per_method"].get("contextual", {})
        all_memory = analysis["per_method"].get("all_memory", {})

        if rag_vector and contextual:
            token_reduction = (contextual.get("avg_total_tokens", 0) - rag_vector.get("avg_total_tokens", 0)) / contextual.get("avg_total_tokens", 1) * 100
            report_lines.append(f"- **Token Efficiency**: RAG+Vector uses {token_reduction:.1f}% fewer tokens than contextual memory")

        if rag_vector and all_memory:
            conflict_improvement = rag_vector.get("conflict_resolution_accuracy", 0) - all_memory.get("conflict_resolution_accuracy", 0)
            report_lines.append(f"- **Conflict Resolution**: RAG+Vector improves by {conflict_improvement:.1%} over all-memory baseline")

        if rag_vector:
            report_lines.append(f"- **Over-personalization**: RAG+Vector rate: {rag_vector.get('over_personalization_rate', 0):.1%}")

        # Save report
        report_path = self.output_dir / "report.md"
        with open(report_path, "w") as f:
            f.write("\n".join(report_lines))

        logger.info(f"Report saved to {report_path}")


def main():
    parser = argparse.ArgumentParser(description="Run personalization experiments")
    parser.add_argument("--config", type=str, help="Path to config YAML file")
    parser.add_argument("--methods", type=str, default="vanilla,contextual,rag,rag_vector",
                        help="Comma-separated list of methods to compare")
    parser.add_argument("--datasets", type=str, default="math-hard,math-500,bigcodebench",
                        help="Comma-separated list of datasets")
    parser.add_argument("--n-profiles", type=int, default=200, help="Number of user profiles")
    parser.add_argument("--n-sessions", type=int, default=30, help="Sessions per profile")
    parser.add_argument("--max-turns", type=int, default=15, help="Max turns per session")
    parser.add_argument("--output-dir", type=str, default="results", help="Output directory")
    parser.add_argument("--profile-path", type=str, help="Path to pre-generated profiles")
    parser.add_argument("--start-profile", type=int, default=0,
                        help="Start profile index (inclusive, 0-indexed)")
    parser.add_argument("--end-profile", type=int, default=None,
                        help="End profile index (exclusive). If not set, runs all profiles from start")

    # vLLM and parallel processing options
    parser.add_argument("--use-vllm", action="store_true",
                        help="Use vLLM servers for inference (much faster)")
    parser.add_argument("--vllm-user-url", type=str, default="http://localhost:8004/v1",
                        help="vLLM server URL for user simulator (70B)")
    parser.add_argument("--vllm-agent-url", type=str, default="http://localhost:8003/v1",
                        help="vLLM server URL for agent (8B)")
    # OpenAI user agent options
    parser.add_argument("--use-openai-user", action="store_true",
                        help="Use OpenAI API (GPT-5) for user simulation instead of vLLM")
    parser.add_argument("--openai-user-model", type=str, default="gpt-5",
                        help="OpenAI model name for user simulator (default: gpt-5)")
    parser.add_argument("--reward-mode", type=str, default="keyword", choices=["keyword", "llm", "llm_local"],
                        help="Reward mode: 'keyword' (user signals), 'llm' (GPT-4o-mini), or 'llm_local' (local vLLM)")
    parser.add_argument("--reward-vllm-url", type=str, default="http://localhost:8005/v1",
                        help="vLLM server URL for local reward model (when --reward-mode=llm_local)")

    parser.add_argument("--parallel-profiles", type=int, default=50,
                        help="Number of profiles to process in parallel (requires --use-vllm)")
    parser.add_argument("--use-batch-processing", action="store_true", default=True,
                        help="Use turn-synchronous batch processing for vanilla/all_memory")
    parser.add_argument("--no-batch-processing", action="store_false", dest="use_batch_processing",
                        help="Disable batch processing")
    parser.add_argument("--batch-size", type=int, default=50,
                        help="Number of conversations to batch together")
    parser.add_argument("--continue-from", type=str, default=None,
                        help="Path to existing output directory to continue from (for extending sessions)")

    args = parser.parse_args()

    # Load or create config
    if args.config and Path(args.config).exists():
        with open(args.config) as f:
            config_dict = yaml.safe_load(f)
        config = ExperimentConfig(**config_dict)
    else:
        config = ExperimentConfig(
            methods=args.methods.split(","),
            datasets=args.datasets.split(","),
            n_profiles=args.n_profiles,
            n_sessions_per_profile=args.n_sessions,
            max_turns_per_session=args.max_turns,
            output_dir=args.output_dir,
            profile_path=args.profile_path,
            start_profile=args.start_profile,
            end_profile=args.end_profile,
            use_vllm=args.use_vllm,
            vllm_user_url=args.vllm_user_url,
            vllm_agent_url=args.vllm_agent_url,
            use_openai_user=args.use_openai_user,
            openai_user_model=args.openai_user_model,
            reward_mode=args.reward_mode,
            reward_vllm_url=args.reward_vllm_url,
            parallel_profiles=args.parallel_profiles,
            use_batch_processing=args.use_batch_processing,
            batch_size_conversations=args.batch_size,
            continue_from=args.continue_from,
        )

    # Run experiments
    runner = ExperimentRunner(config)
    analysis = runner.run_all()

    print("\n" + "=" * 60)
    print("EXPERIMENT COMPLETE")
    print("=" * 60)
    print(f"\nResults saved to: {runner.output_dir}")
    if analysis.get("comparison"):
        print("\nBest methods per metric:")
        for metric, data in analysis["comparison"].items():
            print(f"  {metric}: {data['best_method']} ({data['best_value']:.3f})")
    else:
        print("\nNo comparison data available (sessions may have failed)")


if __name__ == "__main__":
    main()