#!/usr/bin/env python3
"""LLM-as-a-Judge evaluation — rubric-based scoring with bias mitigation.

Demonstrates: building a structured rubric, scoring multiple axes independently,
swapping position to detect position bias, and computing agreement metrics.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any


@dataclass
class RubricLevel:
    score: int
    label: str
    description: str


class Rubric:
    """A scoring rubric with concrete, mutually exclusive levels."""

    def __init__(self, name: str, levels: list[RubricLevel]):
        self.name = name
        self.levels = sorted(levels, key=lambda l: l.score)

    def get_prompt(self) -> str:
        lines = [f"Score the response on **{self.name}** using this rubric:"]
        for level in self.levels:
            lines.append(f"  Score {level.score}: {level.label} — {level.description}")
        return "\n".join(lines)


# ── Pre-built rubrics ───────────────────────────────────────────────────

FAITHFULNESS_RUBRIC = Rubric("faithfulness", [
    RubricLevel(1, "Hallucinated", "Response contains claims not supported by the source text."),
    RubricLevel(2, "Partially faithful", "Most claims are supported but some are unsupported or distorted."),
    RubricLevel(3, "Faithful", "All claims are directly supported by the source text."),
])

HELPFULNESS_RUBRIC = Rubric("helpfulness", [
    RubricLevel(1, "Unhelpful", "Response is irrelevant, evasive, or nonsensical."),
    RubricLevel(2, "Partially helpful", "Response addresses the question but misses key aspects."),
    RubricLevel(3, "Helpful", "Response fully addresses the question with clear, actionable information."),
    RubricLevel(4, "Very helpful", "Response is complete, well-structured, and anticipates follow-up needs."),
])


# ── Mock judge (replace with real LLM API call) ─────────────────────────

def mock_judge_score(rubric: Rubric, response: str) -> int:
    """Simulate an LLM judge scoring a response against a rubric."""
    # Heuristic: longer responses often score higher (this IS the length bias)
    length = len(response)
    if length > 200:
        return rubric.levels[-1].score
    elif length > 80:
        return rubric.levels[len(rubric.levels) // 2].score
    return rubric.levels[0].score


# ── Bias detection: position swap test ──────────────────────────────────

def position_swap_test(responses: list[tuple[str, str]], rubric: Rubric) -> dict[str, Any]:
    """Compare scores when response order is swapped to detect position bias."""
    results = {"swapped_verdicts": 0, "total": len(responses), "details": []}

    for resp_a, resp_b in responses:
        # Score in original order
        score_a1 = mock_judge_score(rubric, resp_a)
        score_b1 = mock_judge_score(rubric, resp_b)
        winner_1 = "A" if score_a1 > score_b1 else "B" if score_b1 > score_a1 else "tie"

        # Score in swapped order
        score_b2 = mock_judge_score(rubric, resp_b)
        score_a2 = mock_judge_score(rubric, resp_a)
        winner_2 = "A" if score_a2 > score_b2 else "B" if score_b2 > score_a2 else "tie"

        if winner_1 != winner_2:
            results["swapped_verdicts"] += 1

        results["details"].append({
            "order_AB": {"score_A": score_a1, "score_B": score_b1, "winner": winner_1},
            "order_BA": {"score_A": score_a2, "score_B": score_b2, "winner": winner_2},
            "consistent": winner_1 == winner_2,
        })

    results["swap_rate"] = results["swapped_verdicts"] / max(results["total"], 1)
    return results


# ── Compute agreement ───────────────────────────────────────────────────

def compute_kappa(judge_scores: list[int], human_scores: list[int]) -> float:
    """Simple Cohen's kappa for ordinal categories."""
    from collections import Counter

    n = len(judge_scores)
    if n == 0:
        return 1.0

    # Observed agreement
    po = sum(1 for j, h in zip(judge_scores, human_scores) if j == h) / n

    # Expected agreement
    j_counts = Counter(judge_scores)
    h_counts = Counter(human_scores)
    pe = sum(j_counts[s] * h_counts[s] for s in set(judge_scores + human_scores)) / (n * n)

    return (po - pe) / (1 - pe) if pe < 1 else 1.0


# ── Demo ────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("=== Rubric Definition ===")
    print(FAITHFULNESS_RUBRIC.get_prompt())
    print(f"\n{HELPFULNESS_RUBRIC.get_prompt()}")

    print("\n=== Position Bias Test ===")
    responses = [
        (
            "The sky is blue because of Rayleigh scattering of sunlight.",
            "Blue light scatters more than other colors due to its shorter wavelength, "
            "which is why the sky appears blue during the day. This phenomenon is called Rayleigh scattering.",
        ),
        (
            "Python is good.",
            "Python is a versatile programming language known for its readability, "
            "extensive standard library, and strong ecosystem for web development, data science, and automation.",
        ),
    ]

    test_results = position_swap_test(responses, HELPFULNESS_RUBRIC)
    for i, detail in enumerate(test_results["details"]):
        print(f"\n  Pair {i+1}:")
        print(f"    Order A->B: {detail['order_AB']}")
        print(f"    Order B->A: {detail['order_BA']}")
        print(f"    Consistent: {detail['consistent']}")

    swap_rate = test_results["swap_rate"]
    print(f"\n  Position swap disagreement rate: {swap_rate:.0%}")
    if swap_rate > 0.05:
        print("  ⚠ Position bias detected (>5% swap rate)")

    print("\n=== Agreement Example ===")
    # Simulated judge vs human scores
    judge = [3, 2, 4, 1, 3, 2, 4, 3, 2, 4]
    human = [3, 2, 3, 1, 3, 2, 4, 3, 1, 4]
    kappa = compute_kappa(judge, human)
    print(f"  Cohen's kappa: {kappa:.3f}")
    print(f"  Interpretation: {'Excellent' if kappa > 0.75 else 'Good' if kappa > 0.6 else 'Moderate' if kappa > 0.4 else 'Poor'}")

    print("\n✓ LLM-as-Judge evaluation demo complete.")