from __future__ import annotations

from dataclasses import dataclass
from math import floor, pi, sqrt, cos, sin
from hashlib import sha256
import json
from typing import Any, Iterable

PHI = (1.0 + sqrt(5.0)) / 2.0
OMEGA0 = PHI + pi
TAU = 1.0 / OMEGA0
LAMBDA = PHI / OMEGA0
MU = pi / OMEGA0

SCHEMA = "OMEGA_REFERENCE_GEOMETRY/0.1"


def omega_pair_encode(x: float, y: float) -> tuple[float, float]:
    """Θ(x,y) = ((x+y)/Ω₀, x/(x+y)), for x,y > 0."""
    if x <= 0 or y <= 0:
        raise ValueError("omega_pair_encode requires x>0 and y>0")
    total = x + y
    return total / OMEGA0, x / total


def omega_pair_decode(c: float, alpha: float) -> tuple[float, float]:
    """Inverse of omega_pair_encode."""
    if c <= 0:
        raise ValueError("c must be > 0")
    if not 0.0 <= alpha <= 1.0:
        raise ValueError("alpha must lie in [0,1]")
    total = c * OMEGA0
    return total * alpha, total * (1.0 - alpha)


def srsd_encode(position: Iterable[float], scale: Iterable[float]):
    """
    Frozen SRSD primitive consumed as an external audit/resolution rule.
    Coordinatewise d=floor(y/s), rho=y-s*d.
    """
    y = list(position)
    s = list(scale)
    if len(y) != len(s) or not y:
        raise ValueError("position and scale must have the same nonzero dimension")
    if any(si <= 0 for si in s):
        raise ValueError("all scales must be positive")

    d = [floor(yi / si) for yi, si in zip(y, s)]
    rho = [yi - si * di for yi, si, di in zip(y, s, d)]
    return d, rho


def srsd_decode(scale: Iterable[float], cell: Iterable[int], residual: Iterable[float]):
    s, d, rho = list(scale), list(cell), list(residual)
    if not (len(s) == len(d) == len(rho)) or not s:
        raise ValueError("scale, cell, residual dimensions must agree")
    for si, ri in zip(s, rho):
        if si <= 0 or not (0 <= ri < si):
            raise ValueError("residual must satisfy 0 <= rho_i < s_i")
    return [si * di + ri for si, di, ri in zip(s, d, rho)]


def phase_to_unit_circle(u: float) -> tuple[float, float]:
    if not 0.0 <= u < 1.0:
        raise ValueError("phase must lie in [0,1)")
    theta = 2.0 * pi * u
    return cos(theta), sin(theta)


def canonical_bytes(state: dict[str, Any]) -> bytes:
    return json.dumps(
        state,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
        allow_nan=False,
    ).encode("utf-8")


def digest(state: dict[str, Any]) -> str:
    return sha256(canonical_bytes(state)).hexdigest()


def make_state(
    position: Iterable[float] | None = None,
    scale: Iterable[float] | None = None,
    phase: Iterable[float] | None = None,
    omega_pair: tuple[float, float] | None = None,
    payload: Any = None,
) -> dict[str, Any]:
    if position is None and scale is None:
        raise ValueError("prototype requires position and scale")
    if position is None or scale is None:
        raise ValueError("position and scale must be supplied together")

    position = list(position)
    scale = list(scale)
    d, rho = srsd_encode(position, scale)

    state: dict[str, Any] = {
        "schema": SCHEMA,
        "dimension": len(position),
        "position": position,
        "scale": scale,
        "cell": d,
        "residual": rho,
    }

    if phase is not None:
        p = list(phase)
        if any(not 0.0 <= u < 1.0 for u in p):
            raise ValueError("all phase coordinates must lie in [0,1)")
        state["phase"] = p

    if omega_pair is not None:
        x, y = omega_pair
        c, alpha = omega_pair_encode(x, y)
        state["omega_pair"] = {"c": c, "alpha": alpha}

    if payload is not None:
        state["payload"] = payload

    return state


def self_check() -> dict[str, Any]:
    eps = 1e-12

    assert abs((LAMBDA + MU) - 1.0) < eps
    assert abs((LAMBDA * LAMBDA - TAU * LAMBDA - TAU * TAU)) < eps

    c, a = omega_pair_encode(PHI, pi)
    assert abs(c - 1.0) < eps
    assert abs(a - LAMBDA) < eps

    x, y = omega_pair_decode(c, a)
    assert abs(x - PHI) < eps
    assert abs(y - pi) < eps

    state = make_state(
        position=[1.618, 3.1415, 4.7596, -0.25],
        scale=[0.01, 0.01, 0.01, 0.01],
        phase=[0.0, 0.25, 0.5, 0.75],
        omega_pair=(PHI, pi),
        payload={"kind": "prototype-point"},
    )
    recovered = srsd_decode(state["scale"], state["cell"], state["residual"])
    assert all(abs(a-b) < eps for a, b in zip(recovered, state["position"]))

    return {
        "schema": SCHEMA,
        "lambda_plus_mu": LAMBDA + MU,
        "distinguished_pair": [c, a],
        "digest": digest(state),
        "pass": True,
    }


if __name__ == "__main__":
    print(json.dumps(self_check(), indent=2, sort_keys=True))