#!/usr/bin/env python3
"""
OMEGA_DRIFT_EXPERIMENT_v0_1

Purpose
-------
Test, rather than assume, the hypothesis that a 4D OMEGA reference frame can
stabilize a drifting state representation.

Important scope boundary
------------------------
OMEGA_REFERENCE_GEOMETRY v0.2, as declared by the user, is an exact reference
geometry. It does NOT itself define a memory-drift correction operator.

Therefore this benchmark separates:

E0. Frozen SRSD exact encode/decode: representation-only control.
E1. A NEW experimental OMEGA spiral operator:
        R_OMEGA = R(2*pi*lambda) direct-sum R(2*pi*mu)
    with lambda = phi/(phi+pi), mu = pi/(phi+pi).
    Since lambda+mu=1, the two angles are counter-rotations modulo 2*pi.
E2. Adversarial model mismatch.
E3. Synthetic embedding-memory proxy under low precision.
E4. Random-dynamics control to test whether OMEGA is privileged absent
    OMEGA-matched dynamics.

No claim of AI-memory improvement is built into the code.
"""

from __future__ import annotations
import json, math, random, hashlib
from fractions import Fraction
from pathlib import Path
import numpy as np

SEED = 20260831
OUT = Path(__file__).with_name("results.json")

phi = (1.0 + math.sqrt(5.0)) / 2.0
omega0 = phi + math.pi
lam = phi / omega0
mu = math.pi / omega0
alpha = 2.0 * math.pi * lam
beta = 2.0 * math.pi * mu

def floor_frac(x: Fraction) -> int:
    return x.numerator // x.denominator

def srsd_encode(x, s):
    q, r = [], []
    for xi, si in zip(x, s):
        qi = floor_frac(xi / si)
        ri = xi - si * qi
        q.append(qi)
        r.append(ri)
    return tuple(q), tuple(r)

def srsd_decode(q, r, s):
    return tuple(si * qi + ri for qi, ri, si in zip(q, r, s))

def rotate4_batch(x, a, b, dtype=np.float64):
    x = np.asarray(x, dtype=dtype)
    ca, sa = dtype(math.cos(a)), dtype(math.sin(a))
    cb, sb = dtype(math.cos(b)), dtype(math.sin(b))
    y = np.empty_like(x)
    y[:, 0] = ca*x[:, 0] - sa*x[:, 1]
    y[:, 1] = sa*x[:, 0] + ca*x[:, 1]
    y[:, 2] = cb*x[:, 2] - sb*x[:, 3]
    y[:, 3] = sb*x[:, 2] + cb*x[:, 3]
    return y

def rotate4_perrow(x, a, b):
    ca, sa, cb, sb = np.cos(a), np.sin(a), np.cos(b), np.sin(b)
    y = np.empty_like(x, dtype=np.float64)
    y[:, 0] = ca*x[:, 0] - sa*x[:, 1]
    y[:, 1] = sa*x[:, 0] + ca*x[:, 1]
    y[:, 2] = cb*x[:, 2] - sb*x[:, 3]
    y[:, 3] = sb*x[:, 2] + cb*x[:, 3]
    return y

def omega_rotate_nd(X, a, b, dtype=np.float32):
    X = np.asarray(X, dtype=dtype)
    Y = X.copy()
    ca, sa = dtype(math.cos(a)), dtype(math.sin(a))
    cb, sb = dtype(math.cos(b)), dtype(math.sin(b))
    for j in range(0, X.shape[1], 4):
        x0, x1, x2, x3 = X[:, j], X[:, j+1], X[:, j+2], X[:, j+3]
        Y[:, j]   = ca*x0 - sa*x1
        Y[:, j+1] = sa*x0 + ca*x1
        Y[:, j+2] = cb*x2 - sb*x3
        Y[:, j+3] = sb*x2 + cb*x3
    return Y

def cosine_diag(A, B):
    A, B = A.astype(np.float64), B.astype(np.float64)
    return np.sum(A*B, axis=1) / (np.linalg.norm(A, axis=1) * np.linalg.norm(B, axis=1))

def topk_neighbors(X, k=10):
    Xd = X.astype(np.float64)
    Xn = Xd / np.maximum(np.linalg.norm(Xd, axis=1, keepdims=True), 1e-30)
    S = Xn @ Xn.T
    np.fill_diagonal(S, -np.inf)
    inds = np.argpartition(-S, kth=k-1, axis=1)[:, :k]
    rows = np.arange(X.shape[0])[:, None]
    order = np.argsort(-S[rows, inds], axis=1)
    return inds[rows, order]

def neighbor_retention(base, out, k=10):
    A, B = topk_neighbors(base, k), topk_neighbors(out, k)
    return float(np.mean([len(set(a).intersection(b))/k for a, b in zip(A, B)]))

def main():
    results = {
        "experiment": "OMEGA_DRIFT_EXPERIMENT_v0_1",
        "seed": SEED,
        "constants_approx_noncanonical": {
            "phi": phi,
            "omega0": omega0,
            "lambda": lam,
            "mu": mu,
            "alpha": alpha,
            "beta": beta,
            "alpha_plus_beta_over_2pi": (alpha+beta)/(2*math.pi),
        },
        "scope": {
            "v0_2_reference_geometry_is_not_modified": True,
            "srsd_is_frozen_external": True,
            "omega_spiral_operator_is_new_experimental_operator": True,
            "real_ai_embedding_model_used": False,
            "embedding_test_is_synthetic_proxy": True,
        }
    }

    # E0 exact SRSD control
    rng_py = random.Random(SEED)
    failures = 0
    trials = 10000
    for _ in range(trials):
        x, s = [], []
        for _j in range(4):
            x.append(Fraction(rng_py.randint(-10**6, 10**6), rng_py.randint(1, 1000)))
            s.append(Fraction(rng_py.randint(1, 10**4), rng_py.randint(1, 1000)))
        q, r = srsd_encode(x, s)
        x2 = srsd_decode(q, r, s)
        if tuple(x) != x2 or any(not (Fraction(0) <= ri < si) for ri, si in zip(r, s)):
            failures += 1
    results["E0_exact_SRSD"] = {
        "trials": trials,
        "failures": failures,
        "interpretation": "Exact encode/decode is lossless; representation alone cannot correct or create drift."
    }

    # Shared 4D states
    rng = np.random.default_rng(SEED)
    M = 512
    x0 = rng.normal(size=(M, 4))
    x0 /= np.linalg.norm(x0, axis=1, keepdims=True)
    checkpoints = [1, 10, 100, 1000, 10000, 100000]

    # E1 matched OMEGA dynamics
    x_iter = x0.astype(np.float32)
    matched = []
    for k in range(1, max(checkpoints)+1):
        x_iter = rotate4_batch(x_iter, alpha, beta, dtype=np.float32)
        if k in checkpoints:
            oracle = rotate4_batch(
                x0,
                (k*alpha) % (2*math.pi),
                (k*beta) % (2*math.pi),
                dtype=np.float64
            )
            anchor = rotate4_batch(
                x0.astype(np.float32),
                (k*alpha) % (2*math.pi),
                (k*beta) % (2*math.pi),
                dtype=np.float32
            )
            ei = np.linalg.norm(x_iter.astype(np.float64)-oracle, axis=1)
            ea = np.linalg.norm(anchor.astype(np.float64)-oracle, axis=1)
            ni = np.abs(np.linalg.norm(x_iter.astype(np.float64), axis=1)-1)
            na = np.abs(np.linalg.norm(anchor.astype(np.float64), axis=1)-1)
            matched.append({
                "steps": k,
                "iterative_float32_mean_state_error": float(ei.mean()),
                "direct_anchor_float32_mean_state_error": float(ea.mean()),
                "iterative_float32_mean_norm_error": float(ni.mean()),
                "direct_anchor_float32_mean_norm_error": float(na.mean()),
            })
    results["E1_matched_omega_spiral"] = {
        "operator": "R(2*pi*lambda) direct-sum R(2*pi*mu)",
        "checkpoints": matched,
        "interpretation": "Direct procedural reconstruction prevents cumulative roundoff when the dynamics are known exactly. A generic direct renderer with the same true dynamics gives the same advantage."
    }

    # E2 adversarial model mismatch
    da, db = 1e-5, -1.7e-5
    a_true, b_true = alpha+da, beta+db
    x_iter = x0.astype(np.float32)
    mismatch = []
    for k in range(1, max(checkpoints)+1):
        x_iter = rotate4_batch(x_iter, a_true, b_true, dtype=np.float32)
        if k in checkpoints:
            oracle = rotate4_batch(
                x0,
                (k*a_true) % (2*math.pi),
                (k*b_true) % (2*math.pi),
                dtype=np.float64
            )
            omega_anchor = rotate4_batch(
                x0.astype(np.float32),
                (k*alpha) % (2*math.pi),
                (k*beta) % (2*math.pi),
                dtype=np.float32
            )
            true_anchor = rotate4_batch(
                x0.astype(np.float32),
                (k*a_true) % (2*math.pi),
                (k*b_true) % (2*math.pi),
                dtype=np.float32
            )
            mismatch.append({
                "steps": k,
                "iterative_float32_mean_state_error": float(np.linalg.norm(x_iter.astype(np.float64)-oracle,axis=1).mean()),
                "omega_anchor_wrong_model_mean_state_error": float(np.linalg.norm(omega_anchor.astype(np.float64)-oracle,axis=1).mean()),
                "generic_true_model_anchor_mean_state_error": float(np.linalg.norm(true_anchor.astype(np.float64)-oracle,axis=1).mean()),
            })
    results["E2_adversarial_mismatch"] = {
        "delta_alpha": da,
        "delta_beta": db,
        "checkpoints": mismatch,
        "interpretation": "A fixed OMEGA anchor becomes wrong if the actual dynamics are not OMEGA-aligned. The correct generic model remains stable."
    }

    # E3 synthetic embedding proxy
    rng = np.random.default_rng(123)
    N, D, C, T = 800, 64, 20, 5000
    centers = rng.normal(size=(C, D))
    centers /= np.linalg.norm(centers, axis=1, keepdims=True)
    labels = rng.integers(0, C, size=N)
    X0 = centers[labels] + 0.08*rng.normal(size=(N, D))
    X0 /= np.linalg.norm(X0, axis=1, keepdims=True)

    Xi = X0.astype(np.float16)
    for _ in range(T):
        Xi = omega_rotate_nd(Xi, alpha, beta, dtype=np.float16)
    Xt = omega_rotate_nd(X0, (T*alpha)%(2*math.pi), (T*beta)%(2*math.pi), dtype=np.float64)
    Xa = omega_rotate_nd(X0.astype(np.float16), (T*alpha)%(2*math.pi), (T*beta)%(2*math.pi), dtype=np.float16)

    matched_embed = {}
    for name, X in [("iterative_float16", Xi), ("direct_omega_anchor_float16", Xa)]:
        matched_embed[name] = {
            "mean_state_error": float(np.linalg.norm(X.astype(np.float64)-Xt,axis=1).mean()),
            "mean_cosine_to_true": float(np.mean(cosine_diag(X, Xt))),
            "knn10_retention": neighbor_retention(Xt, X, 10),
            "mean_norm_error": float(np.abs(np.linalg.norm(X.astype(np.float64),axis=1)-1).mean()),
        }

    a2, b2 = alpha+1e-4, beta-1.3e-4
    Xi2 = X0.astype(np.float16)
    for _ in range(T):
        Xi2 = omega_rotate_nd(Xi2, a2, b2, dtype=np.float16)
    Xt2 = omega_rotate_nd(X0, (T*a2)%(2*math.pi), (T*b2)%(2*math.pi), dtype=np.float64)
    Xomega = omega_rotate_nd(X0.astype(np.float16), (T*alpha)%(2*math.pi), (T*beta)%(2*math.pi), dtype=np.float16)
    Xtrue = omega_rotate_nd(X0.astype(np.float16), (T*a2)%(2*math.pi), (T*b2)%(2*math.pi), dtype=np.float16)

    mismatch_embed = {}
    for name, X in [
        ("iterative_float16", Xi2),
        ("omega_anchor_wrong_model", Xomega),
        ("generic_true_model_anchor", Xtrue),
    ]:
        mismatch_embed[name] = {
            "mean_state_error": float(np.linalg.norm(X.astype(np.float64)-Xt2,axis=1).mean()),
            "mean_cosine_to_true": float(np.mean(cosine_diag(X, Xt2))),
            "knn10_retention": neighbor_retention(Xt2, X, 10),
            "mean_norm_error": float(np.abs(np.linalg.norm(X.astype(np.float64),axis=1)-1).mean()),
        }

    results["E3_synthetic_embedding_proxy"] = {
        "N": N, "D": D, "clusters": C, "steps": T,
        "matched_omega_dynamics": matched_embed,
        "mismatched_dynamics": mismatch_embed,
        "interpretation": "Relational geometry (kNN) can remain nearly intact even while absolute state is wrong under a common wrong rotation. Absolute fidelity and relational fidelity must be scored separately."
    }

    # E4 random-dynamics adversarial control
    rng = np.random.default_rng(SEED)
    num_tasks, Tadv = 2000, 1000
    Xtask = rng.normal(size=(num_tasks,4))
    Xtask /= np.linalg.norm(Xtask,axis=1,keepdims=True)
    atrue = rng.uniform(0,2*math.pi,size=num_tasks)
    btrue = rng.uniform(0,2*math.pi,size=num_tasks)
    oracle = rotate4_perrow(
        Xtask,
        (Tadv*atrue)%(2*math.pi),
        (Tadv*btrue)%(2*math.pi)
    )
    def score(a,b):
        pred = rotate4_batch(
            Xtask,
            (Tadv*a)%(2*math.pi),
            (Tadv*b)%(2*math.pi),
            dtype=np.float64
        )
        return float(np.linalg.norm(pred-oracle,axis=1).mean())

    omega_score = score(alpha,beta)
    random_pairs = rng.uniform(0,2*math.pi,size=(200,2))
    scores = np.array([score(a,b) for a,b in random_pairs])
    results["E4_uniform_random_dynamics_control"] = {
        "tasks": num_tasks,
        "random_fixed_anchor_controls": 200,
        "omega_anchor_mean_state_error": omega_score,
        "random_anchor_mean_of_means": float(scores.mean()),
        "random_anchor_std_of_means": float(scores.std()),
        "fraction_random_anchors_better_than_omega": float(np.mean(scores < omega_score)),
        "best_random_anchor_mean_state_error": float(scores.min()),
        "worst_random_anchor_mean_state_error": float(scores.max()),
        "interpretation": "Across uniformly random dynamics, OMEGA performs like an ordinary fixed anchor; no universal privilege is detected."
    }

    OUT.write_text(json.dumps(results, indent=2, sort_keys=True), encoding="utf-8")
    print(json.dumps(results, indent=2, sort_keys=True))

if __name__ == "__main__":
    main()