#!/usr/bin/env python3
"""
A/B Test Results Comparator for RAGAS Evaluation

Compares two RAGAS evaluation result files to determine if a change
(e.g., orphan connections) improved or degraded retrieval quality.

Usage:
    python lightrag/evaluation/compare_results.py baseline.json experiment.json
    python lightrag/evaluation/compare_results.py results_a.json results_b.json --output comparison.json
"""

import argparse
import json
import math
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any


@dataclass
class MetricComparison:
    """Comparison of a single metric between two runs."""

    metric_name: str
    baseline_value: float
    experiment_value: float
    absolute_change: float
    relative_change_percent: float
    improved: bool
    significant: bool  # > 5% change


def safe_float(value: Any, default: float = 0.0) -> float:
    """Safely convert a value to float, handling NaN."""
    if value is None:
        return default
    try:
        f = float(value)
        if math.isnan(f):
            return default
        return f
    except (ValueError, TypeError):
        return default


def compare_metrics(baseline: dict, experiment: dict) -> list[MetricComparison]:
    """
    Compare metrics between baseline and experiment.

    Args:
        baseline: Benchmark stats from baseline run
        experiment: Benchmark stats from experiment run

    Returns:
        List of MetricComparison objects
    """
    comparisons = []

    baseline_avg = baseline.get('average_metrics', {})
    experiment_avg = experiment.get('average_metrics', {})

    metrics_to_compare = [
        ('faithfulness', 'Faithfulness'),
        ('answer_relevance', 'Answer Relevance'),
        ('context_recall', 'Context Recall'),
        ('context_precision', 'Context Precision'),
        ('ragas_score', 'RAGAS Score'),
    ]

    for metric_key, metric_name in metrics_to_compare:
        b_val = safe_float(baseline_avg.get(metric_key, 0))
        e_val = safe_float(experiment_avg.get(metric_key, 0))

        abs_change = e_val - b_val
        rel_change = (abs_change / b_val * 100) if b_val > 0 else 0

        comparisons.append(
            MetricComparison(
                metric_name=metric_name,
                baseline_value=b_val,
                experiment_value=e_val,
                absolute_change=abs_change,
                relative_change_percent=rel_change,
                improved=abs_change > 0,
                significant=abs(rel_change) > 5,  # > 5% is significant
            )
        )

    return comparisons


def analyze_results(baseline_path: Path, experiment_path: Path) -> dict:
    """
    Perform comprehensive A/B analysis.

    Args:
        baseline_path: Path to baseline results JSON
        experiment_path: Path to experiment results JSON

    Returns:
        Analysis results dictionary
    """
    # Load results
    with open(baseline_path) as f:
        baseline = json.load(f)
    with open(experiment_path) as f:
        experiment = json.load(f)

    baseline_stats = baseline.get('benchmark_stats', {})
    experiment_stats = experiment.get('benchmark_stats', {})

    # Compare metrics
    comparisons = compare_metrics(baseline_stats, experiment_stats)

    # Calculate overall verdict
    improvements = sum(1 for c in comparisons if c.improved)
    regressions = sum(1 for c in comparisons if not c.improved and c.absolute_change != 0)
    significant_improvements = sum(1 for c in comparisons if c.improved and c.significant)
    significant_regressions = sum(1 for c in comparisons if not c.improved and c.significant)

    # Determine verdict
    ragas_comparison = next((c for c in comparisons if c.metric_name == 'RAGAS Score'), None)

    if ragas_comparison:
        if ragas_comparison.improved and ragas_comparison.significant:
            verdict = 'SIGNIFICANT_IMPROVEMENT'
            verdict_description = f'RAGAS Score improved by {ragas_comparison.relative_change_percent:.1f}%'
        elif ragas_comparison.improved:
            verdict = 'MINOR_IMPROVEMENT'
            verdict_description = f'RAGAS Score slightly improved by {ragas_comparison.relative_change_percent:.1f}%'
        elif ragas_comparison.significant:
            verdict = 'SIGNIFICANT_REGRESSION'
            verdict_description = f'RAGAS Score regressed by {abs(ragas_comparison.relative_change_percent):.1f}%'
        elif ragas_comparison.absolute_change == 0:
            verdict = 'NO_CHANGE'
            verdict_description = 'No measurable difference between runs'
        else:
            verdict = 'MINOR_REGRESSION'
            verdict_description = (
                f'RAGAS Score slightly regressed by {abs(ragas_comparison.relative_change_percent):.1f}%'
            )
    else:
        verdict = 'UNKNOWN'
        verdict_description = 'Could not determine RAGAS score comparison'

    return {
        'analysis_timestamp': datetime.now().isoformat(),
        'baseline_file': str(baseline_path),
        'experiment_file': str(experiment_path),
        'verdict': verdict,
        'verdict_description': verdict_description,
        'summary': {
            'metrics_improved': improvements,
            'metrics_regressed': regressions,
            'significant_improvements': significant_improvements,
            'significant_regressions': significant_regressions,
        },
        'metrics': [
            {
                'name': c.metric_name,
                'baseline': round(c.baseline_value, 4),
                'experiment': round(c.experiment_value, 4),
                'change': round(c.absolute_change, 4),
                'change_percent': round(c.relative_change_percent, 2),
                'improved': c.improved,
                'significant': c.significant,
            }
            for c in comparisons
        ],
        'baseline_summary': {
            'total_tests': baseline_stats.get('total_tests', 0),
            'successful_tests': baseline_stats.get('successful_tests', 0),
            'success_rate': baseline_stats.get('success_rate', 0),
        },
        'experiment_summary': {
            'total_tests': experiment_stats.get('total_tests', 0),
            'successful_tests': experiment_stats.get('successful_tests', 0),
            'success_rate': experiment_stats.get('success_rate', 0),
        },
    }


def print_comparison_report(analysis: dict):
    """Print a formatted comparison report to stdout."""
    print('=' * 70)
    print('A/B TEST COMPARISON REPORT')
    print('=' * 70)
    print(f'Baseline:    {analysis["baseline_file"]}')
    print(f'Experiment:  {analysis["experiment_file"]}')
    print('-' * 70)

    # Verdict
    verdict = analysis['verdict']
    verdict_icon = {
        'SIGNIFICANT_IMPROVEMENT': 'PASS',
        'MINOR_IMPROVEMENT': 'PASS',
        'NO_CHANGE': '~',
        'MINOR_REGRESSION': 'WARN',
        'SIGNIFICANT_REGRESSION': 'FAIL',
        'UNKNOWN': '?',
    }.get(verdict, '?')

    print(f'\n[{verdict_icon}] VERDICT: {verdict}')
    print(f'    {analysis["verdict_description"]}')

    # Metrics table
    print('\n' + '-' * 70)
    print(f'{"Metric":<20} {"Baseline":>10} {"Experiment":>10} {"Change":>10} {"Status":>10}')
    print('-' * 70)

    for metric in analysis['metrics']:
        name = metric['name']
        baseline = f'{metric["baseline"]:.4f}'
        experiment = f'{metric["experiment"]:.4f}'

        change = metric['change']
        change_pct = metric['change_percent']
        if change > 0:
            change_str = f'+{change:.4f}'
            status = f'+{change_pct:.1f}%'
        elif change < 0:
            change_str = f'{change:.4f}'
            status = f'{change_pct:.1f}%'
        else:
            change_str = '0.0000'
            status = '0.0%'

        if metric['significant']:
            status = f'[UP] {status}' if metric['improved'] else f'[DOWN] {status}'
        else:
            status = f'      {status}'

        print(f'{name:<20} {baseline:>10} {experiment:>10} {change_str:>10} {status:>10}')

    print('-' * 70)

    # Summary
    summary = analysis['summary']
    print(f'\nSummary: {summary["metrics_improved"]} improved, {summary["metrics_regressed"]} regressed')
    print(
        f'         {summary["significant_improvements"]} significant improvements, {summary["significant_regressions"]} significant regressions'
    )

    # Test counts
    b_summary = analysis['baseline_summary']
    e_summary = analysis['experiment_summary']
    print(
        f'\nBaseline:    {b_summary["successful_tests"]}/{b_summary["total_tests"]} tests ({b_summary["success_rate"]:.1f}% success)'
    )
    print(
        f'Experiment:  {e_summary["successful_tests"]}/{e_summary["total_tests"]} tests ({e_summary["success_rate"]:.1f}% success)'
    )

    print('=' * 70)


def main():
    parser = argparse.ArgumentParser(
        description='Compare RAGAS evaluation results from two runs',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Compare baseline vs experiment
  python lightrag/evaluation/compare_results.py baseline.json experiment.json

  # Save comparison to file
  python lightrag/evaluation/compare_results.py baseline.json experiment.json --output comparison.json

  # Compare with/without orphan connections
  python lightrag/evaluation/compare_results.py results_without_orphans.json results_with_orphans.json
        """,
    )

    parser.add_argument(
        'baseline',
        type=str,
        help='Path to baseline results JSON file',
    )

    parser.add_argument(
        'experiment',
        type=str,
        help='Path to experiment results JSON file',
    )

    parser.add_argument(
        '--output',
        '-o',
        type=str,
        default=None,
        help='Output path for comparison JSON (optional)',
    )

    args = parser.parse_args()

    baseline_path = Path(args.baseline)
    experiment_path = Path(args.experiment)

    # Validate files exist
    if not baseline_path.exists():
        print(f'Error: Baseline file not found: {baseline_path}')
        sys.exit(1)
    if not experiment_path.exists():
        print(f'Error: Experiment file not found: {experiment_path}')
        sys.exit(1)

    # Run analysis
    analysis = analyze_results(baseline_path, experiment_path)

    # Print report
    print_comparison_report(analysis)

    # Save to file if requested
    if args.output:
        output_path = Path(args.output)
        with open(output_path, 'w') as f:
            json.dump(analysis, f, indent=2)
        print(f'\nComparison saved to: {output_path}')

    # Exit with status based on verdict
    if analysis['verdict'] in ('SIGNIFICANT_REGRESSION',):
        sys.exit(1)
    sys.exit(0)


if __name__ == '__main__':
    main()