#!/usr/bin/env python3
"""Multi-sample variance experiment for router training data.

This script runs an experiment to determine whether LLM response variance
is informative for router training. It explores two sources of variance:

1. LLM Response Variance: Same query → different responses (temperature > 0)
2. Embedding Variance: Different embedding models/instructions → different representations

The experiment:
1. Samples N random (query, LLM) pairs from existing data
2. Runs K samples per pair with temperature > 0
3. Computes variance statistics
4. Analyzes whether variance is informative for routing decisions

Usage:
    # Run experiment with defaults (1000 pairs, 5 samples each)
    python -m embergraphrouter.data_processing.multi_sample_experiment run

    # Analyze existing results
    python -m embergraphrouter.data_processing.multi_sample_experiment analyze

    # Custom parameters
    python -m embergraphrouter.data_processing.multi_sample_experiment run \
        --n-pairs 500 --n-samples 3 --temperature 0.8

    # Resume interrupted experiment
    python -m embergraphrouter.data_processing.multi_sample_experiment run --resume
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import random
import threading
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
from tqdm import tqdm

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

_MODULE_DIR = Path(__file__).resolve().parent.parent


def _get_provider(llm_name: str) -> str:
    """Extract provider from LLM name for rate limiting."""
    llm_lower = llm_name.lower()
    if "gpt" in llm_lower or "openai" in llm_lower:
        return "openai"
    elif "claude" in llm_lower or "anthropic" in llm_lower:
        return "anthropic"
    elif "gemini" in llm_lower or "google" in llm_lower:
        return "google"
    return "unknown"


class AdaptiveRateLimiter:
    """Per-provider adaptive rate limiter using AIMD (TCP-style congestion control).

    Features:
    - Tracks success/failure rates per provider
    - Backs off on rate limits (multiplicative decrease)
    - Gradually recovers on success (additive increase)
    - Thread-safe
    """

    def __init__(
        self,
        base_delay: float = 0.1,
        max_delay: float = 30.0,
        window_size: int = 20,
    ):
        self._lock = threading.Lock()
        self._provider_state: dict[str, dict] = {}
        self._base_delay = base_delay
        self._max_delay = max_delay
        self._window_size = window_size

    def _get_state(self, provider: str) -> dict:
        """Get or create state for a provider."""
        if provider not in self._provider_state:
            self._provider_state[provider] = {
                "delay": self._base_delay,
                "window": [],
                "last_request": 0.0,
                "in_flight": 0,
            }
        return self._provider_state[provider]

    def acquire(self, provider: str) -> None:
        """Wait before making a request (respecting rate limits)."""
        with self._lock:
            state = self._get_state(provider)
            delay = state["delay"]

            # Extra delay if error rate is high
            if len(state["window"]) >= 5:
                error_rate = 1 - sum(state["window"]) / len(state["window"])
                if error_rate > 0.25:
                    delay *= 1 + error_rate

            state["in_flight"] += 1

        if delay > 0:
            time.sleep(delay)

    def release_success(self, provider: str) -> None:
        """Record a successful request."""
        with self._lock:
            state = self._get_state(provider)
            state["window"].append(True)
            if len(state["window"]) > self._window_size:
                state["window"].pop(0)
            # Additive increase
            state["delay"] = max(self._base_delay, state["delay"] * 0.95)
            state["in_flight"] -= 1

    def release_error(self, provider: str, is_rate_limit: bool = False) -> None:
        """Record a failed request."""
        with self._lock:
            state = self._get_state(provider)
            state["window"].append(False)
            if len(state["window"]) > self._window_size:
                state["window"].pop(0)
            # Multiplicative decrease
            if is_rate_limit:
                state["delay"] = min(state["delay"] * 2.0, self._max_delay)
            else:
                state["delay"] = min(state["delay"] * 1.2, self._max_delay)
            state["in_flight"] -= 1

    def get_stats(self) -> dict[str, dict]:
        """Get statistics for all providers."""
        with self._lock:
            return {
                provider: {
                    "delay": state["delay"],
                    "success_rate": sum(state["window"]) / len(state["window"])
                    if state["window"]
                    else 1.0,
                    "in_flight": state["in_flight"],
                }
                for provider, state in self._provider_state.items()
            }


@dataclass
class ExperimentConfig:
    """Configuration for multi-sample experiment."""

    n_pairs: int = 1000  # Number of (query, LLM) pairs to sample
    n_samples: int = 5  # Samples per pair
    temperature: float = 0.7  # Sampling temperature
    max_tokens: int = 16384  # Max output tokens
    max_workers: int = 20  # Concurrent API calls
    max_workers_per_provider: int = 5  # Per-provider concurrency
    seed: int = 42  # Random seed for reproducibility
    stratify_by_task: bool = True  # Ensure coverage across tasks
    stratify_by_llm: bool = True  # Ensure coverage across LLMs
    checkpoint_interval: int = 50  # Checkpoint every N samples
    save_interval: int = 10  # Save cache every N samples


@dataclass
class SampleResult:
    """Result from a single sample."""

    query: str
    llm: str
    sample_id: int
    response: str
    effect: float
    latency: float
    timestamp: float


@dataclass
class PairStatistics:
    """Aggregated statistics for a (query, LLM) pair."""

    query: str
    llm: str
    task_id: str
    ground_truth: str
    n_samples: int
    effects: list[float]
    effect_mean: float
    effect_std: float
    effect_min: float
    effect_max: float
    effect_range: float
    n_unique_answers: int
    is_high_variance: bool  # std > 0.1
    is_bimodal: bool  # mix of 0s and 1s


def _query_sample_hash(query: str, llm_name: str, sample_id: int) -> str:
    """Generate cache key for (query, llm, sample_id) triple."""
    return hashlib.sha256(
        f"{query}||{llm_name}||{sample_id}".encode()
    ).hexdigest()[:16]


class MultiSampleExperiment:
    """Runs multi-sample variance experiment with adaptive parallelism."""

    def __init__(
        self,
        config: ExperimentConfig,
        data_path: Path | None = None,
        llm_descriptions_path: Path | None = None,
    ):
        self.config = config
        self.data_path = data_path or (_MODULE_DIR / "data" / "router_data.csv")
        self.llm_descriptions_path = llm_descriptions_path or (
            _MODULE_DIR / "configs" / "LLM_Descriptions.json"
        )

        # Experiment output directory
        self.output_dir = _MODULE_DIR / "data" / ".multi_sample_experiment"
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Cache and checkpoint files
        self.cache_file = self.output_dir / "sample_cache.json"
        self.checkpoint_file = self.output_dir / "checkpoint.json"
        self.results_file = self.output_dir / "experiment_results.json"
        self.statistics_file = self.output_dir / "pair_statistics.csv"
        self.pairs_file = self.output_dir / "selected_pairs.csv"

        # Thread-safe cache access
        self._cache_lock = threading.Lock()
        self._unsaved_count = 0

        # Load cache
        self.cache = self._load_cache()

        # Adaptive rate limiter
        self.rate_limiter = AdaptiveRateLimiter()

        # LLM engine (lazy loaded)
        self._llm_engine = None
        self._llm_descriptions = None

    def _load_cache(self) -> dict:
        """Load existing sample cache."""
        if self.cache_file.exists():
            try:
                return json.loads(self.cache_file.read_text(encoding="utf-8"))
            except json.JSONDecodeError:
                logger.warning("Corrupted cache file, starting fresh")
                return {}
        return {}

    def _save_cache(self, force: bool = False) -> None:
        """Save sample cache to disk."""
        with self._cache_lock:
            if force or self._unsaved_count >= self.config.save_interval:
                self.cache_file.write_text(
                    json.dumps(self.cache, indent=2), encoding="utf-8"
                )
                self._unsaved_count = 0

    def _save_checkpoint(self, completed_keys: set[str]) -> None:
        """Save checkpoint for resumption."""
        checkpoint = {
            "completed_keys": list(completed_keys),
            "timestamp": time.time(),
            "config": {
                "n_pairs": self.config.n_pairs,
                "n_samples": self.config.n_samples,
                "temperature": self.config.temperature,
            },
        }
        self.checkpoint_file.write_text(
            json.dumps(checkpoint, indent=2), encoding="utf-8"
        )

    def _load_checkpoint(self) -> set[str]:
        """Load checkpoint for resumption."""
        if self.checkpoint_file.exists():
            try:
                checkpoint = json.loads(self.checkpoint_file.read_text())
                return set(checkpoint.get("completed_keys", []))
            except json.JSONDecodeError:
                return set()
        return set()

    def _load_llm_engine(self):
        """Lazy load LLM engine."""
        if self._llm_engine is None:
            from .llm_engine import LLMEngine

            with open(self.llm_descriptions_path) as f:
                self._llm_descriptions = json.load(f)

            self._llm_engine = LLMEngine(
                llm_names=list(self._llm_descriptions.keys()),
                llm_description=self._llm_descriptions,
            )
        return self._llm_engine

    def select_pairs(self, df: pd.DataFrame, resume: bool = False) -> pd.DataFrame:
        """Select (query, LLM) pairs for the experiment.

        Uses stratified sampling to ensure coverage across tasks and LLMs.
        """
        # Check for saved pairs (for resumption)
        if resume and self.pairs_file.exists():
            logger.info("Loading previously selected pairs from %s", self.pairs_file)
            return pd.read_csv(self.pairs_file)

        random.seed(self.config.seed)
        np.random.seed(self.config.seed)

        # Get unique (query, llm, task_id, ground_truth, metric) combinations
        unique_pairs = df.drop_duplicates(subset=["query", "llm"])[
            ["query", "llm", "task_id", "ground_truth", "metric"]
        ].copy()

        logger.info("Total unique (query, LLM) pairs: %d", len(unique_pairs))

        if self.config.stratify_by_task and self.config.stratify_by_llm:
            # Stratified sampling: equal samples per (task, llm) combination
            groups = unique_pairs.groupby(["task_id", "llm"])
            n_groups = groups.ngroups
            samples_per_group = max(1, self.config.n_pairs // n_groups)

            selected = []
            for (task, llm), group in groups:
                n_select = min(samples_per_group, len(group))
                selected.append(
                    group.sample(n=n_select, random_state=self.config.seed)
                )

            selected_df = pd.concat(selected, ignore_index=True)

            # If we need more samples, random sample from remainder
            if len(selected_df) < self.config.n_pairs:
                remaining = unique_pairs[~unique_pairs.index.isin(selected_df.index)]
                n_more = self.config.n_pairs - len(selected_df)
                if len(remaining) >= n_more:
                    more = remaining.sample(n=n_more, random_state=self.config.seed)
                    selected_df = pd.concat([selected_df, more], ignore_index=True)

        else:
            # Simple random sampling
            selected_df = unique_pairs.sample(
                n=min(self.config.n_pairs, len(unique_pairs)),
                random_state=self.config.seed,
            )

        # Save for resumption
        selected_df.to_csv(self.pairs_file, index=False)
        logger.info("Selected %d pairs for experiment", len(selected_df))

        return selected_df

    def _run_single_sample(
        self,
        query: str,
        llm_name: str,
        sample_id: int,
        ground_truth: str,
        metric: str,
    ) -> SampleResult:
        """Run a single sample and return result."""
        cache_key = _query_sample_hash(query, llm_name, sample_id)

        # Check cache first (thread-safe)
        with self._cache_lock:
            if cache_key in self.cache:
                cached = self.cache[cache_key]
                return SampleResult(
                    query=query,
                    llm=llm_name,
                    sample_id=sample_id,
                    response=cached["response"],
                    effect=cached["effect"],
                    latency=cached["latency"],
                    timestamp=cached["timestamp"],
                )

        # Load LLM engine
        engine = self._load_llm_engine()

        # Get model info
        model_id = engine.llm_description[llm_name]["model"]
        provider = _get_provider(llm_name)

        # Import model_prompting with temperature support
        from .utils import model_prompting

        # Acquire rate limit slot
        self.rate_limiter.acquire(provider)

        # Run API call with temperature
        start = time.perf_counter()
        try:
            response = model_prompting(
                llm_model=model_id,
                prompt=query,
                max_token_num=self.config.max_tokens,
                temperature=self.config.temperature,
            )
            self.rate_limiter.release_success(provider)
        except Exception as e:
            error_msg = str(e).lower()
            is_rate_limit = any(
                x in error_msg for x in ["rate", "limit", "429", "quota"]
            )
            self.rate_limiter.release_error(provider, is_rate_limit)
            logger.warning("API error for %s sample %d: %s", llm_name, sample_id, e)
            response = ""

        latency = time.perf_counter() - start

        # Evaluate response
        effect = engine.eval(
            prediction=response,
            ground_truth=ground_truth,
            metric=metric,
        )

        # Cache result (thread-safe)
        result_data = {
            "response": response,
            "effect": effect,
            "latency": latency,
            "timestamp": time.time(),
        }
        with self._cache_lock:
            self.cache[cache_key] = result_data
            self._unsaved_count += 1

        # Periodic save
        self._save_cache()

        return SampleResult(
            query=query,
            llm=llm_name,
            sample_id=sample_id,
            response=response,
            effect=effect,
            latency=latency,
            timestamp=result_data["timestamp"],
        )

    def run_experiment(self, resume: bool = True) -> list[PairStatistics]:
        """Run the multi-sample experiment with adaptive parallelism."""
        # Load data
        logger.info("Loading router data from %s", self.data_path)
        df = pd.read_csv(self.data_path)

        # Select pairs (or load from checkpoint)
        pairs_df = self.select_pairs(df, resume=resume)

        # Load checkpoint for resumption
        completed_keys = self._load_checkpoint() if resume else set()
        if completed_keys:
            logger.info("Resuming from checkpoint with %d completed", len(completed_keys))

        # Prepare work items (skip already cached)
        work_items = []
        for _, row in pairs_df.iterrows():
            for sample_id in range(self.config.n_samples):
                cache_key = _query_sample_hash(row["query"], row["llm"], sample_id)
                if cache_key not in self.cache and cache_key not in completed_keys:
                    work_items.append(
                        {
                            "query": row["query"],
                            "llm": row["llm"],
                            "sample_id": sample_id,
                            "ground_truth": row["ground_truth"],
                            "metric": row["metric"],
                            "task_id": row["task_id"],
                            "cache_key": cache_key,
                        }
                    )

        total_samples = len(pairs_df) * self.config.n_samples
        cached_samples = total_samples - len(work_items)
        logger.info(
            "Total samples: %d (cached: %d, to collect: %d)",
            total_samples,
            cached_samples,
            len(work_items),
        )

        if work_items:
            # Shuffle for better provider distribution
            random.shuffle(work_items)

            # Collect samples with progress bar
            logger.info(
                "Collecting %d samples with %d workers...",
                len(work_items),
                self.config.max_workers,
            )

            checkpoint_count = 0

            with ThreadPoolExecutor(max_workers=self.config.max_workers) as executor:
                futures = {}
                for item in work_items:
                    future = executor.submit(
                        self._run_single_sample,
                        item["query"],
                        item["llm"],
                        item["sample_id"],
                        item["ground_truth"],
                        item["metric"],
                    )
                    futures[future] = item

                with tqdm(total=len(work_items), desc="Sampling") as pbar:
                    for future in as_completed(futures):
                        item = futures[future]
                        try:
                            future.result()
                            completed_keys.add(item["cache_key"])
                            checkpoint_count += 1

                            # Periodic checkpoint
                            if checkpoint_count >= self.config.checkpoint_interval:
                                self._save_checkpoint(completed_keys)
                                self._save_cache(force=True)
                                checkpoint_count = 0

                        except Exception as e:
                            logger.warning(
                                "Failed: %s sample %d: %s",
                                item["llm"],
                                item["sample_id"],
                                e,
                            )
                        pbar.update(1)

                        # Log rate limiter stats periodically
                        if pbar.n % 100 == 0:
                            stats = self.rate_limiter.get_stats()
                            logger.debug("Rate limiter: %s", stats)

            # Final save
            self._save_cache(force=True)
            self._save_checkpoint(completed_keys)

        # Compute statistics
        logger.info("Computing pair statistics...")
        statistics = self._compute_statistics(pairs_df)

        # Save results
        self._save_results(statistics)

        # Clean up checkpoint on success
        if self.checkpoint_file.exists():
            self.checkpoint_file.unlink()

        return statistics

    def _compute_statistics(self, pairs_df: pd.DataFrame) -> list[PairStatistics]:
        """Compute variance statistics for each pair."""
        statistics = []

        for _, row in pairs_df.iterrows():
            query = row["query"]
            llm = row["llm"]
            task_id = row["task_id"]
            ground_truth = row["ground_truth"]

            # Collect effects from all samples
            effects = []
            responses = []
            for sample_id in range(self.config.n_samples):
                cache_key = _query_sample_hash(query, llm, sample_id)
                if cache_key in self.cache:
                    cached = self.cache[cache_key]
                    effects.append(cached["effect"])
                    responses.append(cached["response"][:100])  # First 100 chars

            if not effects:
                continue

            effects_array = np.array(effects)
            n_unique_answers = len(set(responses))

            # Check for bimodal distribution (mix of 0s and 1s)
            is_bimodal = (0.0 in effects and 1.0 in effects) or (
                effects_array.std() > 0.3 and len(set(effects)) > 1
            )

            stats = PairStatistics(
                query=query[:200],  # Truncate for storage
                llm=llm,
                task_id=task_id,
                ground_truth=ground_truth[:50],
                n_samples=len(effects),
                effects=effects,
                effect_mean=float(effects_array.mean()),
                effect_std=float(effects_array.std()),
                effect_min=float(effects_array.min()),
                effect_max=float(effects_array.max()),
                effect_range=float(effects_array.max() - effects_array.min()),
                n_unique_answers=n_unique_answers,
                is_high_variance=effects_array.std() > 0.1,
                is_bimodal=is_bimodal,
            )
            statistics.append(stats)

        return statistics

    def _save_results(self, statistics: list[PairStatistics]) -> None:
        """Save experiment results."""
        # Save as CSV for easy analysis
        stats_dicts = []
        for s in statistics:
            stats_dicts.append(
                {
                    "query": s.query,
                    "llm": s.llm,
                    "task_id": s.task_id,
                    "n_samples": s.n_samples,
                    "effect_mean": s.effect_mean,
                    "effect_std": s.effect_std,
                    "effect_min": s.effect_min,
                    "effect_max": s.effect_max,
                    "effect_range": s.effect_range,
                    "n_unique_answers": s.n_unique_answers,
                    "is_high_variance": s.is_high_variance,
                    "is_bimodal": s.is_bimodal,
                }
            )

        stats_df = pd.DataFrame(stats_dicts)
        stats_df.to_csv(self.statistics_file, index=False)
        logger.info("Saved statistics to %s", self.statistics_file)

        # Save summary
        summary = self.analyze_results(statistics)
        self.results_file.write_text(json.dumps(summary, indent=2), encoding="utf-8")
        logger.info("Saved summary to %s", self.results_file)

    def analyze_results(
        self, statistics: list[PairStatistics] | None = None
    ) -> dict[str, Any]:
        """Analyze experiment results and generate summary."""
        if statistics is None:
            # Load from file
            if not self.statistics_file.exists():
                raise FileNotFoundError(
                    f"No results found at {self.statistics_file}. "
                    "Run the experiment first."
                )
            stats_df = pd.read_csv(self.statistics_file)
        else:
            stats_df = pd.DataFrame(
                [
                    {
                        "llm": s.llm,
                        "task_id": s.task_id,
                        "effect_std": s.effect_std,
                        "effect_range": s.effect_range,
                        "is_high_variance": s.is_high_variance,
                        "is_bimodal": s.is_bimodal,
                        "n_unique_answers": s.n_unique_answers,
                    }
                    for s in statistics
                ]
            )

        n_pairs = len(stats_df)
        n_high_variance = stats_df["is_high_variance"].sum()
        n_bimodal = stats_df["is_bimodal"].sum()

        summary = {
            "experiment_config": {
                "n_pairs": self.config.n_pairs,
                "n_samples": self.config.n_samples,
                "temperature": self.config.temperature,
            },
            "overall": {
                "total_pairs": n_pairs,
                "high_variance_pairs": int(n_high_variance),
                "high_variance_pct": float(n_high_variance / n_pairs * 100)
                if n_pairs > 0
                else 0,
                "bimodal_pairs": int(n_bimodal),
                "bimodal_pct": float(n_bimodal / n_pairs * 100) if n_pairs > 0 else 0,
                "mean_std": float(stats_df["effect_std"].mean()),
                "median_std": float(stats_df["effect_std"].median()),
                "mean_range": float(stats_df["effect_range"].mean()),
            },
            "by_llm": {},
            "by_task": {},
            "recommendation": "",
        }

        # By LLM
        for llm in stats_df["llm"].unique():
            llm_df = stats_df[stats_df["llm"] == llm]
            summary["by_llm"][llm] = {
                "n_pairs": len(llm_df),
                "high_variance_pct": float(
                    llm_df["is_high_variance"].sum() / len(llm_df) * 100
                ),
                "mean_std": float(llm_df["effect_std"].mean()),
                "bimodal_pct": float(llm_df["is_bimodal"].sum() / len(llm_df) * 100),
            }

        # By task
        for task in stats_df["task_id"].unique():
            task_df = stats_df[stats_df["task_id"] == task]
            summary["by_task"][task] = {
                "n_pairs": len(task_df),
                "high_variance_pct": float(
                    task_df["is_high_variance"].sum() / len(task_df) * 100
                ),
                "mean_std": float(task_df["effect_std"].mean()),
            }

        # Generate recommendation
        high_var_pct = summary["overall"]["high_variance_pct"]
        if high_var_pct > 20:
            summary["recommendation"] = (
                f"RECOMMENDED: Multi-sample collection. "
                f"{high_var_pct:.1f}% of pairs have σ > 0.1, indicating "
                f"significant variance that is informative for routing."
            )
        elif high_var_pct > 10:
            summary["recommendation"] = (
                f"CONSIDER: Multi-sample collection may be beneficial. "
                f"{high_var_pct:.1f}% of pairs have high variance. "
                f"Consider selective sampling for high-variance LLMs/tasks."
            )
        else:
            summary["recommendation"] = (
                f"NOT RECOMMENDED: Only {high_var_pct:.1f}% of pairs have "
                f"σ > 0.1. LLM responses are mostly consistent. "
                f"Single sampling is sufficient."
            )

        return summary

    def print_analysis(self, summary: dict[str, Any] | None = None) -> None:
        """Print human-readable analysis."""
        if summary is None:
            summary = self.analyze_results()

        print("\n" + "=" * 80)
        print("MULTI-SAMPLE VARIANCE EXPERIMENT RESULTS")
        print("=" * 80)

        cfg = summary["experiment_config"]
        print(
            f"\nConfig: {cfg['n_pairs']} pairs × {cfg['n_samples']} samples "
            f"@ temperature={cfg['temperature']}"
        )

        overall = summary["overall"]
        print("\n--- OVERALL ---")
        print(f"Total pairs analyzed: {overall['total_pairs']}")
        print(
            f"High variance (σ > 0.1): {overall['high_variance_pairs']} "
            f"({overall['high_variance_pct']:.1f}%)"
        )
        print(f"Bimodal responses: {overall['bimodal_pairs']} ({overall['bimodal_pct']:.1f}%)")
        print(f"Mean σ: {overall['mean_std']:.4f}")
        print(f"Median σ: {overall['median_std']:.4f}")
        print(f"Mean range: {overall['mean_range']:.4f}")

        print("\n--- BY LLM ---")
        print(
            f"{'LLM':<30} {'Pairs':>6} {'High Var%':>10} {'Mean σ':>8} {'Bimodal%':>10}"
        )
        print("-" * 70)
        for llm, stats in sorted(
            summary["by_llm"].items(),
            key=lambda x: x[1]["mean_std"],
            reverse=True,
        ):
            print(
                f"{llm:<30} {stats['n_pairs']:>6} "
                f"{stats['high_variance_pct']:>9.1f}% "
                f"{stats['mean_std']:>8.4f} "
                f"{stats['bimodal_pct']:>9.1f}%"
            )

        print("\n--- TOP 5 HIGH-VARIANCE TASKS ---")
        top_tasks = sorted(
            summary["by_task"].items(),
            key=lambda x: x[1]["mean_std"],
            reverse=True,
        )[:5]
        for task, stats in top_tasks:
            print(
                f"  {task}: σ={stats['mean_std']:.4f}, "
                f"{stats['high_variance_pct']:.1f}% high var"
            )

        print("\n--- RECOMMENDATION ---")
        print(f"  {summary['recommendation']}")
        print("\n" + "=" * 80)


def main():
    """CLI entry point."""
    import argparse

    parser = argparse.ArgumentParser(
        description="Multi-sample variance experiment for router training"
    )
    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # Run command
    run_parser = subparsers.add_parser("run", help="Run experiment")
    run_parser.add_argument(
        "--n-pairs", type=int, default=1000, help="Number of (query, LLM) pairs"
    )
    run_parser.add_argument("--n-samples", type=int, default=5, help="Samples per pair")
    run_parser.add_argument(
        "--temperature", type=float, default=0.7, help="Sampling temperature"
    )
    run_parser.add_argument(
        "--max-workers", type=int, default=20, help="Concurrent API calls"
    )
    run_parser.add_argument("--seed", type=int, default=42, help="Random seed")
    run_parser.add_argument(
        "--resume", action="store_true", help="Resume from checkpoint"
    )
    run_parser.add_argument(
        "--no-resume",
        action="store_true",
        help="Start fresh (ignore checkpoint)",
    )

    # Analyze command
    subparsers.add_parser("analyze", help="Analyze existing results")

    # Status command
    subparsers.add_parser("status", help="Show experiment status")

    args = parser.parse_args()

    if args.command == "run":
        config = ExperimentConfig(
            n_pairs=args.n_pairs,
            n_samples=args.n_samples,
            temperature=args.temperature,
            max_workers=args.max_workers,
            seed=args.seed,
        )
        experiment = MultiSampleExperiment(config)
        resume = not args.no_resume  # Resume by default
        statistics = experiment.run_experiment(resume=resume)
        summary = experiment.analyze_results(statistics)
        experiment.print_analysis(summary)

    elif args.command == "analyze":
        config = ExperimentConfig()  # Defaults don't matter for analysis
        experiment = MultiSampleExperiment(config)
        experiment.print_analysis()

    elif args.command == "status":
        config = ExperimentConfig()
        experiment = MultiSampleExperiment(config)

        print("\n--- EXPERIMENT STATUS ---")
        print(f"Cache file: {experiment.cache_file}")
        print(f"Cache entries: {len(experiment.cache)}")

        if experiment.checkpoint_file.exists():
            checkpoint = json.loads(experiment.checkpoint_file.read_text())
            print(f"Checkpoint: {len(checkpoint.get('completed_keys', []))} completed")
            print(f"Checkpoint time: {time.ctime(checkpoint.get('timestamp', 0))}")
        else:
            print("Checkpoint: None")

        if experiment.statistics_file.exists():
            stats_df = pd.read_csv(experiment.statistics_file)
            print(f"Results: {len(stats_df)} pairs analyzed")
        else:
            print("Results: Not yet computed")

    else:
        parser.print_help()


if __name__ == "__main__":
    main()
