#!/usr/bin/env python3
# SPDX-FileCopyrightText: 2026 本项目 large-language-model-editor
# SPDX-License-Identifier: CC-BY-4.0
# Creator: large-language-model-editor
"""Deterministically generate and independently verify the Chapter 18 LLM case."""

from __future__ import annotations

import argparse
import json
import math
import pathlib
from typing import Any

VOCABULARY = ["花", "雨", "猫", "。"]
POSITION_INPUTS = [
    {"position": 1, "weights": [8, 4, 2, 1], "target": "花"},
    {"position": 2, "weights": [1, 2, 1, 6], "target": "。"},
]
ABS_TOLERANCE = 1e-12
REL_TOLERANCE = 0.0
ROOT_KEYS = {
    "schema", "license", "creator", "numerical_policy", "vocabulary",
    "next_token_positions", "sequence_metrics", "decoding",
}


def _distribution(weights: list[int], temperature: float = 1.0) -> list[float]:
    """Normalize positive unscaled weights after temperature scaling."""
    if not math.isfinite(temperature) or temperature <= 0.0:
        raise ValueError("temperature must be finite and greater than zero")
    if not weights or any(isinstance(w, bool) or not isinstance(w, (int, float)) or not math.isfinite(w) or w <= 0 for w in weights):
        raise ValueError("weights must be a non-empty list of positive finite numbers")
    # Shift log-weights before exponentiating: softmax is invariant to this common
    # offset, and the shift prevents overflow for every finite T > 0.
    log_weights = [math.log(float(weight)) for weight in weights]
    pivot = max(log_weights)
    scaled = [math.exp((log_weight - pivot) / temperature) for log_weight in log_weights]
    denominator = math.fsum(scaled)
    # binary64 can underflow mathematically positive tail probabilities to zero at
    # extremely small positive T (as T approaches 0 from above); the published fixture is restricted to T=1.
    return [weight / denominator for weight in scaled]


def _rank(probabilities: list[float]) -> list[int]:
    """Rank descending, resolving equal probabilities by vocabulary order."""
    return sorted(range(len(probabilities)), key=lambda index: (-probabilities[index], index))


def _truncate(probabilities: list[float], method: str, threshold: float, exact_weights: list[int] | None = None) -> dict[str, Any]:
    ranking = _rank(probabilities)
    if method == "top_k":
        if isinstance(threshold, bool) or not isinstance(threshold, int) or not 1 <= threshold <= len(probabilities):
            raise ValueError("top-k requires integer 1 <= k <= vocabulary size")
        selected = ranking[:threshold]
    elif method == "top_p":
        if not isinstance(threshold, (int, float)) or isinstance(threshold, bool) or not 0.0 < threshold <= 1.0:
            raise ValueError("top-p requires 0 < p <= 1")
        selected = []
        cumulative = 0.0
        cumulative_units = 0
        total_units = sum(exact_weights) if exact_weights is not None else None
        for index in ranking:
            selected.append(index)
            cumulative += probabilities[index]
            if exact_weights is not None:
                cumulative_units += exact_weights[index]
                p_num, p_den = float(threshold).as_integer_ratio()
                reached = cumulative_units * p_den >= p_num * total_units
            else:
                reached = cumulative >= threshold
            if reached:
                break
    else:
        raise ValueError(f"unknown truncation method: {method}")
    retained_mass = sum(probabilities[index] for index in selected)
    selected_set = set(selected)
    renormalized = [probability / retained_mass if index in selected_set else 0.0 for index, probability in enumerate(probabilities)]
    cumulative = 0.0
    ranked_prefix = []
    for index in selected:
        cumulative += probabilities[index]
        ranked_prefix.append({
            "token": VOCABULARY[index],
            "original_probability": probabilities[index],
            "cumulative_probability": cumulative,
        })
    return {
        "selected_tokens_in_rank_order": [VOCABULARY[index] for index in selected],
        "ranked_prefix": ranked_prefix,
        "retained_mass_before_renormalization": retained_mass,
        "renormalized_probabilities_in_vocabulary_order": renormalized,
    }


def build_case() -> dict[str, Any]:
    """Build all published values from the fixed primitive inputs."""
    positions = []
    token_nll = []
    for primitive in POSITION_INPUTS:
        probabilities = _distribution(primitive["weights"], 1.0)
        target_index = VOCABULARY.index(primitive["target"])
        nll = -math.log(probabilities[target_index])
        token_nll.append(nll)
        positions.append({
            **primitive,
            "weight_sum": sum(primitive["weights"]),
            "probabilities": probabilities,
            "target_probability": probabilities[target_index],
            "target_nll_nats": nll,
        })

    total_nll = sum(token_nll)
    cross_entropy = total_nll / len(token_nll)
    cases = []
    for position in positions:
        probabilities = position["probabilities"]
        cases.append({
            "position": position["position"],
            "source_probabilities_in_vocabulary_order": probabilities,
            "top_k": {"k": 2, **_truncate(probabilities, "top_k", 2)},
            "top_p": {"p": 0.75, **_truncate(probabilities, "top_p", 0.75, position["weights"])},
        })

    return {
        "schema": 1,
        "license": "CC BY 4.0",
        "creator": "本项目 large-language-model-editor",
        "numerical_policy": {
            "logarithm": "natural",
            "loss_unit": "nat",
            "absolute_tolerance": ABS_TOLERANCE,
            "relative_tolerance": REL_TOLERANCE,
            "conditions": [
                "Vocabulary order is fixed as [花, 雨, 猫, 。].",
                "Every weight must be finite and strictly positive; temperature T must be finite and T > 0.",
                "Targets must occur in the fixed vocabulary; sequence CE is arithmetic mean token NLL.",
                "Top-k requires integer 1 <= k <= vocabulary size.",
                "Top-p requires 0 < p <= 1 and selects the smallest descending-probability prefix whose cumulative mass is at least p.",
                "Equal probabilities are ordered by fixed vocabulary order; truncate first, then renormalize retained mass to one.",
                "The published JSON fixture is exact at every field; 1e-12 applies only to independent cross-runtime recomputation diagnostics for derived reported values, never to fixture identity, control fields, or top-p prefix decisions.",
            ],
        },
        "vocabulary": VOCABULARY,
        "next_token_positions": positions,
        "sequence_metrics": {
            "targets": [primitive["target"] for primitive in POSITION_INPUTS],
            "token_nll_nats": token_nll,
            "total_nll_nats": total_nll,
            "cross_entropy_mean_nats": cross_entropy,
            "perplexity": math.exp(cross_entropy),
        },
        "decoding": {
            "temperature": 1.0,
            "top_k_k": 2,
            "top_p_p": 0.75,
            "operation_order": ["temperature_scale", "normalize", "rank", "truncate", "renormalize"],
            "tie_break": "fixed_vocabulary_order",
            "cases": cases,
        },
    }


def _mapping(value: Any, path: str, keys: set[str]) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise ValueError(f"{path}: expected object")
    if set(value) != keys:
        raise ValueError(f"{path}: key mismatch; missing={sorted(keys - set(value))}, extra={sorted(set(value) - keys)}")
    return value


def _number(value: Any, path: str) -> float:
    if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
        raise ValueError(f"{path}: expected finite number")
    return float(value)


def _close(actual: Any, expected: float, path: str) -> None:
    value = _number(actual, path)
    if not math.isclose(value, expected, rel_tol=REL_TOLERANCE, abs_tol=ABS_TOLERANCE):
        raise ValueError(f"{path}: stored {value!r}, recomputed {expected!r}")


def _verify_tolerance_boundaries() -> None:
    below = math.nextafter(ABS_TOLERANCE, 0.0)
    above = math.nextafter(ABS_TOLERANCE, math.inf)
    if not math.isclose(ABS_TOLERANCE, 0.0, rel_tol=REL_TOLERANCE, abs_tol=ABS_TOLERANCE):
        raise ValueError("tolerance boundary must be inclusive")
    if not math.isclose(below, 0.0, rel_tol=REL_TOLERANCE, abs_tol=ABS_TOLERANCE):
        raise ValueError("nextafter below tolerance rejected")
    if math.isclose(above, 0.0, rel_tol=REL_TOLERANCE, abs_tol=ABS_TOLERANCE):
        raise ValueError("nextafter above tolerance accepted")


def _same_tree(actual: Any, expected: Any, path: str = "$") -> None:
    """Seal inputs, labels, order, metadata, and deterministic binary64 fixtures."""
    if isinstance(expected, dict):
        obj = _mapping(actual, path, set(expected))
        for key, expected_value in expected.items():
            _same_tree(obj[key], expected_value, f"{path}.{key}")
    elif isinstance(expected, list):
        if not isinstance(actual, list) or len(actual) != len(expected):
            raise ValueError(f"{path}: expected list length {len(expected)}")
        for index, (actual_value, expected_value) in enumerate(zip(actual, expected)):
            _same_tree(actual_value, expected_value, f"{path}[{index}]")
    elif isinstance(expected, float):
        value = _number(actual, path)
        if value != expected:
            raise ValueError(f"{path}: stored {value!r}, fixture {expected!r}")
    elif type(actual) is not type(expected) or actual != expected:
        raise ValueError(f"{path}: expected {expected!r}, got {actual!r}")


def verify_recomputation(payload: Any) -> None:
    """Structurally validate, independently recompute, and fixture-seal a payload."""
    _verify_tolerance_boundaries()
    root = _mapping(payload, "$", ROOT_KEYS)
    if root["vocabulary"] != VOCABULARY:
        raise ValueError("$.vocabulary: fixed vocabulary or order changed")
    positions = root["next_token_positions"]
    if not isinstance(positions, list) or len(positions) != 2:
        raise ValueError("$.next_token_positions: expected two positions")
    recomputed_nll = []
    for index, (stored, primitive) in enumerate(zip(positions, POSITION_INPUTS)):
        path = f"$.next_token_positions[{index}]"
        row = _mapping(stored, path, {"position", "weights", "target", "weight_sum", "probabilities", "target_probability", "target_nll_nats"})
        if row["position"] != primitive["position"] or row["weights"] != primitive["weights"] or row["target"] != primitive["target"]:
            raise ValueError(f"{path}: fixed primitive input changed")
        probabilities = _distribution(row["weights"], 1.0)
        if not isinstance(row["probabilities"], list) or len(row["probabilities"]) != len(VOCABULARY):
            raise ValueError(f"{path}.probabilities: expected four values")
        for j, probability in enumerate(probabilities):
            _close(row["probabilities"][j], probability, f"{path}.probabilities[{j}]")
        _close(sum(_number(v, f"{path}.probabilities") for v in row["probabilities"]), 1.0, f"{path}.probability_sum")
        target_probability = probabilities[VOCABULARY.index(row["target"])]
        nll = -math.log(target_probability)
        _close(row["target_probability"], target_probability, f"{path}.target_probability")
        _close(row["target_nll_nats"], nll, f"{path}.target_nll_nats")
        recomputed_nll.append(nll)

    metrics = root["sequence_metrics"]
    if not isinstance(metrics, dict):
        raise ValueError("$.sequence_metrics: expected object")
    total = sum(recomputed_nll)
    mean = total / len(recomputed_nll)
    for i, value in enumerate(recomputed_nll):
        _close(metrics["token_nll_nats"][i], value, f"$.sequence_metrics.token_nll_nats[{i}]")
    _close(metrics["total_nll_nats"], total, "$.sequence_metrics.total_nll_nats")
    _close(metrics["cross_entropy_mean_nats"], mean, "$.sequence_metrics.cross_entropy_mean_nats")
    _close(metrics["perplexity"], math.exp(mean), "$.sequence_metrics.perplexity")

    decoding = root["decoding"]
    if decoding.get("temperature") != 1.0 or decoding.get("top_k_k") != 2 or decoding.get("top_p_p") != 0.75:
        raise ValueError("$.decoding: T, k, or p changed")
    if not isinstance(decoding.get("cases"), list) or len(decoding["cases"]) != 2:
        raise ValueError("$.decoding.cases: expected two position cases")
    for index, probabilities in enumerate(_distribution(p["weights"]) for p in POSITION_INPUTS):
        case = decoding["cases"][index]
        _same_tree(case["top_k"], {"k": 2, **_truncate(probabilities, "top_k", 2)}, f"$.decoding.cases[{index}].top_k")
        _same_tree(case["top_p"], {"p": 0.75, **_truncate(probabilities, "top_p", 0.75)}, f"$.decoding.cases[{index}].top_p")

    _same_tree(root, build_case())


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--json", action="store_true", help="print deterministic JSON")
    parser.add_argument("--check", type=pathlib.Path, help="recompute and verify a published JSON file")
    args = parser.parse_args()
    case = build_case()
    verify_recomputation(case)
    if args.check is not None:
        try:
            verify_recomputation(json.loads(args.check.read_text(encoding="utf-8")))
        except (OSError, json.JSONDecodeError, KeyError, IndexError, TypeError, ValueError, ZeroDivisionError) as error:
            raise SystemExit(f"CH18 DATA FAILED: {args.check}: {error}") from error
    if args.json:
        print(json.dumps(case, ensure_ascii=False, indent=2))
    else:
        print("CH18 DATA PASSED: natural-log CE/NLL/PPL and T=1 top-k/top-p truncation were recomputed structurally.")


if __name__ == "__main__":
    main()