#!/usr/bin/env python3
"""
Analyze BD8 vs BD4 short answer QA performance.

This script compares VLMEvalKit results between BD4 and BD8 experiments,
identifying error patterns where BD8 fails but BD4 succeeds.

Usage:
    python analyze_short_answer_qa.py --bd4_dir <path> --bd8_dir <path> --output_dir <path>
"""

import argparse
import json
import os
from pathlib import Path
from typing import Dict, List, Tuple, Optional
from collections import Counter, defaultdict
import pandas as pd
import numpy as np


# Benchmark directories mapping
BENCHMARK_DIRS = {
    "TextVQA": "TextVQA_VAL",
    "DocVQA": "DocVQA_VAL", 
    "ChartQA": "ChartQA_TEST",
    "GQA": "GQA_TestDev_Balanced",
}


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"


DEFAULT_VLM_ROOT = get_vlmeval_root(PROJECT_ROOT)
DEFAULT_BD4_DIR = (
    DEFAULT_VLM_ROOT
    / "output_models"
    / "sharegpt4v_8node_bsz256_bd4_fix_eos_flexible_lr_1e-6_unfreeze_vision"
    / "eval_vlmeval"
    / "checkpoint-2300"
)
DEFAULT_BD8_DIR = (
    DEFAULT_VLM_ROOT
    / "output_models"
    / "sharegpt4v_8node_bsz256_bd8_fix_eos_flexible_lr_1e-6_unfreeze_vision"
    / "eval_vlmeval"
    / "checkpoint-2100"
)


def find_results_xlsx(benchmark_dir: Path) -> Optional[Path]:
    """Find the *_results.xlsx file in a benchmark directory."""
    # Look for results xlsx file recursively, return the one that actually exists
    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_results(checkpoint_dir: Path, benchmark_name: str) -> Optional[pd.DataFrame]:
    """Load results for a specific benchmark from a checkpoint directory."""
    benchmark_subdir = BENCHMARK_DIRS.get(benchmark_name)
    if not benchmark_subdir:
        print(f"Unknown benchmark: {benchmark_name}")
        return None
    
    benchmark_path = checkpoint_dir / benchmark_subdir
    if not benchmark_path.exists():
        print(f"Benchmark directory not found: {benchmark_path}")
        return None
    
    xlsx_file = find_results_xlsx(benchmark_path)
    if not xlsx_file:
        print(f"No results xlsx found in: {benchmark_path}")
        return None
    
    print(f"Loading: {xlsx_file}")
    df = pd.read_excel(xlsx_file)
    return df


def classify_error_pattern(bd4_pred: str, bd8_pred: str, gt: str) -> str:
    """
    Classify the type of error BD8 made.
    
    Categories:
    - too_long: BD8 answer is significantly longer than BD4/GT
    - too_short: BD8 answer is too short/truncated
    - format_issue: Different formatting (capitalization, punctuation, extra chars)
    - partial_match: BD8 has part of the answer but not complete
    - completely_wrong: Entirely different answer
    - extra_words: BD8 adds unnecessary words
    """
    bd4_pred = str(bd4_pred).strip().lower()
    bd8_pred = str(bd8_pred).strip().lower()
    gt_lower = str(gt).strip().lower()
    
    # Length analysis
    len_bd4 = len(bd4_pred)
    len_bd8 = len(bd8_pred)
    len_gt = len(gt_lower) if gt_lower else len_bd4
    
    # Check if BD8 is much longer
    if len_bd8 > len_bd4 * 2 and len_bd8 > len_gt * 2:
        return "too_long"
    
    # Check if BD8 is much shorter
    if len_bd8 < len_bd4 * 0.5 and len_bd8 < len_gt * 0.5 and len_bd8 < 3:
        return "too_short"
    
    # Check for extra words (BD8 has the answer but with extra stuff)
    if gt_lower in bd8_pred and len(bd8_pred) > len(gt_lower) + 5:
        return "extra_words"
    
    # Check for partial match
    if len(gt_lower) > 3:
        if gt_lower[:len(gt_lower)//2] in bd8_pred or bd8_pred in gt_lower:
            return "partial_match"
    
    # Check format issues (same content but different format)
    # Remove common punctuation and compare
    import re
    bd4_clean = re.sub(r'[^\w\s]', '', bd4_pred)
    bd8_clean = re.sub(r'[^\w\s]', '', bd8_pred)
    gt_clean = re.sub(r'[^\w\s]', '', gt_lower)
    
    if bd4_clean == bd8_clean or bd8_clean == gt_clean:
        return "format_issue"
    
    # Default to completely wrong
    return "completely_wrong"


def analyze_benchmark(bd4_df: pd.DataFrame, bd8_df: pd.DataFrame, benchmark_name: str) -> Dict:
    """
    Analyze differences between BD4 and BD8 for a benchmark.
    
    Returns dict with:
    - total_samples: total number of samples
    - bd4_correct: samples BD4 got correct
    - bd8_correct: samples BD8 got correct
    - bd4_only_correct: samples where only BD4 is correct
    - bd8_only_correct: samples where only BD8 is correct
    - error_patterns: classification of BD8 errors
    - error_cases: list of (idx, question, gt, bd4_pred, bd8_pred, pattern) tuples
    """
    # Merge on index
    merged = bd4_df.merge(
        bd8_df, 
        on='index', 
        suffixes=('_bd4', '_bd8'),
        how='inner'
    )
    
    print(f"  Merged samples: {len(merged)}")
    
    # Get scores
    bd4_scores = merged['eval_score_bd4'].values
    bd8_scores = merged['eval_score_bd8'].values
    
    # Binary correct/incorrect (threshold at 0.5)
    bd4_correct = bd4_scores >= 0.5
    bd8_correct = bd8_scores >= 0.5
    
    # Find cases where BD4 is correct but BD8 is wrong
    bd4_only = bd4_correct & ~bd8_correct
    bd8_only = ~bd4_correct & bd8_correct
    
    # Analyze error patterns
    error_patterns = Counter()
    error_cases = []
    
    bd4_only_indices = np.where(bd4_only)[0]
    for idx in bd4_only_indices:
        row = merged.iloc[idx]
        question = row['question_bd4']
        gt = row['answer_bd4']  # Ground truth
        bd4_pred = row['prediction_bd4']
        bd8_pred = row['prediction_bd8']
        image_path = row.get('image_path_bd4', row.get('image_path_bd8', ''))
        
        # Get a representative GT for pattern classification
        gt_str = gt[0] if isinstance(gt, list) and len(gt) > 0 else str(gt)
        
        pattern = classify_error_pattern(bd4_pred, bd8_pred, gt_str)
        error_patterns[pattern] += 1
        
        error_cases.append({
            'index': str(row['index']),  # Keep as string since some benchmarks use string IDs
            'question': str(question),
            'ground_truth': gt if isinstance(gt, list) else [str(gt)],
            'bd4_prediction': str(bd4_pred),
            'bd8_prediction': str(bd8_pred),
            'bd4_score': float(row['eval_score_bd4']),
            'bd8_score': float(row['eval_score_bd8']),
            'error_pattern': pattern,
            'image_path': str(image_path),
        })
    
    return {
        'benchmark': benchmark_name,
        'total_samples': len(merged),
        'bd4_correct_count': int(bd4_correct.sum()),
        'bd8_correct_count': int(bd8_correct.sum()),
        'bd4_accuracy': float(bd4_correct.mean() * 100),
        'bd8_accuracy': float(bd8_correct.mean() * 100),
        'bd4_only_correct': int(bd4_only.sum()),
        'bd8_only_correct': int(bd8_only.sum()),
        'both_correct': int((bd4_correct & bd8_correct).sum()),
        'both_wrong': int((~bd4_correct & ~bd8_correct).sum()),
        'error_patterns': dict(error_patterns),
        'error_cases': error_cases,
    }


def generate_report(results: Dict[str, Dict], output_dir: Path) -> str:
    """Generate a markdown report from analysis results."""
    report = ["# BD8 Short Answer QA Analysis Report\n"]
    report.append("## Overview\n")
    report.append("This report analyzes cases where BD8 performs worse than BD4 on short answer QA benchmarks.\n")
    
    # Summary table
    report.append("## Summary\n")
    report.append("| Benchmark | BD4 Acc | BD8 Acc | Gap | BD4 Only Correct | BD8 Only Correct |")
    report.append("|-----------|---------|---------|-----|------------------|------------------|")
    
    total_bd4_only = 0
    total_bd8_only = 0
    
    for benchmark, data in results.items():
        gap = data['bd4_accuracy'] - data['bd8_accuracy']
        report.append(
            f"| {benchmark} | {data['bd4_accuracy']:.1f}% | {data['bd8_accuracy']:.1f}% | "
            f"{gap:+.1f}% | {data['bd4_only_correct']} | {data['bd8_only_correct']} |"
        )
        total_bd4_only += data['bd4_only_correct']
        total_bd8_only += data['bd8_only_correct']
    
    report.append(f"\n**Total BD4-only correct cases: {total_bd4_only}**")
    report.append(f"**Total BD8-only correct cases: {total_bd8_only}**\n")
    
    # Error pattern analysis
    report.append("## Error Pattern Analysis\n")
    report.append("Classification of BD8 errors (cases where BD4 is correct but BD8 is wrong):\n")
    
    for benchmark, data in results.items():
        report.append(f"### {benchmark}\n")
        if data['error_patterns']:
            report.append("| Error Pattern | Count | Percentage |")
            report.append("|---------------|-------|------------|")
            total_errors = sum(data['error_patterns'].values())
            for pattern, count in sorted(data['error_patterns'].items(), key=lambda x: -x[1]):
                pct = count / total_errors * 100 if total_errors > 0 else 0
                report.append(f"| {pattern} | {count} | {pct:.1f}% |")
        else:
            report.append("No BD8-specific errors found.\n")
        report.append("")
    
    # Sample error cases
    report.append("## Sample Error Cases\n")
    report.append("Below are sample cases where BD8 failed but BD4 succeeded:\n")
    
    for benchmark, data in results.items():
        report.append(f"### {benchmark}\n")
        # Show up to 10 examples per benchmark
        for i, case in enumerate(data['error_cases'][:10]):
            report.append(f"#### Case {i+1} (Pattern: {case['error_pattern']})\n")
            report.append(f"- **Question**: {case['question']}")
            gt_display = case['ground_truth'][:3] if len(case['ground_truth']) > 3 else case['ground_truth']
            report.append(f"- **Ground Truth**: {gt_display}")
            report.append(f"- **BD4 Prediction**: `{case['bd4_prediction']}` (score: {case['bd4_score']:.2f})")
            report.append(f"- **BD8 Prediction**: `{case['bd8_prediction']}` (score: {case['bd8_score']:.2f})")
            report.append(f"- **Image**: `{case['image_path']}`\n")
        report.append("")
    
    return "\n".join(report)


def export_error_cases(results: Dict[str, Dict], output_dir: Path):
    """Export error cases to JSON files for LLM analysis."""
    error_cases_dir = output_dir / "error_cases"
    error_cases_dir.mkdir(parents=True, exist_ok=True)
    
    for benchmark, data in results.items():
        # Export all error cases
        output_file = error_cases_dir / f"{benchmark}_errors.json"
        with open(output_file, 'w') as f:
            json.dump({
                'benchmark': benchmark,
                'summary': {
                    'total_errors': len(data['error_cases']),
                    'error_patterns': data['error_patterns'],
                },
                'cases': data['error_cases'],
            }, f, indent=2, ensure_ascii=False)
        print(f"Exported {len(data['error_cases'])} error cases to {output_file}")
    
    # Also create a combined summary for LLM analysis
    combined = []
    for benchmark, data in results.items():
        for case in data['error_cases'][:50]:  # Top 50 per benchmark
            combined.append({
                'benchmark': benchmark,
                **case
            })
    
    combined_file = error_cases_dir / "combined_errors_for_llm.json"
    with open(combined_file, 'w') as f:
        json.dump(combined, f, indent=2, ensure_ascii=False)
    print(f"Exported {len(combined)} combined error cases to {combined_file}")


def main():
    parser = argparse.ArgumentParser(description="Analyze BD8 vs BD4 short answer QA performance")
    parser.add_argument(
        "--bd4_dir",
        type=str,
        default=str(DEFAULT_BD4_DIR),
        help="Path to BD4 checkpoint evaluation directory"
    )
    parser.add_argument(
        "--bd8_dir",
        type=str,
        default=str(DEFAULT_BD8_DIR),
        help="Path to BD8 checkpoint evaluation directory"
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="chengyue_script/dvlm/evaluation/analysis_results",
        help="Output directory for analysis results"
    )
    parser.add_argument(
        "--benchmarks",
        type=str,
        nargs="+",
        default=["TextVQA", "DocVQA", "ChartQA", "GQA"],
        help="Benchmarks to analyze"
    )
    
    args = parser.parse_args()
    
    bd4_dir = Path(args.bd4_dir)
    bd8_dir = Path(args.bd8_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"BD4 directory: {bd4_dir}")
    print(f"BD8 directory: {bd8_dir}")
    print(f"Output directory: {output_dir}")
    print(f"Benchmarks: {args.benchmarks}")
    print()
    
    results = {}
    
    for benchmark in args.benchmarks:
        print(f"Analyzing {benchmark}...")
        
        bd4_df = load_benchmark_results(bd4_dir, benchmark)
        bd8_df = load_benchmark_results(bd8_dir, benchmark)
        
        if bd4_df is None or bd8_df is None:
            print(f"  Skipping {benchmark} - missing data")
            continue
        
        analysis = analyze_benchmark(bd4_df, bd8_df, benchmark)
        results[benchmark] = analysis
        
        print(f"  BD4 accuracy: {analysis['bd4_accuracy']:.1f}%")
        print(f"  BD8 accuracy: {analysis['bd8_accuracy']:.1f}%")
        print(f"  BD4 only correct: {analysis['bd4_only_correct']}")
        print(f"  Error patterns: {analysis['error_patterns']}")
        print()
    
    # Generate report
    report = generate_report(results, output_dir)
    report_file = output_dir / "bd8_short_answer_analysis.md"
    with open(report_file, 'w') as f:
        f.write(report)
    print(f"Report saved to: {report_file}")
    
    # Export error cases
    export_error_cases(results, output_dir)
    
    print("\nAnalysis complete!")


if __name__ == "__main__":
    main()

