#!/usr/bin/env python3
"""Benchmark-scale hard/soft agreement audit on all public MATH solutions.

The paper-specific Atomic Dependency Graph annotations and learned factuality
scorer are not public.  This audit therefore makes a deliberately scoped proxy:
each gold solution is split into reasoning steps, the steps form a chain DAG,
and a deterministic structural risk score is assigned to each step.  The DCF
prediction relaxation itself follows the equations audited for Theorem 3.2.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import re
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq


DATASET = "https://huggingface.co/datasets/qwedsacf/competition_math"
DATASET_REVISION = "e839825f9ec5c6cfa585c654a59610969ec13993"
ALPHAS = tuple(round(i / 100, 2) for i in range(1, 11))
THRESHOLDS = tuple(round(i / 20, 2) for i in range(1, 20))
TEMPERATURE = 0.01
ANCESTOR_WEIGHT = 1.0


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


def split_steps(solution: str) -> list[str]:
    text = re.sub(r"\s+", " ", solution).strip()
    parts = [part.strip() for part in re.split(r"(?<=[.!?])\s+|(?=\\boxed)|(?=\\Rightarrow)", text)]
    parts = [part for part in parts if len(part) >= 8]
    if not parts:
        parts = [text or "empty"]
    # Bound pathological TeX fragments while retaining full benchmark coverage.
    if len(parts) > 24:
        parts = parts[:23] + [" ".join(parts[23:])]
    return parts


def structural_risk(step: str, row_index: int, step_index: int) -> float:
    """Deterministic, non-learned proxy risk in [0.05, 0.95]."""
    tokens = re.findall(r"[A-Za-z0-9]+|\\[A-Za-z]+", step)
    equation_density = min(1.0, (step.count("=") + step.count("\\")) / 12.0)
    length_term = min(1.0, len(tokens) / 80.0)
    conclusion_bonus = 1.0 if "boxed" in step or "therefore" in step.lower() else 0.0
    digest = hashlib.blake2b(
        f"{row_index}:{step_index}:{step}".encode(), digest_size=8
    ).digest()
    jitter = int.from_bytes(digest, "big") / (2**64 - 1)
    score = 0.10 + 0.42 * length_term + 0.28 * equation_density - 0.10 * conclusion_bonus + 0.20 * jitter
    return min(0.95, max(0.05, score))


def ancestors_for_chain(nodes: int) -> list[list[int]]:
    return [list(range(node)) for node in range(nodes)]


def log_sigmoid(value: float) -> float:
    if value >= 0:
        return -math.log1p(math.exp(-value))
    return value - math.log1p(math.exp(value))


def hard_prediction(risks: list[float], ancestors: list[list[int]], tau_alpha: float) -> list[int]:
    tau_star = max(tau for tau in THRESHOLDS if tau < tau_alpha)
    return [
        int(risks[node] <= tau_star and all(risks[parent] <= tau_star for parent in ancestors[node]))
        for node in range(len(risks))
    ]


def soft_prediction(risks: list[float], ancestors: list[list[int]], tau_alpha: float) -> tuple[list[int], float]:
    # Theorem 3.2 schedule: a=1, b=3, beta=T^-a, tau_z=T^(ab).
    beta = 1.0 / TEMPERATURE
    tau_z = TEMPERATURE**3
    log_weights = [
        beta * tau + log_sigmoid((tau_alpha - tau - math.sqrt(tau_z)) / tau_z)
        for tau in THRESHOLDS
    ]
    maximum = max(log_weights)
    weights = np.exp(np.asarray(log_weights) - maximum)
    weights /= weights.sum()
    memberships = []
    for tau in THRESHOLDS:
        probabilities = [
            1.0 / (1.0 + math.exp(max(-700.0, min(700.0, (risk - tau) / TEMPERATURE))))
            for risk in risks
        ]
        coherent = []
        for node, parent_nodes in enumerate(ancestors):
            indices = [node] + parent_nodes
            node_weights = [1.0] + [ANCESTOR_WEIGHT] * len(parent_nodes)
            denominator = sum(node_weights)
            coherent.append(
                math.exp(
                    sum(
                        weight * math.log(max(probabilities[index], 1e-300))
                        for index, weight in zip(indices, node_weights)
                    )
                    / denominator
                )
            )
        memberships.append(coherent)
    mixed = weights @ np.asarray(memberships)
    return [int(value >= 0.5) for value in mixed], float(weights.max())


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", type=Path, default=Path("competition_math.parquet"))
    parser.add_argument("--output", type=Path, default=Path("outputs/math_agreement_v1"))
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)

    table = pq.read_table(args.data, columns=["solution", "level", "type"])
    records = table.to_pylist()
    examples = []
    all_risks = []
    step_counts = []
    for row_index, record in enumerate(records):
        steps = split_steps(record["solution"])
        risks = [structural_risk(step, row_index, step_index) for step_index, step in enumerate(steps)]
        examples.append((risks, ancestors_for_chain(len(risks)), record["level"], record["type"]))
        all_risks.extend(risks)
        step_counts.append(len(steps))

    output_rows = []
    for alpha in ALPHAS:
        # Quantile calibration is performed once over all public solution-step risks.
        tau_alpha = float(np.quantile(np.asarray(all_risks), 1.0 - alpha, method="higher"))
        # Ensure the prediction gate has at least one feasible threshold and is
        # not exactly on a grid point, as in the theorem's separated regime.
        tau_alpha = min(0.999, max(0.051, tau_alpha + 1e-7))
        exact_matches = 0
        node_matches = 0
        node_total = 0
        hard_retained = 0
        soft_retained = 0
        minimum_selected_mass = 1.0
        for risks, ancestors, _, _ in examples:
            hard = hard_prediction(risks, ancestors, tau_alpha)
            soft, selected_mass = soft_prediction(risks, ancestors, tau_alpha)
            exact_matches += int(hard == soft)
            node_matches += sum(int(left == right) for left, right in zip(hard, soft))
            node_total += len(hard)
            hard_retained += sum(hard)
            soft_retained += sum(soft)
            minimum_selected_mass = min(minimum_selected_mass, selected_mass)
        output_rows.append(
            {
                "alpha": alpha,
                "calibrated_tau_alpha": tau_alpha,
                "examples": len(examples),
                "claim_nodes": node_total,
                "exact_set_matches": exact_matches,
                "exact_set_agreement_percent": 100.0 * exact_matches / len(examples),
                "node_matches": node_matches,
                "node_agreement_percent": 100.0 * node_matches / node_total,
                "hard_retained_nodes": hard_retained,
                "soft_retained_nodes": soft_retained,
                "minimum_selected_threshold_mass": minimum_selected_mass,
            }
        )

    with (args.output / "agreement_by_alpha.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(output_rows[0]))
        writer.writeheader()
        writer.writerows(output_rows)
    report = {
        "claim": "DCF soft/hard agreement across alpha=0.01..0.10",
        "scope": "benchmark-scale proxy; not the unreleased paper scorer/ADGs",
        "dataset": DATASET,
        "dataset_revision": DATASET_REVISION,
        "dataset_file_sha256": sha256(args.data),
        "examples": len(examples),
        "solution_claim_nodes": sum(step_counts),
        "step_count": {
            "minimum": min(step_counts),
            "median": float(np.median(step_counts)),
            "maximum": max(step_counts),
        },
        "graph_construction": "sentence/TeX-step split; sequential chain with transitive ancestors",
        "risk_proxy": "deterministic structural features plus pinned BLAKE2 jitter; not learned factuality",
        "threshold_grid": list(THRESHOLDS),
        "temperature": TEMPERATURE,
        "ancestor_weight": ANCESTOR_WEIGHT,
        "rows": output_rows,
        "minimum_exact_set_agreement_percent": min(row["exact_set_agreement_percent"] for row in output_rows),
        "maximum_exact_set_agreement_percent": max(row["exact_set_agreement_percent"] for row in output_rows),
        "minimum_node_agreement_percent": min(row["node_agreement_percent"] for row in output_rows),
        "maximum_node_agreement_percent": max(row["node_agreement_percent"] for row in output_rows),
    }
    (args.output / "summary.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()