#!/usr/bin/env python3
"""
Analyze DocVQA checkpoint regressions: samples that were correct at an early checkpoint
but became wrong at a later checkpoint.

This helps identify "forgetting" or "derailing" patterns during training.

Usage:
    python analyze_docvqa_checkpoint_regressions.py \
        --early_ckpt_dir <path_to_checkpoint-250_eval> \
        --late_ckpt_dir <path_to_checkpoint-1500_eval> \
        --output_dir <output_path>

Example:
    python chengyue_script/dvlm/evaluation/analyze_docvqa_checkpoint_regressions.py \
        --early_ckpt_dir VLMEvalKit_local/output_models/sharegpt4v_text_8node_bsz256_bd8_block_causal_lr_1e-6_unfreeze_vision/eval_vlmeval/checkpoint-250/DocVQA_VAL \
        --late_ckpt_dir VLMEvalKit_local/output_models/sharegpt4v_text_8node_bsz256_bd8_block_causal_lr_1e-6_unfreeze_vision/eval_vlmeval/checkpoint-1500/DocVQA_VAL \
        --output_dir output/docvqa_regression_analysis
"""

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


# Error classification thresholds
NOISE_COVERAGE_LENGTH_RATIO = 2.0  # If pred is >2x longer than GT and contains GT
YES_NO_KEYWORDS = {"yes", "no", "true", "false"}

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_EARLY_CKPT = (
    DEFAULT_VLM_ROOT
    / "output_models"
    / "sharegpt4v_text_8node_bsz256_bd8_block_causal_lr_1e-6_unfreeze_vision"
    / "eval_vlmeval"
    / "checkpoint-250"
    / "DocVQA_VAL"
)
DEFAULT_LATE_CKPT = (
    DEFAULT_VLM_ROOT
    / "output_models"
    / "sharegpt4v_text_8node_bsz256_bd8_block_causal_lr_1e-6_unfreeze_vision"
    / "eval_vlmeval"
    / "checkpoint-1500"
    / "DocVQA_VAL"
)


def find_results_xlsx(benchmark_dir: Path) -> Optional[Path]:
    """Find the *_results.xlsx file in a benchmark directory (recursive)."""
    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_results(ckpt_dir: Path) -> Optional[pd.DataFrame]:
    """Load results from a checkpoint evaluation directory."""
    xlsx_file = find_results_xlsx(ckpt_dir)
    if not xlsx_file:
        print(f"[ERROR] No *_results.xlsx found in: {ckpt_dir}")
        return None
    print(f"[INFO] Loading: {xlsx_file}")
    return pd.read_excel(xlsx_file)


def normalize_text(text: str) -> str:
    """Normalize text for comparison (lowercase, strip, remove extra whitespace)."""
    if pd.isna(text):
        return ""
    text = str(text).strip().lower()
    text = re.sub(r'\s+', ' ', text)
    return text


def get_gt_string(gt: Any) -> str:
    """Extract a representative GT string from possibly a list."""
    if isinstance(gt, list):
        return str(gt[0]) if gt else ""
    return str(gt) if not pd.isna(gt) else ""


def compute_edit_distance(s1: str, s2: str) -> int:
    """Compute Levenshtein edit distance between two strings."""
    if len(s1) < len(s2):
        return compute_edit_distance(s2, s1)
    if len(s2) == 0:
        return len(s1)
    
    prev_row = range(len(s2) + 1)
    for i, c1 in enumerate(s1):
        curr_row = [i + 1]
        for j, c2 in enumerate(s2):
            # Cost is 0 if chars match, 1 otherwise
            insertions = prev_row[j + 1] + 1
            deletions = curr_row[j] + 1
            substitutions = prev_row[j] + (c1 != c2)
            curr_row.append(min(insertions, deletions, substitutions))
        prev_row = curr_row
    return prev_row[-1]


def classify_regression_error(
    early_pred: str, 
    late_pred: str, 
    gt: str,
    early_score: float,
    late_score: float
) -> Dict[str, Any]:
    """
    Classify the type of regression error.
    
    Categories (from meeting note):
    1. noise_coverage: 答案本来在输出里，但因为后续生成了噪声/冗余 token，导致 evaluator/parse 失败
    2. total_derail: 模型在大 block 一次性生成过多 token 时失焦/跑偏，导致输出整体偏离正确答案
    3. yes_no_flip: 针对 yes/no 类问题，答案被翻转
    4. partial_forget: 早期正确但晚期只有部分正确（partial match）
    5. other: 其他类型
    
    Returns dict with:
    - category: the error category
    - confidence: how confident we are in this classification
    - details: additional diagnostic info
    """
    early_norm = normalize_text(early_pred)
    late_norm = normalize_text(late_pred)
    gt_norm = normalize_text(gt)
    
    result = {
        "category": "other",
        "confidence": 0.0,
        "details": {}
    }
    
    # Length analysis
    len_early = len(early_norm)
    len_late = len(late_norm)
    len_gt = len(gt_norm)
    
    result["details"]["len_early"] = len_early
    result["details"]["len_late"] = len_late
    result["details"]["len_gt"] = len_gt
    result["details"]["len_ratio_late_gt"] = len_late / len_gt if len_gt > 0 else float('inf')
    
    # Check for yes/no flip
    early_yn = early_norm in YES_NO_KEYWORDS or any(kw in early_norm.split() for kw in YES_NO_KEYWORDS)
    late_yn = late_norm in YES_NO_KEYWORDS or any(kw in late_norm.split() for kw in YES_NO_KEYWORDS)
    gt_yn = gt_norm in YES_NO_KEYWORDS or any(kw in gt_norm.split() for kw in YES_NO_KEYWORDS)
    
    if gt_yn:
        # Check if early was correct yes/no and late flipped
        early_has_yes = "yes" in early_norm or "true" in early_norm
        early_has_no = "no" in early_norm or "false" in early_norm
        late_has_yes = "yes" in late_norm or "true" in late_norm
        late_has_no = "no" in late_norm or "false" in late_norm
        gt_has_yes = "yes" in gt_norm or "true" in gt_norm
        gt_has_no = "no" in gt_norm or "false" in gt_norm
        
        # Detect flip: early matched GT, late is opposite
        if gt_has_yes and early_has_yes and late_has_no:
            result["category"] = "yes_no_flip"
            result["confidence"] = 0.95
            result["details"]["flip_type"] = "yes→no"
            return result
        if gt_has_no and early_has_no and late_has_yes:
            result["category"] = "yes_no_flip"
            result["confidence"] = 0.95
            result["details"]["flip_type"] = "no→yes"
            return result
    
    # Check for noise coverage: late output contains GT but is much longer
    if gt_norm in late_norm:
        if len_late > len_gt * NOISE_COVERAGE_LENGTH_RATIO:
            result["category"] = "noise_coverage"
            result["confidence"] = 0.9
            result["details"]["gt_in_late"] = True
            result["details"]["noise_ratio"] = len_late / len_gt if len_gt > 0 else float('inf')
            return result
    
    # Check for partial coverage (GT partially in late output)
    if len_gt > 5:
        # Check if significant portion of GT is in late
        gt_words = gt_norm.split()
        late_words = late_norm.split()
        if len(gt_words) > 0:
            overlap = len(set(gt_words) & set(late_words))
            overlap_ratio = overlap / len(gt_words)
            result["details"]["word_overlap_ratio"] = overlap_ratio
            
            if 0.3 < overlap_ratio < 0.8 and len_late > len_gt * 1.5:
                result["category"] = "noise_coverage"
                result["confidence"] = 0.7
                result["details"]["partial_gt_coverage"] = True
                return result
    
    # Check for total derail: late output completely different from GT
    edit_dist = compute_edit_distance(late_norm[:100], gt_norm[:100])  # Cap length for efficiency
    max_len = max(len(late_norm[:100]), len(gt_norm[:100]))
    edit_dist_ratio = edit_dist / max_len if max_len > 0 else 0
    
    result["details"]["edit_distance"] = edit_dist
    result["details"]["edit_distance_ratio"] = edit_dist_ratio
    
    if edit_dist_ratio > 0.7:  # More than 70% different
        result["category"] = "total_derail"
        result["confidence"] = 0.85
        return result
    
    # Check for partial forget
    if edit_dist_ratio > 0.3 and edit_dist_ratio <= 0.7:
        result["category"] = "partial_forget"
        result["confidence"] = 0.6
        return result
    
    # Default: other
    result["category"] = "other"
    result["confidence"] = 0.5
    return result


def analyze_regressions(early_df: pd.DataFrame, late_df: pd.DataFrame) -> Dict:
    """
    Analyze regressions between two checkpoints.
    
    Returns dict with:
    - total_samples: total matched samples
    - early_correct: count correct at early checkpoint
    - late_correct: count correct at late checkpoint
    - regressions: samples that went from correct to wrong
    - improvements: samples that went from wrong to correct
    - regression_cases: detailed list of regression cases
    """
    # Merge on index
    merged = early_df.merge(
        late_df,
        on='index',
        suffixes=('_early', '_late'),
        how='inner'
    )
    
    print(f"[INFO] Merged samples: {len(merged)}")
    
    # Get scores
    early_scores = merged['eval_score_early'].values
    late_scores = merged['eval_score_late'].values
    
    # Binary correct/incorrect (threshold at 0.5 for VQA)
    early_correct = early_scores >= 0.5
    late_correct = late_scores >= 0.5
    
    # Regressions: early correct, late wrong
    regressions = early_correct & ~late_correct
    # Improvements: early wrong, late correct
    improvements = ~early_correct & late_correct
    
    # Analyze regression cases
    error_categories = Counter()
    regression_cases = []
    
    regression_indices = np.where(regressions)[0]
    for idx in regression_indices:
        row = merged.iloc[idx]
        
        question = row['question_early']
        gt = row['answer_early']
        gt_str = get_gt_string(gt)
        early_pred = str(row['prediction_early'])
        late_pred = str(row['prediction_late'])
        early_score = float(row['eval_score_early'])
        late_score = float(row['eval_score_late'])
        image_path = row.get('image_path_early', row.get('image_path_late', ''))
        
        # Classify error
        classification = classify_regression_error(
            early_pred, late_pred, gt_str, early_score, late_score
        )
        error_categories[classification["category"]] += 1
        
        regression_cases.append({
            'index': str(row['index']),
            'question': str(question),
            'ground_truth': gt if isinstance(gt, list) else [str(gt)],
            'gt_string': gt_str,
            'early_prediction': early_pred,
            'late_prediction': late_pred,
            'early_score': early_score,
            'late_score': late_score,
            'score_drop': early_score - late_score,
            'error_category': classification["category"],
            'error_confidence': classification["confidence"],
            'error_details': classification["details"],
            'image_path': str(image_path),
        })
    
    # Sort by score drop (largest first)
    regression_cases.sort(key=lambda x: -x['score_drop'])
    
    return {
        'total_samples': len(merged),
        'early_correct_count': int(early_correct.sum()),
        'late_correct_count': int(late_correct.sum()),
        'early_accuracy': float(early_correct.mean() * 100),
        'late_accuracy': float(late_correct.mean() * 100),
        'accuracy_drop': float((early_correct.mean() - late_correct.mean()) * 100),
        'regression_count': int(regressions.sum()),
        'improvement_count': int(improvements.sum()),
        'both_correct': int((early_correct & late_correct).sum()),
        'both_wrong': int((~early_correct & ~late_correct).sum()),
        'error_categories': dict(error_categories),
        'regression_cases': regression_cases,
    }


def generate_lookahead_audit() -> str:
    """
    Generate the look-ahead bias audit section based on code path analysis.
    
    This documents the attention mask behavior during eval/inference and
    identifies conditions under which look-ahead bias could occur.
    """
    audit = []
    audit.append("## Look-Ahead Bias Audit (Code Path Analysis)\n")
    audit.append("Based on meeting note 1/15/2026: 复核 look-ahead bias 是否仍然存在。\n")
    
    audit.append("### 1. Eval Attention Mask Logic\n")
    audit.append("**File:** `src/lmflow/models/qwen2_5_vl_mdm/modeling.py`\n")
    audit.append("**Function:** `eval_mask(seqlen, block_size, cache_seq_len, update_kv_cache, use_block_causal_mask)`\n")
    audit.append("""
```python
def eval_mask(self, seqlen, block_size, cache_seq_len, update_kv_cache=False, use_block_causal_mask=False):
    q_indices = torch.arange(seqlen) + cache_seq_len
    k_indices = torch.arange(seqlen + cache_seq_len)
    if use_block_causal_mask and update_kv_cache:
        mask = eval_causal_mask(q_indices[:, None], k_indices[None, :])  # q >= k (causal)
    else:
        mask = eval_block_diff_mask(...)  # Returns torch.ones_like(...) → FULL ATTENTION
    return mask
```
""")
    audit.append("**Key Observation:**")
    audit.append("- When `use_block_causal_mask=False` OR `update_kv_cache=False`: uses `eval_block_diff_mask()` which returns **all-ones mask (full attention)**")
    audit.append("- When `use_block_causal_mask=True` AND `update_kv_cache=True`: uses `eval_causal_mask()` which enforces **causal (q >= k)**\n")
    
    audit.append("### 2. Sampling Strategy: `force_left_to_right`\n")
    audit.append("**File:** `chengyue_script/dvlm/models/generation_functions.py`\n")
    audit.append("**Function:** `block_causal_sample_cache(..., force_left_to_right=False)`\n")
    audit.append("""
```python
if force_left_to_right:
    # select the mask token to be predicted from the left to the right
    unmask_idx = torch.zeros_like(mask_idx, dtype=torch.bool)
    unmask_id = mask_idx.nonzero(as_tuple=False)[0]
    unmask_idx[unmask_id[0], unmask_id[1]] = True  # Only unmask leftmost MASK
else:
    unmask_idx = (x1_p > threshold)  # Unmask ALL tokens above threshold (parallel)
```
""")
    audit.append("**Key Observation:**")
    audit.append("- `force_left_to_right=True`: Only unmasks the **leftmost** MASK token per step → **No look-ahead possible** (sequential decoding)")
    audit.append("- `force_left_to_right=False` (default): Unmasks **multiple** tokens in parallel based on confidence threshold → **Look-ahead risk if eval mask is full attention**\n")
    
    audit.append("### 3. Risk Conditions Summary\n")
    audit.append("| Condition | Eval Mask | Sampling | Look-Ahead Risk |")
    audit.append("|-----------|-----------|----------|-----------------|")
    audit.append("| `use_block_causal_mask=True` + `update_kv_cache=True` + `force_left_to_right=True` | Causal | Sequential | **None** |")
    audit.append("| `use_block_causal_mask=True` + `update_kv_cache=True` + `force_left_to_right=False` | Causal | Parallel | **None** (causal mask prevents) |")
    audit.append("| `use_block_causal_mask=False` OR `update_kv_cache=False` + `force_left_to_right=True` | Full | Sequential | **None** (sequential prevents) |")
    audit.append("| `use_block_causal_mask=False` OR `update_kv_cache=False` + `force_left_to_right=False` | Full | Parallel | **HIGH RISK** |")
    audit.append("")
    
    audit.append("### 4. Current VLMEvalKit Wrapper Status\n")
    audit.append("**File:** `chengyue_script/dvlm/evaluation/vlmeval/vlmeval_qwen_vl_mdm.py`\n")
    audit.append("**Issue:** The wrapper does NOT pass `threshold` or `force_left_to_right` to the model's generate function.\n")
    audit.append("- `generate_vlmeval_config.py` supports `THRESHOLD` env var but it's not wired through")
    audit.append("- Default behavior depends on model config's `use_block_causal_mask` setting\n")
    
    audit.append("### 5. Recommendations\n")
    audit.append("1. **For safe parallel decoding (threshold > 0):** Ensure `use_block_causal_mask=True` in model config")
    audit.append("2. **For debugging/comparison:** Use `force_left_to_right=True` to guarantee no look-ahead")
    audit.append("3. **Wire threshold through VLMEvalKit:** Update `vlmeval_qwen_vl_mdm.py` to pass `threshold` param")
    audit.append("4. **Verify current experiments:** Check if `use_block_causal_mask` is enabled in the best model config\n")
    
    return "\n".join(audit)


def generate_report(results: Dict, early_name: str, late_name: str) -> str:
    """Generate a markdown report from analysis results."""
    report = [f"# DocVQA Checkpoint Regression Analysis\n"]
    report.append(f"**Comparing:** {early_name} → {late_name}\n")
    
    # Summary
    report.append("## Summary\n")
    report.append(f"| Metric | Value |")
    report.append(f"|--------|-------|")
    report.append(f"| Total Samples | {results['total_samples']} |")
    report.append(f"| Early Checkpoint Accuracy | {results['early_accuracy']:.2f}% |")
    report.append(f"| Late Checkpoint Accuracy | {results['late_accuracy']:.2f}% |")
    report.append(f"| **Accuracy Drop** | **{results['accuracy_drop']:.2f}%** |")
    report.append(f"| Regressions (early✓ → late✗) | {results['regression_count']} |")
    report.append(f"| Improvements (early✗ → late✓) | {results['improvement_count']} |")
    report.append(f"| Both Correct | {results['both_correct']} |")
    report.append(f"| Both Wrong | {results['both_wrong']} |")
    report.append("")
    
    # Error category breakdown
    report.append("## Error Category Breakdown\n")
    report.append("Categories (from meeting note 1/15/2026):\n")
    report.append("- **noise_coverage**: 答案本来在输出里，但后续生成了噪声/冗余 token，导致 evaluator/parse 失败")
    report.append("- **total_derail**: 模型失焦/跑偏，输出整体偏离正确答案")
    report.append("- **yes_no_flip**: yes/no 类问题，答案被翻转")
    report.append("- **partial_forget**: 早期正确但晚期只有部分正确")
    report.append("- **other**: 其他类型\n")
    
    if results['error_categories']:
        report.append("| Error Category | Count | Percentage |")
        report.append("|----------------|-------|------------|")
        total = sum(results['error_categories'].values())
        for cat, count in sorted(results['error_categories'].items(), key=lambda x: -x[1]):
            pct = count / total * 100 if total > 0 else 0
            report.append(f"| {cat} | {count} | {pct:.1f}% |")
    else:
        report.append("No regression cases found.\n")
    report.append("")
    
    # Sample cases by category
    report.append("## Sample Regression Cases\n")
    
    # Group by category
    cases_by_cat = {}
    for case in results['regression_cases']:
        cat = case['error_category']
        if cat not in cases_by_cat:
            cases_by_cat[cat] = []
        cases_by_cat[cat].append(case)
    
    for cat in ['noise_coverage', 'total_derail', 'yes_no_flip', 'partial_forget', 'other']:
        if cat not in cases_by_cat:
            continue
        
        report.append(f"### {cat.upper()} ({len(cases_by_cat[cat])} cases)\n")
        
        # Show top 5 examples per category
        for i, case in enumerate(cases_by_cat[cat][:5]):
            report.append(f"#### Example {i+1}\n")
            report.append(f"- **Index**: {case['index']}")
            report.append(f"- **Question**: {case['question'][:200]}{'...' if len(case['question']) > 200 else ''}")
            report.append(f"- **Ground Truth**: `{case['gt_string'][:100]}{'...' if len(case['gt_string']) > 100 else ''}`")
            report.append(f"- **Early Prediction** (score={case['early_score']:.2f}): `{case['early_prediction'][:150]}{'...' if len(case['early_prediction']) > 150 else ''}`")
            report.append(f"- **Late Prediction** (score={case['late_score']:.2f}): `{case['late_prediction'][:150]}{'...' if len(case['late_prediction']) > 150 else ''}`")
            report.append(f"- **Score Drop**: {case['score_drop']:.2f}")
            if case['error_details']:
                details_str = ", ".join(f"{k}={v:.2f}" if isinstance(v, float) else f"{k}={v}" 
                                        for k, v in case['error_details'].items() 
                                        if not isinstance(v, bool) or v)
                report.append(f"- **Details**: {details_str}")
            report.append(f"- **Image**: `{case['image_path']}`\n")
        report.append("")
    
    # Top-K largest drops
    report.append("## Top-10 Largest Score Drops\n")
    for i, case in enumerate(results['regression_cases'][:10]):
        report.append(f"{i+1}. **{case['error_category']}** | drop={case['score_drop']:.2f} | Q: {case['question'][:80]}...")
    report.append("")
    
    # Add look-ahead bias audit
    report.append(generate_lookahead_audit())
    
    return "\n".join(report)


def export_results(results: Dict, output_dir: Path, early_name: str, late_name: str):
    """Export results to json, csv, and markdown files."""
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # JSON export
    json_file = output_dir / "regressions.json"
    with open(json_file, 'w', encoding='utf-8') as f:
        json.dump({
            'comparison': f'{early_name} → {late_name}',
            'summary': {
                'total_samples': results['total_samples'],
                'early_accuracy': results['early_accuracy'],
                'late_accuracy': results['late_accuracy'],
                'accuracy_drop': results['accuracy_drop'],
                'regression_count': results['regression_count'],
                'improvement_count': results['improvement_count'],
                'error_categories': results['error_categories'],
            },
            'cases': results['regression_cases'],
        }, f, indent=2, ensure_ascii=False)
    print(f"[INFO] Exported JSON: {json_file}")
    
    # CSV export
    csv_file = output_dir / "regressions.csv"
    df = pd.DataFrame([
        {
            'index': c['index'],
            'question': c['question'][:200],
            'ground_truth': c['gt_string'],
            'early_prediction': c['early_prediction'][:200],
            'late_prediction': c['late_prediction'][:200],
            'early_score': c['early_score'],
            'late_score': c['late_score'],
            'score_drop': c['score_drop'],
            'error_category': c['error_category'],
            'image_path': c['image_path'],
        }
        for c in results['regression_cases']
    ])
    df.to_csv(csv_file, index=False, encoding='utf-8')
    print(f"[INFO] Exported CSV: {csv_file}")
    
    # Markdown report
    report = generate_report(results, early_name, late_name)
    md_file = output_dir / "report.md"
    with open(md_file, 'w', encoding='utf-8') as f:
        f.write(report)
    print(f"[INFO] Exported Report: {md_file}")


def main():
    parser = argparse.ArgumentParser(
        description="Analyze DocVQA checkpoint regressions (early correct → late wrong)"
    )
    parser.add_argument(
        "--early_ckpt_dir",
        type=str,
        default=str(DEFAULT_EARLY_CKPT),
        help="Path to early checkpoint DocVQA evaluation directory"
    )
    parser.add_argument(
        "--late_ckpt_dir",
        type=str,
        default=str(DEFAULT_LATE_CKPT),
        help="Path to late checkpoint DocVQA evaluation directory"
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="output/docvqa_regression_analysis",
        help="Output directory for analysis results"
    )
    parser.add_argument(
        "--early_name",
        type=str,
        default=None,
        help="Display name for early checkpoint (default: derived from path)"
    )
    parser.add_argument(
        "--late_name",
        type=str,
        default=None,
        help="Display name for late checkpoint (default: derived from path)"
    )
    
    args = parser.parse_args()
    
    early_dir = Path(args.early_ckpt_dir)
    late_dir = Path(args.late_ckpt_dir)
    output_dir = Path(args.output_dir)
    
    # Derive checkpoint names from paths
    early_name = args.early_name or early_dir.parent.name if early_dir.parent.name.startswith("checkpoint") else early_dir.name
    late_name = args.late_name or late_dir.parent.name if late_dir.parent.name.startswith("checkpoint") else late_dir.name
    
    print(f"[INFO] Early checkpoint: {early_dir}")
    print(f"[INFO] Late checkpoint: {late_dir}")
    print(f"[INFO] Output directory: {output_dir}")
    print(f"[INFO] Comparing: {early_name} → {late_name}")
    print()
    
    # Load results
    early_df = load_results(early_dir)
    late_df = load_results(late_dir)
    
    if early_df is None or late_df is None:
        print("[ERROR] Failed to load one or both checkpoint results. Exiting.")
        return 1
    
    print(f"[INFO] Early checkpoint samples: {len(early_df)}")
    print(f"[INFO] Late checkpoint samples: {len(late_df)}")
    print()
    
    # Analyze regressions
    results = analyze_regressions(early_df, late_df)
    
    print(f"[RESULT] Early accuracy: {results['early_accuracy']:.2f}%")
    print(f"[RESULT] Late accuracy: {results['late_accuracy']:.2f}%")
    print(f"[RESULT] Accuracy drop: {results['accuracy_drop']:.2f}%")
    print(f"[RESULT] Regressions: {results['regression_count']}")
    print(f"[RESULT] Improvements: {results['improvement_count']}")
    print(f"[RESULT] Error categories: {results['error_categories']}")
    print()
    
    # Export results
    export_results(results, output_dir, early_name, late_name)
    
    print("\n[DONE] Analysis complete!")
    return 0


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

