#!/usr/bin/env python3
"""Self-authored, bounded reaggregation for the NinueNAODD evidence run.

This file is intentionally input-only: it never downloads data, imports paper
or peer code, starts a subprocess, or writes outside the explicit --out path.
It verifies the exact release artifacts before reading them.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import time
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Iterable


EXPECTED = {
    "screened-00000-of-00001.parquet": {
        "bytes": 91141,
        "sha256": "cf8fb398c421291cad2331fd5b20a9baadf0d8e20fa43fce2813eb8665900da4",
        "rows": 1920,
    },
    "unscreened-00000-of-00001.parquet": {
        "bytes": 91114,
        "sha256": "c1f0199a69d3a0d56eea6230b2bd582e69aee7c2cb4b1e7ee8a52b10647fe7ee",
        "rows": 1921,
    },
    "feedback_dataset.parquet": {
        "bytes": 1888747,
        "sha256": "372d432b76340a793a1f9be053de73b3d6d23ecd9ccb4058023ce7d4569002e5",
        "rows": 257280,
    },
}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def verify_release(path: Path) -> dict[str, Any]:
    spec = EXPECTED.get(path.name)
    if spec is None:
        raise ValueError(f"unexpected input basename: {path.name}")
    actual_bytes = path.stat().st_size
    actual_sha = sha256_file(path)
    if actual_bytes != spec["bytes"] or actual_sha != spec["sha256"]:
        raise ValueError(
            f"hash gate failed for {path}: bytes={actual_bytes} sha256={actual_sha}"
        )
    return {"path": str(path), "bytes": actual_bytes, "sha256": actual_sha}


def read_parquet(path: Path) -> list[dict[str, Any]]:
    try:
        import pyarrow.parquet as parquet
    except ImportError as exc:  # pragma: no cover - dependency is checked by permit
        raise RuntimeError("pyarrow is required; no network fallback is allowed") from exc
    table = parquet.read_table(path)
    return table.to_pylist()


def finite_number(value: Any) -> bool:
    return isinstance(value, (int, float)) and math.isfinite(float(value))


def clic_observations(rows: Iterable[dict[str, Any]]) -> tuple[list[dict[str, Any]], Counter]:
    """Flatten each unambiguous two-item 2AFC record into one observation."""
    observations: list[dict[str, Any]] = []
    exclusions: Counter = Counter()
    for row in rows:
        ratings = row.get("ratings") or []
        by_condition: dict[str, dict[str, Any]] = {}
        for rating in ratings:
            condition = rating.get("condition")
            score = rating.get("score")
            if condition is not None and finite_number(score):
                by_condition[str(condition)] = rating
        if len(by_condition) != 2:
            exclusions["not_exactly_two_conditions"] += 1
            continue
        items = sorted(by_condition)
        left, right = by_condition[items[0]], by_condition[items[1]]
        left_score, right_score = float(left["score"]), float(right["score"])
        if left_score == right_score:
            exclusions["tied_scores"] += 1
            continue
        observations.append(
            {
                "rater": str(row.get("rater_id")),
                "left": items[0],
                "right": items[1],
                "winner": items[0] if left_score > right_score else items[1],
            }
        )
    return observations, exclusions


def humaine_observations(rows: Iterable[dict[str, Any]]) -> tuple[list[dict[str, Any]], Counter]:
    """Keep explicit A/B choices and exclude ties or malformed rows."""
    observations: list[dict[str, Any]] = []
    exclusions: Counter = Counter()
    for row in rows:
        a, b, choice = row.get("model_a"), row.get("model_b"), row.get("choice")
        if not a or not b or a == b:
            exclusions["missing_or_identical_pair"] += 1
            continue
        normalized = str(choice).strip().upper()
        if normalized not in {"A", "B"}:
            exclusions["tie_or_unknown_choice"] += 1
            continue
        observations.append(
            {
                "rater": str(row.get("conversation_id", "row")),
                "left": str(a),
                "right": str(b),
                "winner": str(a) if normalized == "A" else str(b),
            }
        )
    return observations, exclusions


def pair_table(observations: Iterable[dict[str, Any]]) -> tuple[list[str], list[dict[str, Any]]]:
    items = sorted({x for row in observations for x in (row["left"], row["right"])})
    index = {item: n for n, item in enumerate(items)}
    out = []
    for row in observations:
        i, j = sorted((index[row["left"]], index[row["right"]]))
        winner_index = index[row["winner"]]
        out.append({"rater": row["rater"], "i": i, "j": j, "winner": winner_index})
    return items, out


def fit_bt(items: list[str], observations: list[dict[str, Any]], iterations: int = 200) -> tuple[list[float], int]:
    """Self-authored Bradley-Terry MM fit with a unit ridge-like stabilizer."""
    k = len(items)
    lam = [1.0] * k
    pair_counts: dict[tuple[int, int], list[int]] = defaultdict(lambda: [0, 0])
    for row in observations:
        pair = (row["i"], row["j"])
        pair_counts[pair][0] += 1
        if row["winner"] == row["i"]:
            pair_counts[pair][1] += 1
    for step in range(1, iterations + 1):
        updated = [1.0] * k
        for i in range(k):
            wins = 0.0
            denom = 1.0
            for (a, b), (n, wins_a) in pair_counts.items():
                if i == a:
                    wins += wins_a
                    denom += n / (lam[a] + lam[b])
                elif i == b:
                    wins += n - wins_a
                    denom += n / (lam[a] + lam[b])
            updated[i] = max(wins / denom, 1e-12)
        scale = sum(updated) / max(k, 1)
        updated = [x / scale for x in updated]
        if max(abs(math.log(a) - math.log(b)) for a, b in zip(lam, updated)) < 1e-8:
            return updated, step
        lam = updated
    return lam, iterations


def fit_bbq(
    items: list[str],
    observations: list[dict[str, Any]],
    alpha: float = 10.0,
    beta: float = 2.0,
    shape: float = 5.0,
    rate: float = 0.1,
    iterations: int = 200,
) -> dict[str, Any]:
    """Self-authored implementation of the paper's Eq. 11--13 updates."""
    k = len(items)
    lam = [1.0] * k
    raters = sorted({row["rater"] for row in observations})
    q = {rater: 0.8 for rater in raters}
    history: list[float] = []
    start = time.perf_counter()
    for step in range(1, iterations + 1):
        gammas: list[float] = []
        for row in observations:
            i, j = row["i"], row["j"]
            p_i = lam[i] / (lam[i] + lam[j])
            p_win = p_i if row["winner"] == i else 1.0 - p_i
            quality = min(max(q[row["rater"]], 1e-12), 1.0 - 1e-12)
            numerator = quality * p_win
            gammas.append(numerator / (numerator + (1.0 - quality) * 0.5))

        by_rater: dict[str, list[float]] = defaultdict(list)
        for row, gamma in zip(observations, gammas):
            by_rater[row["rater"]].append(gamma)
        for rater in raters:
            n = len(by_rater[rater])
            q[rater] = min(max((sum(by_rater[rater]) + alpha - 1.0) / (n + alpha + beta - 2.0), 0.0), 1.0)

        updated = [0.0] * k
        denominators = [rate] * k
        for row, gamma in zip(observations, gammas):
            i, j = row["i"], row["j"]
            updated[row["winner"]] += gamma
            weight = gamma / (lam[i] + lam[j])
            denominators[i] += weight
            denominators[j] += weight
        updated = [(wins + shape - 1.0) / denom for wins, denom in zip(updated, denominators)]
        scale = sum(updated) / max(k, 1)
        updated = [max(x / scale, 1e-12) for x in updated]

        # The observed-data mixture log-likelihood is a diagnostic only; the
        # paper's MAP objective also includes the stated Gamma/Beta priors.
        likelihood = 0.0
        for row in observations:
            i, j = row["i"], row["j"]
            p_i = lam[i] / (lam[i] + lam[j])
            p_win = p_i if row["winner"] == i else 1.0 - p_i
            quality = q[row["rater"]]
            likelihood += math.log(max(quality * p_win + (1.0 - quality) * 0.5, 1e-300))
        history.append(likelihood)
        if max(abs(math.log(a) - math.log(b)) for a, b in zip(lam, updated)) < 1e-8:
            lam = updated
            break
        lam = updated
    return {
        "items": items,
        "skills": dict(zip(items, lam)),
        "q": q,
        "iterations": step,
        "seconds": time.perf_counter() - start,
        "observed_log_likelihood": history[-1] if history else None,
        "likelihood_history": history,
    }


def pearson(xs: list[float], ys: list[float]) -> float:
    if len(xs) != len(ys) or len(xs) < 2:
        raise ValueError("Pearson correlation needs equal-length vectors of at least two values")
    mean_x, mean_y = sum(xs) / len(xs), sum(ys) / len(ys)
    dx, dy = [x - mean_x for x in xs], [y - mean_y for y in ys]
    denom = math.sqrt(sum(x * x for x in dx) * sum(y * y for y in dy))
    if denom == 0.0:
        raise ValueError("constant vector in Pearson correlation")
    return sum(x * y for x, y in zip(dx, dy)) / denom


def rater_agreement(observations: list[dict[str, Any]], skills: dict[str, float]) -> dict[str, float]:
    agreement: dict[str, list[bool]] = defaultdict(list)
    for row in observations:
        left_score, right_score = skills[row["left"]], skills[row["right"]]
        predicted = row["left"] if left_score >= right_score else row["right"]
        agreement[row["rater"]].append(predicted == row["winner"])
    return {rater: sum(values) / len(values) for rater, values in agreement.items() if values}


def permutation_negative_control(q: dict[str, float], agreement: dict[str, float]) -> dict[str, Any]:
    """Deterministic label-shuffle control for the q/agreement correlation."""
    common = sorted(set(q) & set(agreement))
    if len(common) < 3:
        return {"n": len(common), "status": "insufficient_common_raters"}
    q_values = [q[r] for r in common]
    observed = pearson(q_values, [agreement[r] for r in common])
    shuffled = q_values[1:] + q_values[:1]
    shuffled_corr = pearson(shuffled, [agreement[r] for r in common])
    return {
        "n": len(common),
        "observed_pearson": observed,
        "one_step_cyclic_shuffle_pearson": shuffled_corr,
        "status": "pass" if observed > shuffled_corr else "inspect",
    }


def summarize_clic(path: Path) -> dict[str, Any]:
    rows = read_parquet(path)
    expected_rows = EXPECTED[path.name]["rows"]
    if len(rows) != expected_rows:
        raise ValueError(f"row-count gate failed for {path.name}: {len(rows)} != {expected_rows}")
    observations, exclusions = clic_observations(rows)
    items, encoded = pair_table(observations)
    bbq = fit_bbq(items, encoded)
    bt_skills, bt_iterations = fit_bt(items, encoded)
    bt_map = dict(zip(items, bt_skills))
    agreement = rater_agreement(observations, bbq["skills"])
    return {
        "rows": len(rows),
        "raters": len({str(row.get("rater_id")) for row in rows}),
        "items": len(items),
        "usable_2afc_observations": len(encoded),
        "exclusions": dict(exclusions),
        "bt_iterations": bt_iterations,
        "bbq_iterations": bbq["iterations"],
        "bbq_seconds": bbq["seconds"],
        "rater_q_pearson_to_consensus_agreement": pearson(
            [bbq["q"][r] for r in sorted(set(bbq["q"]) & set(agreement))],
            [agreement[r] for r in sorted(set(bbq["q"]) & set(agreement))],
        ),
        "permutation_negative_control": permutation_negative_control(bbq["q"], agreement),
        "bt_top_item": max(bt_map, key=bt_map.get) if bt_map else None,
        "bbq_top_item": max(bbq["skills"], key=bbq["skills"].get) if bbq["skills"] else None,
    }


def summarize_humaine(path: Path) -> dict[str, Any]:
    rows = read_parquet(path)
    expected_rows = EXPECTED[path.name]["rows"]
    if len(rows) != expected_rows:
        raise ValueError(f"row-count gate failed for {path.name}: {len(rows)} != {expected_rows}")
    observations, exclusions = humaine_observations(rows)
    items, encoded = pair_table(observations)
    start = time.perf_counter()
    bt_skills, bt_iterations = fit_bt(items, encoded)
    bt_seconds = time.perf_counter() - start
    bbq = fit_bbq(items, encoded)
    metrics = Counter(str(row.get("metric")) for row in rows)
    return {
        "rows": len(rows),
        "metric_rows": dict(metrics),
        "usable_non_tie_observations": len(encoded),
        "excluded_rows": dict(exclusions),
        "items": len(items),
        "raters_or_conversation_ids": len({row["rater"] for row in observations}),
        "bt_iterations": bt_iterations,
        "bt_seconds": bt_seconds,
        "bbq_iterations": bbq["iterations"],
        "bbq_seconds": bbq["seconds"],
        "bbq_over_bt_seconds_ratio": bbq["seconds"] / bt_seconds if bt_seconds else None,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--screened", type=Path, required=True)
    parser.add_argument("--unscreened", type=Path, required=True)
    parser.add_argument("--humaine-feedback", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()

    inputs = [args.screened, args.unscreened, args.humaine_feedback]
    release_manifest = [verify_release(path) for path in inputs]
    result = {
        "release_inputs": release_manifest,
        "screened": summarize_clic(args.screened),
        "unscreened": summarize_clic(args.unscreened),
        "humaine_feedback": summarize_humaine(args.humaine_feedback),
        "controls": {
            "exact_row_gates": True,
            "all_skills_positive": True,
            "cpu_threads_requested": 2,
            "author_or_peer_code_used": False,
        },
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")


if __name__ == "__main__":
    main()