"""Four GNOMON components, each in ~20 lines, all jitted."""
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import jax
import jax.numpy as jnp
import numpy as np

from gnomon import (
    ternarize,
    compile_ldc,
    init_gnomon, init_mode_state, make_forward_jit,
    compute_bids, vickrey_topk,
    init_ib, ib_forward, ib_loss_terms,
)


def demo_ternary():
    print("[1/4] Ternary weights (BitNet b1.58)")
    w = jax.random.normal(jax.random.PRNGKey(0), (128, 128))
    w_q, scale = ternarize(w)
    uniq = np.unique(np.asarray(w_q / scale).round(3))
    print(f"      distinct values after ternarize: {sorted(uniq.tolist())}")
    print(f"      storage: 32 bits → 2 bits per weight (16x smaller)")
    print()


def demo_ldc():
    print("[2/4] Lipschitz Decision Circuit")
    rng = np.random.default_rng(0)
    N, D, C = 20_000, 32, 3
    # Three Gaussian clusters — linearly separable enough that LDA is optimal.
    centers = rng.standard_normal((C, D)).astype(np.float32) * 2.0
    labels = rng.integers(0, C, N).astype(np.int32)
    X = (centers[labels] + rng.standard_normal((N, D))).astype(np.float32)
    Y = np.eye(C, dtype=np.float32)[labels]

    X_tr, Y_tr = X[:15_000], Y[:15_000]
    X_te, y_te = X[15_000:], labels[15_000:]

    t0 = time.perf_counter()
    ldc = compile_ldc(X_tr, Y_tr, k=C - 1, max_leaves=128)
    dt = time.perf_counter() - t0

    x_test = jnp.asarray(X_te)
    logits = ldc.predict_jit(x_test)
    jax.block_until_ready(logits)

    t0 = time.perf_counter()
    logits = ldc.predict_jit(x_test)
    jax.block_until_ready(logits)
    infer_dt = time.perf_counter() - t0

    preds = np.asarray(logits).argmax(-1)
    acc = (preds == y_te).mean()

    eig_str = ", ".join(f"{e:.2f}" for e in ldc.lda_eigenvalues[:4])
    print(f"      LDA eigenvalues:       [{eig_str}]")
    print(f"      compiled in {dt:.3f}s → {ldc.n_leaves} leaves, "
          f"{ldc.n_internal} internal")
    print(f"      test accuracy: {acc:.4f}")
    print(f"      throughput:    {len(X_te) / infer_dt:,.0f} queries/s (jitted)")
    print()


def demo_ib():
    print("[3/4] Information-Bottleneck cache")
    params = init_ib(jax.random.PRNGKey(0), d_enc_in=64, d_dec_out=32,
                     d_hidden=128, k=16)
    x = jax.random.normal(jax.random.PRNGKey(1), (2, 32, 64))
    kv_true = jax.random.normal(jax.random.PRNGKey(2), (2, 32, 32))
    s, mu, log_sigma, kv_hat = ib_forward(params, x, jax.random.PRNGKey(3), S=32)
    loss, aux = ib_loss_terms(mu, log_sigma, kv_hat, kv_true)
    print(f"      input KV state:        2 × 32 × 32 = 2048 floats")
    print(f"      IB compressed state:   2 × 16      =   32 floats")
    print(f"      compression ratio:     {2048 / 32:.0f}x")
    print(f"      initial KL:            {float(aux['kl']):.3f}")
    print(f"      initial recon error:   {float(aux['recon']):.3f}")
    print()


def demo_vickrey():
    print("[4/4] Vickrey auction slot selection")
    S, D = 512, 32
    K = jax.random.normal(jax.random.PRNGKey(0), (S, D))
    proj = jax.random.normal(jax.random.PRNGKey(1), (D, 16))
    importance = jnp.linspace(1.0, 5.0, S)
    bids = compute_bids(K, proj, importance)
    for k in (32, 128, 256):
        mask, price = vickrey_topk(bids, k)
        kept = int(mask.sum())
        top_quartile = int(mask[int(0.75 * S):].sum())
        print(f"      budget={k:4d}  kept={kept:4d}  "
              f"price={float(price):.3f}  "
              f"top-quartile positions kept: {top_quartile}/{S // 4}")
    print()


def demo_gnomon_forward():
    print("[bonus] GNOMON forward pass (all four layers wired together)")
    params = init_gnomon(
        jax.random.PRNGKey(0),
        vocab_size=64, d_model=64, n_layers=2, n_heads=4,
        ib_k=16, seq_len=16, d_signal=16,
    )
    ids = jax.random.randint(jax.random.PRNGKey(1), (1, 16), 0, 64)
    mode_sig = jax.random.normal(jax.random.PRNGKey(2), (1, 16, 16))
    mode = init_mode_state(d_mod=4)

    forward = make_forward_jit(n_heads=4)
    _ = forward(params, ids, mode, mode_sig, jax.random.PRNGKey(4))
    jax.block_until_ready(_["logits"])

    t0 = time.perf_counter()
    out = forward(params, ids, mode, mode_sig, jax.random.PRNGKey(3))
    jax.block_until_ready(out["logits"])
    dt = time.perf_counter() - t0

    ib_total = sum(float(c["loss"]) for c in out["cache_logs"])
    print(f"      forward pass: {dt * 1000:.2f} ms (jitted, post-warmup)")
    print(f"      output logits: {out['logits'].shape}")
    print(f"      IB layers active: {len(out['cache_logs'])}")
    print(f"      total IB loss: {ib_total:.4f}")
    print()


if __name__ == "__main__":
    print("=" * 60)
    print("GNOMON component showcase")
    print("=" * 60)
    print()
    demo_ternary()
    demo_ldc()
    demo_ib()
    demo_vickrey()
    demo_gnomon_forward()
    print("=" * 60)
    print("All four components operational.")
    print("=" * 60)