#!/usr/bin/env python3
# SPDX-FileCopyrightText: 2026 本项目 generative-ai-editor
# SPDX-License-Identifier: CC-BY-4.0
# Creator: generative-ai-editor
"""Deterministically generate and independently verify Chapter 17 numeric cases."""

from __future__ import annotations

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


ROOT_KEYS = {"schema", "license", "creator", "nll", "discrete_elbo", "diffusion_step", "temperature"}


def build_case() -> dict[str, Any]:
    """Build the published case from its declared primitive inputs."""
    tokens = ["A", "A", "A", "B"]
    token_probability = {"A": 0.75, "B": 0.25}
    token_nll = [-math.log(token_probability[token]) for token in tokens]
    nll_total = sum(token_nll)
    nll_mean = nll_total / len(tokens)

    prior = [0.5, 0.5]
    likelihood = [0.2, 0.8]
    posterior_q = [0.25, 0.75]
    expected_log_likelihood = sum(q * math.log(p) for q, p in zip(posterior_q, likelihood))
    kl_q_prior = sum(q * math.log(q / p) for q, p in zip(posterior_q, prior))
    elbo = expected_log_likelihood - kl_q_prior
    evidence = sum(p * likelihood_z for p, likelihood_z in zip(prior, likelihood))
    log_evidence = math.log(evidence)
    exact_posterior = [p * likelihood_z / evidence for p, likelihood_z in zip(prior, likelihood)]
    kl_gap = sum(q * math.log(q / p) for q, p in zip(posterior_q, exact_posterior))

    x0 = 1.0
    beta = 0.2
    epsilon = -0.5
    epsilon_hat = -0.4
    alpha = 1.0 - beta
    x1 = math.sqrt(alpha) * x0 + math.sqrt(beta) * epsilon
    mse = (epsilon - epsilon_hat) ** 2
    x0_hat = (x1 - math.sqrt(beta) * epsilon_hat) / math.sqrt(alpha)

    logits = [math.log(4.0), 0.0]
    temperatures = [0.5, 1.0, 2.0]
    rows = []
    fractions = [[16, 17], [4, 5], [2, 3]]
    for temperature, fraction in zip(temperatures, fractions):
        scaled = [value / temperature for value in logits]
        pivot = max(scaled)
        weights = [math.exp(value - pivot) for value in scaled]
        denominator = sum(weights)
        probabilities = [weight / denominator for weight in weights]
        rows.append(
            {
                "temperature": temperature,
                "probabilities": probabilities,
                "first_probability_fraction": fraction,
            }
        )

    return {
        "schema": 1,
        "license": "CC BY 4.0",
        "creator": "本项目 generative-ai-editor",
        "nll": {
            "sequence": tokens,
            "token_probability": token_probability,
            "token_nll": token_nll,
            "total_nll": nll_total,
            "mean_nll": nll_mean,
            "perplexity": math.exp(nll_mean),
        },
        "discrete_elbo": {
            "latent_states": [0, 1],
            "prior": prior,
            "likelihood_of_observation": likelihood,
            "q": posterior_q,
            "expected_log_likelihood": expected_log_likelihood,
            "kl_q_prior": kl_q_prior,
            "elbo": elbo,
            "evidence": evidence,
            "log_evidence": log_evidence,
            "exact_posterior": exact_posterior,
            "kl_gap": kl_gap,
        },
        "diffusion_step": {
            "x0": x0,
            "beta": beta,
            "alpha": alpha,
            "epsilon": epsilon,
            "epsilon_hat": epsilon_hat,
            "x1": x1,
            "epsilon_mse": mse,
            "x0_hat": x0_hat,
        },
        "temperature": {
            "logits": logits,
            "classes": ["第一类", "第二类"],
            "rows": rows,
        },
    }


def _mapping(value: Any, path: str, keys: set[str]) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise ValueError(f"{path}: expected object")
    actual_keys = set(value)
    if actual_keys != keys:
        missing = sorted(keys - actual_keys)
        extra = sorted(actual_keys - keys)
        raise ValueError(f"{path}: key mismatch; missing={missing}, extra={extra}")
    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 _same_number(actual: Any, recomputed: float, path: str) -> None:
    """Demand the exact deterministic binary64 result, rejecting even subtle drift."""
    value = _number(actual, path)
    if value != recomputed:
        raise ValueError(f"{path}: stored {value!r}, recomputed {recomputed!r}")


def _same_tree(actual: Any, expected: Any, path: str = "$") -> None:
    """Recursively enforce every key, list position, literal, and numeric fixture."""
    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):
            got = len(actual) if isinstance(actual, list) else type(actual).__name__
            raise ValueError(f"{path}: expected list length {len(expected)}, got {got}")
        for index, (actual_value, expected_value) in enumerate(zip(actual, expected)):
            _same_tree(actual_value, expected_value, f"{path}[{index}]")
    elif isinstance(expected, float):
        _same_number(actual, expected, path)
    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:
    """Validate structure, recompute each result from payload inputs, then seal fixtures."""
    root = _mapping(payload, "$", ROOT_KEYS)

    nll = _mapping(
        root["nll"],
        "$.nll",
        {"sequence", "token_probability", "token_nll", "total_nll", "mean_nll", "perplexity"},
    )
    sequence = nll["sequence"]
    probabilities = _mapping(nll["token_probability"], "$.nll.token_probability", {"A", "B"})
    if not isinstance(sequence, list) or not sequence:
        raise ValueError("$.nll.sequence: expected non-empty list")
    recomputed_token_nll = [-math.log(_number(probabilities[token], f"$.nll.token_probability.{token}")) for token in sequence]
    _same_tree(nll["token_nll"], recomputed_token_nll, "$.nll.token_nll")
    recomputed_total = sum(recomputed_token_nll)
    _same_number(nll["total_nll"], recomputed_total, "$.nll.total_nll")
    recomputed_mean = recomputed_total / len(sequence)
    _same_number(nll["mean_nll"], recomputed_mean, "$.nll.mean_nll")
    _same_number(nll["perplexity"], math.exp(recomputed_mean), "$.nll.perplexity")

    elbo = _mapping(
        root["discrete_elbo"],
        "$.discrete_elbo",
        {"latent_states", "prior", "likelihood_of_observation", "q", "expected_log_likelihood", "kl_q_prior", "elbo", "evidence", "log_evidence", "exact_posterior", "kl_gap"},
    )
    prior = [_number(v, f"$.discrete_elbo.prior[{i}]") for i, v in enumerate(elbo["prior"])]
    likelihood = [_number(v, f"$.discrete_elbo.likelihood_of_observation[{i}]") for i, v in enumerate(elbo["likelihood_of_observation"])]
    q = [_number(v, f"$.discrete_elbo.q[{i}]") for i, v in enumerate(elbo["q"])]
    if not (len(prior) == len(likelihood) == len(q) == 2):
        raise ValueError("$.discrete_elbo: prior, likelihood, and q must each contain two values")
    expected_ll = sum(qi * math.log(li) for qi, li in zip(q, likelihood))
    kl_prior = sum(qi * math.log(qi / pi) for qi, pi in zip(q, prior))
    evidence = sum(pi * li for pi, li in zip(prior, likelihood))
    posterior = [pi * li / evidence for pi, li in zip(prior, likelihood)]
    _same_number(elbo["expected_log_likelihood"], expected_ll, "$.discrete_elbo.expected_log_likelihood")
    _same_number(elbo["kl_q_prior"], kl_prior, "$.discrete_elbo.kl_q_prior")
    _same_number(elbo["elbo"], expected_ll - kl_prior, "$.discrete_elbo.elbo")
    _same_number(elbo["evidence"], evidence, "$.discrete_elbo.evidence")
    _same_number(elbo["log_evidence"], math.log(evidence), "$.discrete_elbo.log_evidence")
    _same_tree(elbo["exact_posterior"], posterior, "$.discrete_elbo.exact_posterior")
    gap = sum(qi * math.log(qi / post) for qi, post in zip(q, posterior))
    _same_number(elbo["kl_gap"], gap, "$.discrete_elbo.kl_gap")
    identity_gap = _number(elbo["log_evidence"], "$.discrete_elbo.log_evidence") - _number(elbo["elbo"], "$.discrete_elbo.elbo")
    if not math.isclose(identity_gap, gap, rel_tol=0.0, abs_tol=1e-15):
        raise ValueError(f"$.discrete_elbo.identity_gap: evidence minus ELBO {identity_gap!r}, KL {gap!r}")

    diffusion = _mapping(
        root["diffusion_step"],
        "$.diffusion_step",
        {"x0", "beta", "alpha", "epsilon", "epsilon_hat", "x1", "epsilon_mse", "x0_hat"},
    )
    x0 = _number(diffusion["x0"], "$.diffusion_step.x0")
    beta = _number(diffusion["beta"], "$.diffusion_step.beta")
    epsilon = _number(diffusion["epsilon"], "$.diffusion_step.epsilon")
    epsilon_hat = _number(diffusion["epsilon_hat"], "$.diffusion_step.epsilon_hat")
    alpha = 1.0 - beta
    x1 = math.sqrt(alpha) * x0 + math.sqrt(beta) * epsilon
    _same_number(diffusion["alpha"], alpha, "$.diffusion_step.alpha")
    _same_number(diffusion["x1"], x1, "$.diffusion_step.x1")
    _same_number(diffusion["epsilon_mse"], (epsilon - epsilon_hat) ** 2, "$.diffusion_step.epsilon_mse")
    _same_number(diffusion["x0_hat"], (x1 - math.sqrt(beta) * epsilon_hat) / math.sqrt(alpha), "$.diffusion_step.x0_hat")

    temperature = _mapping(root["temperature"], "$.temperature", {"logits", "classes", "rows"})
    logits = [_number(v, f"$.temperature.logits[{i}]") for i, v in enumerate(temperature["logits"])]
    rows = temperature["rows"]
    if not isinstance(rows, list):
        raise ValueError("$.temperature.rows: expected list")
    for index, row_value in enumerate(rows):
        row = _mapping(row_value, f"$.temperature.rows[{index}]", {"temperature", "probabilities", "first_probability_fraction"})
        t = _number(row["temperature"], f"$.temperature.rows[{index}].temperature")
        scaled = [value / t for value in logits]
        pivot = max(scaled)
        weights = [math.exp(value - pivot) for value in scaled]
        softmax = [weight / sum(weights) for weight in weights]
        _same_tree(row["probabilities"], softmax, f"$.temperature.rows[{index}].probabilities")
        fraction = row["first_probability_fraction"]
        if not isinstance(fraction, list) or len(fraction) != 2 or fraction[1] == 0:
            raise ValueError(f"$.temperature.rows[{index}].first_probability_fraction: invalid fraction")
        _same_number(softmax[0], fraction[0] / fraction[1], f"$.temperature.rows[{index}].fraction_identity")

    # The recomputation above guards formula integrity; this recursive fixture seal also
    # rejects altered primitive inputs, labels, metadata, order, or extra structure.
    _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:
            payload = json.loads(args.check.read_text(encoding="utf-8"))
            verify_recomputation(payload)
        except (OSError, json.JSONDecodeError, KeyError, TypeError, ValueError, ZeroDivisionError) as error:
            raise SystemExit(f"CH17 DATA FAILED: {args.check}: {error}") from error

    if args.json:
        print(json.dumps(case, ensure_ascii=False, indent=2))
    else:
        print("CH17 DATA PASSED: NLL, discrete ELBO, one diffusion step, and temperature softmax recomputed structurally.")


if __name__ == "__main__":
    main()