#!/usr/bin/env python3
"""
OMEGA-ALPHA REFERENCE IMPLEMENTATION
====================================

Author: Matthew Alan Newman (@NVE024)
Prepared as the executable companion to the master synthesis preprint.
Version: 1.0.0
Python: 3.10+
Dependencies: Python standard library only

Scope
-----
This module implements the finite, exactly checkable parts of the Omega-Alpha
framework:

1. Verification-level evaluation of
       phi = (1 + sqrt(5))/2,
       lambda = phi/(phi + pi),
       kappa = -ln(lambda),
   using Decimal arithmetic and Machin's formula for pi.

2. The exact rational protocol realization
       S = 10^15,
       Q = 339949771344778,
       N = 23.

3. The witnessed integer contraction
       d_next = floor(Q*d/S),
       rho    = Q*d - S*d_next,
   and exact reverse reconstruction from the ordered remainder ledger.

4. Reversible 15-bit MSB-first byte folding with LSB zero padding.

5. The finite monotone codebooks
       T(b) = F^23((b+2)S)
   on the byte and 16-bit seed domains.

6. The corrected Ramanujan ellipse-multiplier coordinate table.

The module does not claim that floating-point or Decimal evaluation is symbolic
exact arithmetic. Exactness claims in this implementation concern integer
arithmetic after S, Q, and N have been fixed.
"""

from __future__ import annotations

from dataclasses import dataclass
from decimal import Decimal, ROUND_FLOOR, localcontext
from hashlib import sha256
import argparse
import json
import struct
from typing import Any, Iterable, Sequence

VERSION = "1.0.0"
S = 10**15
Q = 339_949_771_344_778
N = 23
LIMB_BITS = 15
LIMB_BASE = 1 << LIMB_BITS
LIMB_MAX = LIMB_BASE - 1
FOLD_RULE = "fold15_be_lsbpad_v1"
BYTE_CODEBOOK_HASH = "0fc8922c2d84cfd965c9fe310070e00ffd86ed76117584a9ff2faa85136a0c47"


class OmegaError(ValueError):
    """Raised when an input or witness violates the reference specification."""


def _require_plain_int(value: Any, label: str) -> int:
    if type(value) is not int:  # reject bool, which is an int subclass
        raise OmegaError(f"{label} must be an integer, not {type(value).__name__}")
    return value


def _require_nonnegative_int(value: Any, label: str) -> int:
    value = _require_plain_int(value, label)
    if value < 0:
        raise OmegaError(f"{label} must be nonnegative")
    return value


def _require_limb(value: Any, label: str = "limb") -> int:
    value = _require_plain_int(value, label)
    if not 0 <= value <= LIMB_MAX:
        raise OmegaError(f"{label} must be in [0, {LIMB_MAX}]")
    return value


def _atan_inverse(n: int, precision: int) -> Decimal:
    """Return atan(1/n) using its alternating power series."""
    if n <= 1:
        raise OmegaError("atan reciprocal denominator must be greater than 1")
    with localcontext() as ctx:
        ctx.prec = precision + 12
        x = Decimal(1) / Decimal(n)
        x2 = x * x
        term = x
        total = term
        k = 1
        threshold = Decimal(10) ** Decimal(-(precision + 8))
        sign = -1
        while True:
            term *= x2
            addend = term / Decimal(2 * k + 1)
            if abs(addend) < threshold:
                break
            total = total + addend if sign > 0 else total - addend
            sign *= -1
            k += 1
        ctx.prec = precision
        return +total


def decimal_constants(precision: int = 80) -> tuple[Decimal, Decimal, Decimal, Decimal]:
    """Evaluate phi, pi, lambda, and kappa to the requested Decimal precision."""
    precision = _require_plain_int(precision, "precision")
    if precision < 30:
        raise OmegaError("precision must be at least 30 decimal digits")
    with localcontext() as ctx:
        ctx.prec = precision + 12
        phi = (Decimal(1) + Decimal(5).sqrt()) / Decimal(2)
        pi = Decimal(16) * _atan_inverse(5, precision + 10) - Decimal(4) * _atan_inverse(239, precision + 10)
        lam = phi / (phi + pi)
        kappa = -lam.ln()
        ctx.prec = precision
        return +phi, +pi, +lam, +kappa


def derive_protocol_q(scale: int = S, precision: int = 80) -> int:
    """Re-evaluate floor(scale*lambda) as a reproducibility check."""
    scale = _require_nonnegative_int(scale, "scale")
    if scale == 0:
        raise OmegaError("scale must be positive")
    _, _, lam, _ = decimal_constants(precision)
    with localcontext() as ctx:
        ctx.prec = precision
        return int((Decimal(scale) * lam).to_integral_value(rounding=ROUND_FLOOR))


@dataclass(frozen=True)
class BoundaryWitness:
    """Terminal state and complete ordered Euclidean remainder ledger."""

    d_final: int
    remainders: tuple[int, ...]
    scale: int = S
    multiplier: int = Q
    steps: int = N

    def to_dict(self) -> dict[str, Any]:
        return {
            "d_final": self.d_final,
            "remainders": list(self.remainders),
            "S": self.scale,
            "Q": self.multiplier,
            "N": self.steps,
        }


def contract_integer(d0: int, steps: int = N) -> BoundaryWitness:
    """Contract a nonnegative integer and retain every Euclidean remainder."""
    d = _require_nonnegative_int(d0, "d0")
    steps = _require_plain_int(steps, "steps")
    if steps < 0:
        raise OmegaError("steps must be nonnegative")
    remainders: list[int] = []
    for _ in range(steps):
        product = Q * d
        d_next, rho = divmod(product, S)
        if not 0 <= rho < S:
            raise AssertionError("Euclidean remainder invariant failed")
        remainders.append(rho)
        d = d_next
    return BoundaryWitness(d, tuple(remainders), S, Q, steps)


def reconstruct_integer(witness: BoundaryWitness) -> int:
    """Recover the unique nonnegative initial integer represented by a valid witness."""
    if not isinstance(witness, BoundaryWitness):
        raise OmegaError("witness must be a BoundaryWitness")
    if witness.scale != S or witness.multiplier != Q:
        raise OmegaError("witness constants do not match the reference protocol")
    if witness.steps != len(witness.remainders):
        raise OmegaError("witness step count does not match ledger length")
    d = _require_nonnegative_int(witness.d_final, "d_final")
    for index, rho in reversed(list(enumerate(witness.remainders))):
        rho = _require_plain_int(rho, f"remainders[{index}]")
        if not 0 <= rho < S:
            raise OmegaError(f"remainders[{index}] must be in [0, S)")
        numerator = S * d + rho
        quotient, remainder = divmod(numerator, Q)
        if remainder != 0:
            raise OmegaError(f"nondivisible reverse numerator at ledger index {index}")
        d = quotient
    # Canonical replay rejects fabricated but accidentally divisible ledgers.
    replay = contract_integer(d, witness.steps)
    if replay != witness:
        raise OmegaError("witness fails canonical forward replay")
    return d


def full_witness_identity_holds(d0: int, witness: BoundaryWitness) -> bool:
    """Check Q^N d0 = S^N dN + sum rho_k S^k Q^(N-1-k)."""
    d0 = _require_nonnegative_int(d0, "d0")
    n = witness.steps
    if n != len(witness.remainders):
        return False
    rhs = (S**n) * witness.d_final
    rhs += sum(rho * (S**k) * (Q ** (n - 1 - k)) for k, rho in enumerate(witness.remainders))
    return (Q**n) * d0 == rhs


def steps_to_zero(d0: int) -> int:
    """Return the least k such that F^k(d0)=0 for a nonnegative integer."""
    d = _require_nonnegative_int(d0, "d0")
    steps = 0
    while d:
        d = (Q * d) // S
        steps += 1
    return steps


def fold15(data: bytes | bytearray) -> tuple[list[int], int]:
    """Big-endian MSB-first bytes to 15-bit limbs with right/LSB zero padding."""
    if not isinstance(data, (bytes, bytearray)):
        raise OmegaError("data must be bytes or bytearray")
    data = bytes(data)
    total_bits = len(data) * 8
    padding_bits = (LIMB_BITS - total_bits % LIMB_BITS) % LIMB_BITS
    padded_bits = total_bits + padding_bits
    if padded_bits == 0:
        return [], 0
    value = int.from_bytes(data, "big") << padding_bits
    limbs = [
        (value >> shift) & LIMB_MAX
        for shift in range(padded_bits - LIMB_BITS, -1, -LIMB_BITS)
    ]
    return limbs, padding_bits


def unfold15(limbs: Sequence[int], padding_bits: int, original_byte_length: int) -> bytes:
    """Inverse of fold15, with strict validation of padding and length."""
    if not isinstance(limbs, Sequence) or isinstance(limbs, (str, bytes, bytearray)):
        raise OmegaError("limbs must be a sequence of integers")
    padding_bits = _require_plain_int(padding_bits, "padding_bits")
    original_byte_length = _require_nonnegative_int(original_byte_length, "original_byte_length")
    if not 0 <= padding_bits < LIMB_BITS:
        raise OmegaError("padding_bits must be in [0, 14]")
    checked = [_require_limb(v, f"limbs[{i}]") for i, v in enumerate(limbs)]
    if not checked:
        if padding_bits != 0 or original_byte_length != 0:
            raise OmegaError("empty limb sequence requires zero padding and zero byte length")
        return b""
    value = 0
    for limb in checked:
        value = (value << LIMB_BITS) | limb
    if padding_bits:
        if value & ((1 << padding_bits) - 1):
            raise OmegaError("nonzero LSB padding bits")
        value >>= padding_bits
    unpadded_bits = len(checked) * LIMB_BITS - padding_bits
    if unpadded_bits != original_byte_length * 8:
        raise OmegaError("limb count, padding, and original length are inconsistent")
    return value.to_bytes(original_byte_length, "big")


def codebook_tick(seed: int, steps: int = N) -> int:
    """Evaluate T(seed)=F^steps((seed+2)S)."""
    seed = _require_nonnegative_int(seed, "seed")
    steps = _require_plain_int(steps, "steps")
    if steps < 0:
        raise OmegaError("steps must be nonnegative")
    value = (seed + 2) * S
    for _ in range(steps):
        value = (Q * value) // S
    return value


def generate_codebook(size: int) -> list[int]:
    size = _require_plain_int(size, "size")
    if size < 0:
        raise OmegaError("size must be nonnegative")
    return [codebook_tick(seed) for seed in range(size)]


def decode_byte_tick(tick: int) -> int:
    """Exact nearest-integer decoder for the verified 256-entry byte codebook."""
    tick = _require_nonnegative_int(tick, "tick")
    # floor(tick/16693 + 1/2) - 2, evaluated using integers.
    seed = (2 * tick + 16_693) // (2 * 16_693) - 2
    if not 0 <= seed <= 255 or codebook_tick(seed) != tick:
        raise OmegaError("tick is not a codeword in the byte codebook")
    return seed


def iterative_codebook_decode(tick: int) -> tuple[int, int]:
    """Approximate reverse recurrence followed by exact codeword replay validation."""
    tick = _require_nonnegative_int(tick, "tick")
    t = tick
    for _ in range(N):
        t = (t * S) // Q
    # round(t/S)-2 with nonnegative half-up rounding.
    seed = (2 * t + S) // (2 * S) - 2
    if seed < 0 or codebook_tick(seed) != tick:
        raise OmegaError("tick is not accepted by the iterative codebook decoder")
    return seed, t


def byte_codebook_json_hash(values: Sequence[int]) -> str:
    raw = json.dumps(list(values), separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    return sha256(raw).hexdigest()


def tick_map_u64be_hash(values: Iterable[int]) -> str:
    digest = sha256()
    for i, value in enumerate(values):
        value = _require_nonnegative_int(value, f"values[{i}]")
        if value >= 1 << 64:
            raise OmegaError("tick does not fit unsigned 64-bit encoding")
        digest.update(struct.pack(">Q", value))
    return digest.hexdigest()


def ramanujan_multiplier(aspect_ratio: Decimal, precision: int = 80) -> Decimal:
    """Ramanujan's second-approximation multiplier for r=a/b >= 1."""
    if not isinstance(aspect_ratio, Decimal):
        aspect_ratio = Decimal(str(aspect_ratio))
    if aspect_ratio < 1:
        raise OmegaError("aspect ratio must be at least 1")
    with localcontext() as ctx:
        ctx.prec = precision + 10
        h = ((aspect_ratio - 1) / (aspect_ratio + 1)) ** 2
        value = 1 + (3 * h) / (10 + (4 - 3 * h).sqrt())
        ctx.prec = precision
        return +value


def omega_coordinate(multiplier: Decimal, precision: int = 80) -> Decimal:
    if not isinstance(multiplier, Decimal):
        multiplier = Decimal(str(multiplier))
    _, _, lam, _ = decimal_constants(precision + 10)
    with localcontext() as ctx:
        ctx.prec = precision
        return +(Decimal(1) + (multiplier - Decimal(1)) / lam)


def truncate_decimal(value: Decimal, places: int = 15) -> Decimal:
    places = _require_nonnegative_int(places, "places")
    quantum = Decimal(1).scaleb(-places)
    return value.quantize(quantum, rounding=ROUND_FLOOR)


def corrected_ellipse_table() -> list[dict[str, str]]:
    # The revised ellipse document declares a 15-decimal literal policy for
    # lambda and computes v from the stored 15-decimal multiplier.
    phi, _, _, _ = decimal_constants(90)
    lambda_policy = Decimal("0.339949771344778")
    labels_and_ratios: list[tuple[str, Decimal]] = [
        ("1", Decimal(1)),
        ("phi", phi),
        *[(str(i), Decimal(i)) for i in range(2, 11)],
    ]
    rows: list[dict[str, str]] = []
    for label, ratio in labels_and_ratios:
        w = truncate_decimal(ramanujan_multiplier(ratio, 90), 15)
        with localcontext() as ctx:
            ctx.prec = 90
            v_untruncated = Decimal(1) + (w - Decimal(1)) / lambda_policy
        v = truncate_decimal(v_untruncated, 15)
        rows.append({"aspect_ratio": label, "multiplier": format(w, "f"), "coordinate": format(v, "f")})
    return rows


def _expect_omega_error(fn: Any) -> None:
    try:
        fn()
    except OmegaError:
        return
    raise AssertionError("expected OmegaError was not raised")


def run_reference_tests() -> dict[str, Any]:
    """Run deterministic unit, exhaustive finite-domain, and integrity tests."""
    phi, pi, lam, kappa = decimal_constants(90)
    derived_q_70 = derive_protocol_q(S, 70)
    derived_q_90 = derive_protocol_q(S, 90)
    assert derived_q_70 == Q == derived_q_90

    # Exhaustive 15-bit witness verification.
    max_steps_to_zero = 0
    for limb in range(LIMB_BASE):
        witness = contract_integer(limb)
        assert witness.d_final == 0
        assert reconstruct_integer(witness) == limb
        assert full_witness_identity_holds(limb, witness)
        max_steps_to_zero = max(max_steps_to_zero, steps_to_zero(limb))
    assert max_steps_to_zero == 10

    # Reversible byte folding across deterministic lengths and content patterns.
    folding_vectors = [
        b"",
        b"\x00",
        b"\xff",
        b"JPRV-LEGEND",
        bytes(range(256)),
        bytes((i * 73 + 19) % 256 for i in range(1025)),
    ]
    for vector in folding_vectors:
        limbs, padding = fold15(vector)
        assert unfold15(limbs, padding, len(vector)) == vector

    # Byte codebook: full enumeration and direct decoding.
    byte_ticks = generate_codebook(256)
    byte_diffs = [b - a for a, b in zip(byte_ticks, byte_ticks[1:])]
    assert byte_ticks[0] == 33_386
    assert byte_ticks[65] == 1_118_443
    assert byte_ticks[255] == 4_290_149
    assert min(byte_diffs) == 16_692 and max(byte_diffs) == 16_694
    byte_residuals = [tick - 16_693 * (seed + 2) for seed, tick in enumerate(byte_ticks)]
    assert min(byte_residuals) == 0 and max(byte_residuals) == 48
    for seed, tick in enumerate(byte_ticks):
        assert decode_byte_tick(tick) == seed
    byte_hash = byte_codebook_json_hash(byte_ticks)
    assert byte_hash == BYTE_CODEBOOK_HASH

    # 16-bit codebook: exhaustive monotonicity and iterative decoding.
    ticks16 = generate_codebook(65_536)
    diffs16 = [b - a for a, b in zip(ticks16, ticks16[1:])]
    assert ticks16[0] == 33_386
    assert ticks16[1] == 50_079
    assert ticks16[-1] == 1_094_021_571
    assert min(diffs16) == 16_692 and max(diffs16) == 16_694
    min_error_num = None
    max_error_num = None
    for seed, tick in enumerate(ticks16):
        decoded, expanded = iterative_codebook_decode(tick)
        assert decoded == seed
        error_num = expanded - (seed + 2) * S
        min_error_num = error_num if min_error_num is None else min(min_error_num, error_num)
        max_error_num = error_num if max_error_num is None else max(max_error_num, error_num)
    map_hash = tick_map_u64be_hash(ticks16)

    # Corrected ellipse coordinate anchors.
    ellipse = corrected_ellipse_table()
    by_ratio = {row["aspect_ratio"]: row for row in ellipse}
    assert by_ratio["1"]["multiplier"] == "1.000000000000000"
    assert by_ratio["2"]["multiplier"] == "1.027976282991305"
    assert by_ratio["10"]["multiplier"] == "1.175989279709379"

    # Negative validation tests.
    _expect_omega_error(lambda: contract_integer(True))
    _expect_omega_error(lambda: unfold15([0], 1, 1))
    bad = contract_integer(1234)
    altered = BoundaryWitness(bad.d_final, bad.remainders[:-1] + ((bad.remainders[-1] + 1) % S,), S, Q, N)
    _expect_omega_error(lambda: reconstruct_integer(altered))
    _expect_omega_error(lambda: decode_byte_tick(byte_ticks[0] + 1))

    with localcontext() as ctx:
        ctx.prec = 90
        lambda_power_23 = lam ** N

    return {
        "implementation_version": VERSION,
        "constants": {
            "phi_50": format(phi, ".50f"),
            "pi_50": format(pi, ".50f"),
            "lambda_50": format(lam, ".50f"),
            "kappa_50": format(kappa, ".50f"),
            "lambda_power_23": format(lambda_power_23, ".50E"),
            "S": S,
            "Q": Q,
            "N": N,
            "Q_rederived_at_70_digits": derived_q_70,
            "Q_rederived_at_90_digits": derived_q_90,
        },
        "integer_kernel": {
            "verified_limb_count": LIMB_BASE,
            "limb_domain": [0, LIMB_MAX],
            "all_23_step_terminal_states_zero": True,
            "maximum_steps_to_zero": max_steps_to_zero,
            "exact_round_trip": True,
            "full_identity_checked": True,
        },
        "fold15": {
            "rule": FOLD_RULE,
            "vectors_checked": len(folding_vectors),
            "round_trip": True,
        },
        "byte_codebook": {
            "entries": len(byte_ticks),
            "first": byte_ticks[0],
            "seed_65": byte_ticks[65],
            "last": byte_ticks[-1],
            "increment_min": min(byte_diffs),
            "increment_max": max(byte_diffs),
            "residual_min": min(byte_residuals),
            "residual_max": max(byte_residuals),
            "compact_json_sha256": byte_hash,
            "direct_decoder_exact": True,
        },
        "sixteen_bit_codebook": {
            "entries": len(ticks16),
            "first": ticks16[0],
            "second": ticks16[1],
            "last": ticks16[-1],
            "increment_min": min(diffs16),
            "increment_max": max(diffs16),
            "iterative_decoder_exact_after_replay": True,
            "expanded_error_min": f"{min_error_num}/{S}",
            "expanded_error_max": f"{max_error_num}/{S}",
            "u64be_map_sha256": map_hash,
        },
        "ellipse_table": {
            "rows": ellipse,
            "corrected_r2_multiplier": by_ratio["2"]["multiplier"],
        },
        "negative_tests": "PASS",
        "overall": "PASS",
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Omega-Alpha exact finite reference implementation")
    parser.add_argument("--report", help="write the JSON test report to this path")
    parser.add_argument("--pretty", action="store_true", help="pretty-print the test report")
    args = parser.parse_args()
    report = run_reference_tests()
    encoded = json.dumps(report, sort_keys=True, indent=2 if args.pretty else None, separators=None if args.pretty else (",", ":"), ensure_ascii=False)
    if args.report:
        with open(args.report, "w", encoding="utf-8") as handle:
            handle.write(encoded)
            handle.write("\n")
    print(encoded)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())