"""
babilong_eval.report — Results aggregation, printing, and JSON export.

Provides:
  - BabilongResult: container for all evaluation results
  - SampleResult: per-sample result
  - print_report(): formatted console output
  - save_results(): JSON export for further analysis
"""

from __future__ import annotations

import json
from collections import defaultdict
from dataclasses import dataclass, field, asdict
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING

if TYPE_CHECKING:
    from babilong_eval.harness import BabilongConfig


# ---------------------------------------------------------------------------
# Result types
# ---------------------------------------------------------------------------

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

    Attributes:
        task: Task name (qa1–qa5).
        condition: Preprocessing condition (full/hcp/random).
        retain_ratio: Fraction of context retained (1.0 for full).
        output: Raw LLM output.
        target: Expected answer.
        question: The question asked.
        correct: Whether the output matched the target.
        context_len: Original context length in chars.
        filtered_len: Filtered context length in chars.
    """
    task: str
    condition: str
    retain_ratio: float
    output: str
    target: str
    question: str
    correct: bool
    context_len: int
    filtered_len: int


@dataclass
class BabilongResult:
    """Container for a full BABILong evaluation run.

    Attributes:
        config: The evaluation configuration used.
        samples: All per-sample results.
    """
    config: Any  # BabilongConfig (avoid circular import)
    samples: List[SampleResult] = field(default_factory=list)

    @property
    def n_samples(self) -> int:
        return len(self.samples)

    @property
    def overall_accuracy(self) -> float:
        if not self.samples:
            return 0.0
        return sum(1 for s in self.samples if s.correct) / len(self.samples)


# ---------------------------------------------------------------------------
# Aggregation helpers
# ---------------------------------------------------------------------------

def _group_results(
    samples: List[SampleResult],
) -> Dict[str, Dict[str, Dict[float, List[SampleResult]]]]:
    """Group results by task → condition → retain_ratio.

    Returns:
        Nested dict: task → condition → ratio → [SampleResult, ...]
    """
    grouped: Dict[str, Dict[str, Dict[float, List[SampleResult]]]] = defaultdict(
        lambda: defaultdict(lambda: defaultdict(list))
    )
    for s in samples:
        grouped[s.task][s.condition][s.retain_ratio].append(s)
    return grouped


def _compute_accuracy(results: List[SampleResult]) -> float:
    if not results:
        return 0.0
    return sum(1 for r in results if r.correct) / len(results)


def _compute_compression(results: List[SampleResult]) -> float:
    """Average compression ratio (filtered_len / context_len)."""
    if not results:
        return 1.0
    ratios = [r.filtered_len / r.context_len for r in results if r.context_len > 0]
    return sum(ratios) / len(ratios) if ratios else 1.0


# ---------------------------------------------------------------------------
# Console report
# ---------------------------------------------------------------------------

def print_report(result: BabilongResult) -> None:
    """Print a formatted console report of BABILong evaluation results.

    Shows:
      1. Per-task accuracy comparison across conditions
      2. Compression ratios for filtered conditions
      3. Statistical summary (HCP advantage over baselines)
    """
    grouped = _group_results(result.samples)

    print("\n" + "=" * 80)
    print("BABILong Evaluation Report")
    print("=" * 80)

    # Overall summary
    print(f"\nTotal samples: {result.n_samples}")
    print(f"Overall accuracy: {result.overall_accuracy:.1%}")

    # Per-task breakdown
    for task in sorted(grouped.keys()):
        print(f"\n{'-' * 70}")
        print(f"TASK: {task.upper()}")
        print(f"{'-' * 70}")
        print(f"  {'Condition':<20} {'Ratio':<8} {'Accuracy':<12} {'Compression':<12} {'N':<6}")

        conditions = grouped[task]
        for condition in sorted(conditions.keys()):
            for ratio in sorted(conditions[condition].keys()):
                samples = conditions[condition][ratio]
                acc = _compute_accuracy(samples)
                comp = _compute_compression(samples)
                n = len(samples)
                label = f"{condition}" if condition == "full" else f"{condition}"
                ratio_str = f"{ratio:.0%}" if ratio < 1.0 else "100%"
                print(f"  {label:<20} {ratio_str:<8} {acc:<12.1%} {comp:<12.1%} {n:<6}")

    # Comparative summary
    print(f"\n{'=' * 80}")
    print("COMPARATIVE SUMMARY")
    print(f"{'=' * 80}")

    for task in sorted(grouped.keys()):
        conditions = grouped[task]

        # Full baseline accuracy
        full_results = []
        for ratio_results in conditions.get("full", {}).values():
            full_results.extend(ratio_results)
        full_acc = _compute_accuracy(full_results)

        # HCP accuracies per ratio
        hcp_accs = {}
        for ratio, samples in conditions.get("hcp", {}).items():
            hcp_accs[ratio] = _compute_accuracy(samples)

        # Random accuracies per ratio
        random_accs = {}
        for ratio, samples in conditions.get("random", {}).items():
            random_accs[ratio] = _compute_accuracy(samples)

        print(f"\n  {task.upper()}:")
        print(f"    Full baseline:     {full_acc:.1%}")

        for ratio in sorted(set(list(hcp_accs.keys()) + list(random_accs.keys()))):
            hcp_val = hcp_accs.get(ratio, 0)
            rand_val = random_accs.get(ratio, 0)
            hcp_diff = hcp_val - full_acc
            rand_diff = rand_val - full_acc
            hcp_vs_rand = hcp_val - rand_val

            ratio_str = f"{ratio:.0%}"
            print(f"    HCP@{ratio_str:<6}         {hcp_val:.1%}  (vs full: {hcp_diff:+.1%}, vs random: {hcp_vs_rand:+.1%})")
            print(f"    Random@{ratio_str:<6}      {rand_val:.1%}  (vs full: {rand_diff:+.1%})")

    print(f"\n{'=' * 80}")


# ---------------------------------------------------------------------------
# JSON export
# ---------------------------------------------------------------------------

def save_results(result: BabilongResult, filepath: str = "babilong_results.json") -> None:
    """Save evaluation results to JSON.

    Args:
        result: The BabilongResult to save.
        filepath: Output file path.
    """
    data = {
        "timestamp": datetime.now().isoformat(),
        "config": {
            "tasks": list(result.config.tasks),
            "lengths": list(result.config.lengths),
            "n_samples": result.config.n_samples,
            "retain_ratios": list(result.config.retain_ratios),
            "conditions": list(result.config.conditions),
            "model": result.config.model,
        },
        "summary": {
            "total_samples": result.n_samples,
            "overall_accuracy": round(result.overall_accuracy, 4),
        },
        "per_task": {},
        "samples": [],
    }

    # Per-task summary
    grouped = _group_results(result.samples)
    for task in sorted(grouped.keys()):
        task_data: Dict[str, Any] = {}
        for condition in sorted(grouped[task].keys()):
            for ratio in sorted(grouped[task][condition].keys()):
                samples = grouped[task][condition][ratio]
                key = f"{condition}@{ratio}" if ratio < 1.0 else condition
                task_data[key] = {
                    "accuracy": round(_compute_accuracy(samples), 4),
                    "compression": round(_compute_compression(samples), 4),
                    "n_samples": len(samples),
                }
        data["per_task"][task] = task_data

    # Individual samples (for detailed analysis)
    for s in result.samples:
        data["samples"].append({
            "task": s.task,
            "condition": s.condition,
            "retain_ratio": s.retain_ratio,
            "correct": s.correct,
            "target": s.target,
            "output": s.output,
            "context_len": s.context_len,
            "filtered_len": s.filtered_len,
        })

    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

    print(f"\nResults saved to {filepath}")
