#!/usr/bin/env python3
"""
Error Case Analysis: Compare two experiments

This script compares evaluation results between two experiments, computing:
- Scoreboard and deltas (per-benchmark and overall)
- Error pattern breakdown using the v2 taxonomy (Python-counted from per-sample artifacts)
- Trend analysis (optional)

Usage:
    python error_case_compare.py expA expB
    python error_case_compare.py expA expB --checkpoint_policy best
    python error_case_compare.py expA expB --benchmarks TextVQA,DocVQA,ChartQA,GQA

Output:
    output/analysis/error_case_compare/{expA}_vs_{expB}/summary.md
    output/analysis/error_case_compare/{expA}_vs_{expB}/report.json
"""

import argparse
import json
import os
import re
import sys
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple
from collections import Counter, defaultdict
from dataclasses import dataclass, asdict

import pandas as pd

# Import error pattern module
script_dir = Path(__file__).parent
if str(script_dir) not in sys.path:
    sys.path.insert(0, str(script_dir))

from error_patterns import (
    classify_error_v2,
    count_patterns,
    pattern_summary_table,
    PATTERN_DESCRIPTIONS,
    ErrorPatternResult,
)


# ==============================================================================
# Configuration
# ==============================================================================

PROJECT_ROOT = Path(__file__).resolve().parents[3]


def get_vlmeval_root(project_root: Path) -> Path:
    """Resolve VLMEvalKit root (prefer third_party over VLMEvalKit_local)."""
    candidate = project_root / "third_party" / "VLMEvalKit"
    if candidate.exists():
        return candidate
    return project_root / "VLMEvalKit_local"

# Benchmark directory mappings
BENCHMARK_DIRS = {
    "TextVQA": "TextVQA_VAL",
    "DocVQA": "DocVQA_VAL",
    "ChartQA": "ChartQA_TEST",
    "GQA": "GQA_TestDev_Balanced",
    "POPE": "POPE",
    "MMMU": "MMMU_DEV_VAL",
    "MMBench": "MMBench_DEV_EN_V11",
    "MathVista": "MathVista_MINI",
    "AI2D": "AI2D_TEST",
    "RealWorldQA": "RealWorldQA",
    "SEEDBench2": "SEEDBench2_Plus",
}

# Short-answer QA benchmarks (have per-sample XLSX)
SHORT_ANSWER_BENCHMARKS = ["TextVQA", "DocVQA", "ChartQA", "GQA"]

# Default output directory
DEFAULT_OUTPUT_DIR = PROJECT_ROOT / "output" / "analysis" / "error_case_compare"


# ==============================================================================
# Data Classes
# ==============================================================================

@dataclass
class ExperimentInfo:
    """Information about an experiment."""
    name: str
    vlmeval_dir: Optional[Path] = None  # VLMEvalKit_local/output_models/{exp}
    training_dir: Optional[Path] = None  # output_models/{exp}
    checkpoint_results_path: Optional[Path] = None  # Aggregated results
    checkpoints: List[str] = None  # Available checkpoints

    def __post_init__(self):
        if self.checkpoints is None:
            self.checkpoints = []


@dataclass
class BenchmarkResult:
    """Results for a single benchmark at a checkpoint."""
    benchmark: str
    checkpoint: str
    total_samples: int
    correct_count: int
    accuracy: float
    pattern_counts: Dict[str, int]
    pattern_details: List[Dict[str, Any]]  # Detailed per-case info (sample)
    source_file: Optional[str] = None


@dataclass
class ComparisonResult:
    """Full comparison result between two experiments."""
    exp_a: str
    exp_b: str
    checkpoint_a: str
    checkpoint_b: str
    benchmarks: Dict[str, Dict[str, Any]]  # {benchmark: {expA: ..., expB: ..., delta: ...}}
    sources: List[str]
    generated_at: str


# ==============================================================================
# Experiment Discovery
# ==============================================================================

def find_experiment_paths(exp_name: str) -> ExperimentInfo:
    """
    Find paths for an experiment by name.
    
    Search order:
    1. VLMEvalKit/output_models/{exp} (third_party or VLMEvalKit_local)
    2. output_models/{exp}
    """
    info = ExperimentInfo(name=exp_name)
    
    # Check VLMEvalKit location
    vlmeval_root = get_vlmeval_root(PROJECT_ROOT)
    vlmeval_path = vlmeval_root / "output_models" / exp_name
    if vlmeval_path.exists():
        info.vlmeval_dir = vlmeval_path
        
        # Check for aggregated results
        for fname in ["checkpoint_results.md", "checkpoint_results.csv"]:
            fpath = vlmeval_path / fname
            if fpath.exists():
                info.checkpoint_results_path = fpath
                break
        
        # Find checkpoints
        eval_vlmeval = vlmeval_path / "eval_vlmeval"
        if eval_vlmeval.exists():
            info.checkpoints = sorted([
                d.name for d in eval_vlmeval.iterdir()
                if d.is_dir() and d.name.startswith("checkpoint-")
            ], key=lambda x: int(re.search(r'\d+', x).group()) if re.search(r'\d+', x) else 0)
    
    # Check training output location
    training_path = PROJECT_ROOT / "output_models" / exp_name
    if training_path.exists():
        info.training_dir = training_path
        
        # Also check for checkpoints here
        eval_vlmeval = training_path / "eval_vlmeval"
        if eval_vlmeval.exists() and not info.checkpoints:
            info.checkpoints = sorted([
                d.name for d in eval_vlmeval.iterdir()
                if d.is_dir() and d.name.startswith("checkpoint-")
            ], key=lambda x: int(re.search(r'\d+', x).group()) if re.search(r'\d+', x) else 0)
    
    return info


def get_checkpoint_dir(info: ExperimentInfo, checkpoint: str) -> Optional[Path]:
    """Get the evaluation directory for a specific checkpoint."""
    # Try VLMEvalKit location first
    if info.vlmeval_dir:
        ckpt_dir = info.vlmeval_dir / "eval_vlmeval" / checkpoint
        if ckpt_dir.exists():
            return ckpt_dir
    
    # Try training output location
    if info.training_dir:
        ckpt_dir = info.training_dir / "eval_vlmeval" / checkpoint
        if ckpt_dir.exists():
            return ckpt_dir
    
    return None


def select_checkpoint(info: ExperimentInfo, policy: str) -> Optional[str]:
    """
    Select a checkpoint based on policy.
    
    Policies:
    - best: Use aggregated results to find best Avg
    - latest: Use latest checkpoint
    - step:<n>: Use specific step
    """
    if not info.checkpoints:
        return None
    
    if policy == "latest":
        return info.checkpoints[-1]
    
    if policy.startswith("step:"):
        target_step = int(policy.split(":")[1])
        for ckpt in info.checkpoints:
            match = re.search(r'checkpoint-(\d+)', ckpt)
            if match and int(match.group(1)) == target_step:
                return ckpt
        return None
    
    if policy == "best":
        # Try to find best from aggregated results
        if info.checkpoint_results_path:
            best = find_best_checkpoint_from_results(info.checkpoint_results_path)
            if best:
                return best
        # Fallback to latest
        return info.checkpoints[-1]
    
    return info.checkpoints[-1]


def find_best_checkpoint_from_results(results_path: Path) -> Optional[str]:
    """Find checkpoint with best Avg from aggregated results."""
    if results_path.suffix == ".csv":
        try:
            df = pd.read_csv(results_path)
            if "Avg" in df.columns and "Step" in df.columns:
                best_idx = df["Avg"].idxmax()
                best_step = int(df.loc[best_idx, "Step"])
                return f"checkpoint-{best_step}"
        except Exception:
            pass
    
    elif results_path.suffix == ".md":
        try:
            content = results_path.read_text()
            # Parse markdown table
            lines = content.strip().split('\n')
            best_avg = -1
            best_step = None
            for line in lines:
                if line.startswith('|') and 'Step' not in line and '---' not in line:
                    parts = [p.strip() for p in line.split('|')[1:-1]]
                    if len(parts) >= 2:
                        try:
                            step = int(parts[0])
                            avg = float(parts[-1])  # Avg is usually last column
                            if avg > best_avg:
                                best_avg = avg
                                best_step = step
                        except (ValueError, IndexError):
                            continue
            if best_step:
                return f"checkpoint-{best_step}"
        except Exception:
            pass
    
    return None


# ==============================================================================
# Per-Sample Loading and Analysis
# ==============================================================================

def find_results_xlsx(benchmark_dir: Path) -> Optional[Path]:
    """Find *_results.xlsx file in benchmark directory."""
    for xlsx_file in benchmark_dir.rglob("*_results.xlsx"):
        if xlsx_file.exists() and xlsx_file.is_file():
            return xlsx_file
    return None


def load_benchmark_samples(
    ckpt_dir: Path,
    benchmark: str,
) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
    """
    Load per-sample results for a benchmark.
    
    Returns:
        (DataFrame with samples, source_file_path)
    """
    benchmark_subdir = BENCHMARK_DIRS.get(benchmark)
    if not benchmark_subdir:
        return None, None
    
    benchmark_path = ckpt_dir / benchmark_subdir
    if not benchmark_path.exists():
        return None, None
    
    xlsx_file = find_results_xlsx(benchmark_path)
    if not xlsx_file:
        return None, None
    
    try:
        df = pd.read_excel(xlsx_file)
        return df, str(xlsx_file)
    except Exception as e:
        print(f"[WARN] Failed to load {xlsx_file}: {e}")
        return None, None


def analyze_benchmark_errors(
    df: pd.DataFrame,
    benchmark: str,
    source_file: str,
    checkpoint: str,
    max_details: int = 100,
) -> BenchmarkResult:
    """
    Analyze errors in a benchmark using the v2 pattern taxonomy.
    
    Args:
        df: DataFrame with columns [index, question, answer, prediction, eval_score, ...]
        benchmark: Benchmark name
        source_file: Path to source XLSX
        checkpoint: Checkpoint name
        max_details: Max number of detailed cases to include
    
    Returns:
        BenchmarkResult with pattern counts and sample details
    """
    total = len(df)
    
    # Get score column
    score_col = "eval_score" if "eval_score" in df.columns else "score"
    if score_col not in df.columns:
        # Try to find any score column
        score_cols = [c for c in df.columns if "score" in c.lower()]
        score_col = score_cols[0] if score_cols else None
    
    if score_col is None:
        return BenchmarkResult(
            benchmark=benchmark,
            checkpoint=checkpoint,
            total_samples=total,
            correct_count=0,
            accuracy=0.0,
            pattern_counts={},
            pattern_details=[],
            source_file=source_file,
        )
    
    # Binary correct (threshold 0.5)
    correct = df[score_col] >= 0.5
    correct_count = int(correct.sum())
    accuracy = correct_count / total * 100 if total > 0 else 0.0
    
    # Analyze incorrect cases
    incorrect_mask = ~correct
    incorrect_df = df[incorrect_mask]
    
    pattern_results: List[ErrorPatternResult] = []
    pattern_details: List[Dict[str, Any]] = []
    
    for idx, row in incorrect_df.iterrows():
        pred = str(row.get("prediction", ""))
        gt = row.get("answer", "")
        if isinstance(gt, list):
            gt = gt[0] if gt else ""
        gt = str(gt)
        question = str(row.get("question", ""))
        
        # Classify error
        result = classify_error_v2(pred, gt, question, benchmark)
        pattern_results.append(result)
        
        # Store detail (limited)
        if len(pattern_details) < max_details:
            pattern_details.append({
                "index": str(row.get("index", idx)),
                "question": question[:200],
                "ground_truth": gt[:200],
                "prediction": pred[:200],
                "score": float(row.get(score_col, 0)),
                "pattern": result.pattern,
                "subpattern": result.subpattern,
                "confidence": result.confidence,
            })
    
    # Count patterns
    pattern_counts = count_patterns(pattern_results)
    
    return BenchmarkResult(
        benchmark=benchmark,
        checkpoint=checkpoint,
        total_samples=total,
        correct_count=correct_count,
        accuracy=accuracy,
        pattern_counts=dict(pattern_counts),
        pattern_details=pattern_details,
        source_file=source_file,
    )


# ==============================================================================
# Comparison Logic
# ==============================================================================

def compare_experiments(
    exp_a_info: ExperimentInfo,
    exp_b_info: ExperimentInfo,
    checkpoint_a: str,
    checkpoint_b: str,
    benchmarks: List[str],
) -> ComparisonResult:
    """
    Compare two experiments at specified checkpoints.
    """
    results: Dict[str, Dict[str, Any]] = {}
    sources: List[str] = []
    
    ckpt_dir_a = get_checkpoint_dir(exp_a_info, checkpoint_a)
    ckpt_dir_b = get_checkpoint_dir(exp_b_info, checkpoint_b)
    
    for benchmark in benchmarks:
        results[benchmark] = {
            "exp_a": None,
            "exp_b": None,
            "delta": None,
        }
        
        # Analyze exp_a
        if ckpt_dir_a:
            df_a, src_a = load_benchmark_samples(ckpt_dir_a, benchmark)
            if df_a is not None:
                result_a = analyze_benchmark_errors(df_a, benchmark, src_a, checkpoint_a)
                results[benchmark]["exp_a"] = asdict(result_a)
                if src_a:
                    sources.append(src_a)
        
        # Analyze exp_b
        if ckpt_dir_b:
            df_b, src_b = load_benchmark_samples(ckpt_dir_b, benchmark)
            if df_b is not None:
                result_b = analyze_benchmark_errors(df_b, benchmark, src_b, checkpoint_b)
                results[benchmark]["exp_b"] = asdict(result_b)
                if src_b:
                    sources.append(src_b)
        
        # Compute delta
        if results[benchmark]["exp_a"] and results[benchmark]["exp_b"]:
            acc_a = results[benchmark]["exp_a"]["accuracy"]
            acc_b = results[benchmark]["exp_b"]["accuracy"]
            results[benchmark]["delta"] = {
                "accuracy_diff": acc_a - acc_b,
                "correct_diff": results[benchmark]["exp_a"]["correct_count"] - results[benchmark]["exp_b"]["correct_count"],
            }
    
    return ComparisonResult(
        exp_a=exp_a_info.name,
        exp_b=exp_b_info.name,
        checkpoint_a=checkpoint_a,
        checkpoint_b=checkpoint_b,
        benchmarks=results,
        sources=list(set(sources)),
        generated_at=datetime.now().isoformat(),
    )


# ==============================================================================
# Report Generation
# ==============================================================================

def generate_markdown_report(result: ComparisonResult) -> str:
    """Generate markdown report from comparison result."""
    lines = []
    
    # Header
    lines.append(f"# Error Case Analysis: {result.exp_a} vs {result.exp_b}")
    lines.append("")
    lines.append(f"**Generated**: {result.generated_at}")
    lines.append(f"**Checkpoint A**: {result.checkpoint_a}")
    lines.append(f"**Checkpoint B**: {result.checkpoint_b}")
    lines.append("")
    
    # Scoreboard
    lines.append("## Scoreboard")
    lines.append("")
    lines.append(f"| Benchmark | {result.exp_a[:20]} | {result.exp_b[:20]} | Delta |")
    lines.append("|-----------|" + "-" * 22 + "|" + "-" * 22 + "|-------|")
    
    total_a = 0
    total_b = 0
    count = 0
    
    for benchmark, data in result.benchmarks.items():
        acc_a = data["exp_a"]["accuracy"] if data["exp_a"] else "-"
        acc_b = data["exp_b"]["accuracy"] if data["exp_b"] else "-"
        
        if isinstance(acc_a, (int, float)) and isinstance(acc_b, (int, float)):
            delta = f"{acc_a - acc_b:+.1f}"
            total_a += acc_a
            total_b += acc_b
            count += 1
        else:
            delta = "-"
        
        acc_a_str = f"{acc_a:.1f}" if isinstance(acc_a, (int, float)) else acc_a
        acc_b_str = f"{acc_b:.1f}" if isinstance(acc_b, (int, float)) else acc_b
        
        lines.append(f"| {benchmark} | {acc_a_str} | {acc_b_str} | {delta} |")
    
    # Average
    if count > 0:
        avg_a = total_a / count
        avg_b = total_b / count
        lines.append(f"| **Average** | **{avg_a:.1f}** | **{avg_b:.1f}** | **{avg_a - avg_b:+.1f}** |")
    
    lines.append("")
    
    # Error Pattern Analysis per benchmark
    lines.append("## Error Pattern Analysis")
    lines.append("")
    
    for benchmark, data in result.benchmarks.items():
        lines.append(f"### {benchmark}")
        lines.append("")
        
        for exp_key, exp_label in [("exp_a", result.exp_a), ("exp_b", result.exp_b)]:
            if not data[exp_key]:
                continue
            
            exp_data = data[exp_key]
            lines.append(f"#### {exp_label} ({exp_data['checkpoint']})")
            lines.append("")
            lines.append(f"- Total samples: {exp_data['total_samples']}")
            lines.append(f"- Correct: {exp_data['correct_count']} ({exp_data['accuracy']:.1f}%)")
            lines.append(f"- Incorrect: {exp_data['total_samples'] - exp_data['correct_count']}")
            lines.append("")
            
            pattern_counts = exp_data.get("pattern_counts", {})
            if pattern_counts:
                total_errors = sum(pattern_counts.values())
                lines.append("| Pattern | Count | % of Errors | Description |")
                lines.append("|---------|-------|-------------|-------------|")
                
                for pattern, cnt in sorted(pattern_counts.items(), key=lambda x: -x[1]):
                    pct = cnt / total_errors * 100 if total_errors > 0 else 0
                    desc = PATTERN_DESCRIPTIONS.get(pattern, "")[:40]
                    lines.append(f"| {pattern} | {cnt} | {pct:.1f}% | {desc} |")
                
                lines.append("")
        
        lines.append("")
    
    # Sample Error Cases
    lines.append("## Sample Error Cases")
    lines.append("")
    
    for benchmark, data in result.benchmarks.items():
        for exp_key, exp_label in [("exp_a", result.exp_a), ("exp_b", result.exp_b)]:
            if not data[exp_key]:
                continue
            
            exp_data = data[exp_key]
            details = exp_data.get("pattern_details", [])
            
            if details:
                lines.append(f"### {benchmark} - {exp_label}")
                lines.append("")
                
                # Show up to 5 examples
                for i, case in enumerate(details[:5]):
                    lines.append(f"**Case {i+1}** (Pattern: `{case['pattern']}`)")
                    lines.append(f"- Question: {case['question'][:150]}...")
                    lines.append(f"- GT: `{case['ground_truth'][:100]}`")
                    lines.append(f"- Pred: `{case['prediction'][:100]}`")
                    lines.append(f"- Score: {case['score']:.2f}")
                    lines.append("")
    
    # Sources
    lines.append("## Sources")
    lines.append("")
    lines.append("*All counts computed from per-sample XLSX files:*")
    for src in result.sources:
        lines.append(f"- `{src}`")
    lines.append("")
    
    # Pattern Legend
    lines.append("## Pattern Legend")
    lines.append("")
    lines.append("| Pattern | Description |")
    lines.append("|---------|-------------|")
    for pattern, desc in sorted(PATTERN_DESCRIPTIONS.items()):
        lines.append(f"| {pattern} | {desc} |")
    
    return "\n".join(lines)


def generate_json_report(result: ComparisonResult) -> Dict[str, Any]:
    """Generate JSON report from comparison result."""
    return {
        "exp_a": result.exp_a,
        "exp_b": result.exp_b,
        "checkpoint_a": result.checkpoint_a,
        "checkpoint_b": result.checkpoint_b,
        "benchmarks": result.benchmarks,
        "sources": result.sources,
        "generated_at": result.generated_at,
        "pattern_descriptions": PATTERN_DESCRIPTIONS,
    }


# ==============================================================================
# Main CLI
# ==============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Compare error patterns between two experiments",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python error_case_compare.py exp_bd4 exp_bd8
    python error_case_compare.py exp_bd4 exp_bd8 --checkpoint_policy latest
    python error_case_compare.py exp_bd4 exp_bd8 --benchmarks TextVQA,DocVQA
    python error_case_compare.py exp_bd4 exp_bd8 --checkpoint_policy step:2000
        """,
    )
    parser.add_argument("exp_a", type=str, help="First experiment name")
    parser.add_argument("exp_b", type=str, help="Second experiment name")
    parser.add_argument(
        "--checkpoint_policy",
        type=str,
        default="best",
        help="Checkpoint selection policy: best, latest, or step:<n> (default: best)",
    )
    parser.add_argument(
        "--benchmarks",
        type=str,
        default=",".join(SHORT_ANSWER_BENCHMARKS),
        help=f"Comma-separated benchmarks (default: {','.join(SHORT_ANSWER_BENCHMARKS)})",
    )
    parser.add_argument(
        "--out_dir",
        type=str,
        default=None,
        help="Output directory (default: output/analysis/error_case_compare/{expA}_vs_{expB})",
    )
    
    args = parser.parse_args()
    
    # Parse benchmarks
    benchmarks = [b.strip() for b in args.benchmarks.split(",")]
    
    print(f"[INFO] Comparing: {args.exp_a} vs {args.exp_b}")
    print(f"[INFO] Checkpoint policy: {args.checkpoint_policy}")
    print(f"[INFO] Benchmarks: {benchmarks}")
    print()
    
    # Find experiments
    print("[INFO] Discovering experiment paths...")
    exp_a_info = find_experiment_paths(args.exp_a)
    exp_b_info = find_experiment_paths(args.exp_b)
    
    if not exp_a_info.checkpoints:
        print(f"[ERROR] No checkpoints found for {args.exp_a}")
        return 1
    if not exp_b_info.checkpoints:
        print(f"[ERROR] No checkpoints found for {args.exp_b}")
        return 1
    
    print(f"  {args.exp_a}: {len(exp_a_info.checkpoints)} checkpoints")
    print(f"  {args.exp_b}: {len(exp_b_info.checkpoints)} checkpoints")
    
    # Select checkpoints
    checkpoint_a = select_checkpoint(exp_a_info, args.checkpoint_policy)
    checkpoint_b = select_checkpoint(exp_b_info, args.checkpoint_policy)
    
    if not checkpoint_a:
        print(f"[ERROR] Could not select checkpoint for {args.exp_a}")
        return 1
    if not checkpoint_b:
        print(f"[ERROR] Could not select checkpoint for {args.exp_b}")
        return 1
    
    print(f"  Selected: {checkpoint_a} vs {checkpoint_b}")
    print()
    
    # Run comparison
    print("[INFO] Analyzing error patterns...")
    result = compare_experiments(
        exp_a_info,
        exp_b_info,
        checkpoint_a,
        checkpoint_b,
        benchmarks,
    )
    
    # Determine output directory
    if args.out_dir:
        out_dir = Path(args.out_dir)
    else:
        out_dir = DEFAULT_OUTPUT_DIR / f"{args.exp_a}_vs_{args.exp_b}"
    
    out_dir.mkdir(parents=True, exist_ok=True)
    
    # Write reports
    md_path = out_dir / "summary.md"
    json_path = out_dir / "report.json"
    
    md_content = generate_markdown_report(result)
    with open(md_path, "w") as f:
        f.write(md_content)
    print(f"[INFO] Wrote: {md_path}")
    
    json_content = generate_json_report(result)
    with open(json_path, "w") as f:
        json.dump(json_content, f, indent=2, ensure_ascii=False)
    print(f"[INFO] Wrote: {json_path}")
    
    # Print summary
    print()
    print("=" * 60)
    print("SUMMARY")
    print("=" * 60)
    
    for benchmark, data in result.benchmarks.items():
        if data["exp_a"] and data["exp_b"]:
            acc_a = data["exp_a"]["accuracy"]
            acc_b = data["exp_b"]["accuracy"]
            delta = acc_a - acc_b
            print(f"  {benchmark}: {acc_a:.1f}% vs {acc_b:.1f}% (Δ={delta:+.1f}%)")
    
    print()
    print(f"Full report: {md_path}")
    
    return 0


if __name__ == "__main__":
    sys.exit(main())

