"""Cortex-A 0.5 — live latent-AR inference (CPU, JAX/Flax).

A deep latent-AR PLANNER (32 layers over 5-token chunks) predicts the next chunk
embedding through a KV-cached decode; a shallow causal AR WRITER decodes each token
conditioned on the FULL planner plan — cond_fuse(concat(token, plan)) — with the planner
advancing once per 5-token chunk. The tied/factorized head reads tokens out (standard
next-token decoding). The latent backbone (planner) is ALWAYS on: its plan conditions
every generated token, so generation is the real two-level pipeline, not the writer alone.
"""
import os, time

os.environ.setdefault("JAX_PLATFORMS", "cpu")
os.environ.setdefault("XLA_FLAGS", "--xla_cpu_multi_thread_eigen=true")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")        # accelerated weight download

import numpy as np
import jax, jax.numpy as jnp
from flax import nnx
import gradio as gr
from huggingface_hub import hf_hub_download

from cortex.config import ModelConfig, EOT_ID, BOS_ID, TOKENIZER_VOCAB
from cortex.model import compute_rope
from cortex.latent_diffusion import CortexLatentDiffusion
from cortex.tokenizer import CortexTokenizer

# Latest weights = the Phase-2 continued-pretraining checkpoint. Tokenizer is bundled
# locally (static; canonical source is the Madarabr/CortexSym-xv-0.5 repo) -> no download.
WEIGHTS_REPO = os.environ.get("CORTEX_REPO", "Madarabr/CortexSym-xv-0.5-phase2")
TOKEN = os.environ.get("HF_TOKEN")
MAXT  = int(os.environ.get("CORTEX_MAX_CTX", "1024"))          # context cap (decode attends the whole
                                                              # cache each step -> keep it small for speed)
_HERE = os.path.dirname(os.path.abspath(__file__))

# Mirror the training launch so the param tree matches the checkpoint exactly:
# SWA 2048 + 4 learned sinks + the parallel-writer head + full-bandwidth plan conditioning.
# use_writer_mtp loads the trained recurrent EAGLE draft head (verified 84.6% depth-1 acceptance,
# 3.31x tokens/verify) so the latest checkpoint's writer_mtp params load instead of being ignored.
MCFG = ModelConfig(arch="latent_diffusion", sliding_window=2048, learned_sink=4,
                   use_parallel_writer=True, use_writer_mtp=True, writer_mtp_depth=4)
P, MAXC = MCFG.patch_size, MAXT // MCFG.patch_size

# ----------------------------------------------------------------------------- load
def _load_by_path(params, path):
    """Path-keyed bf16 safetensors load ('<idx>|<keystr>') -> the leaf's dtype; robust to
    leaf ordering, leaves any param absent from the checkpoint at its fresh init."""
    from safetensors.flax import load_file
    by_path = {k.split("|", 1)[1]: v for k, v in load_file(str(path)).items()}
    flat, td = jax.tree_util.tree_flatten_with_path(params)
    out, missing = [], []
    for p, leaf in flat:
        ks = jax.tree_util.keystr(p)
        if ks in by_path:
            out.append(jnp.asarray(by_path[ks], dtype=leaf.dtype))
        else:
            missing.append(ks); out.append(leaf)
    if missing:
        print(f"[cortex] {len(missing)} params not in checkpoint (fresh init), e.g. {missing[:3]}", flush=True)
    return jax.tree_util.tree_unflatten(td, out)

def _download_weights():
    return hf_hub_download(WEIGHTS_REPO, "model.safetensors", repo_type="model", token=TOKEN)

print("[cortex] loading latest weights + tokenizer ...", flush=True)
_wpath = _download_weights()
_m = CortexLatentDiffusion(MCFG, rngs=nnx.Rngs(0))
_GD, _P = nnx.split(_m)
_P = _load_by_path(_P, _wpath)
MODEL = nnx.merge(_GD, _P)
TOK = CortexTokenizer(os.path.join(_HERE, "tokenizer.json"))   # bundled -> instant
N_PARAMS = sum(int(x.size) for x in jax.tree.leaves(_P))

COSC, SINC = compute_rope(MAXC, MCFG.head_dim, MCFG.rope_base, MODEL.bcfg, jnp.float32)  # chunk pos
COST, SINT = compute_rope(MAXT, MCFG.head_dim, MCFG.rope_base, MODEL.wcfg, jnp.float32)  # token pos
_ZC = jnp.zeros((1, 1, MCFG.d_model), jnp.float32)             # null hint (chunk-0 tokens only)

# ------------------------------------------------------------------ jitted primitives
@jax.jit
def _embed(ids):
    return MODEL.embed(ids)

@jax.jit
def _predict(cm, pos, cache):                                 # chunk mean -> next-chunk plan (backbone)
    return MODEL.predict_chunk(cm, pos, cache, COSC, SINC)

@jax.jit
def _write(tok, cond, pos, wcache):                           # token + plan -> next-token logits [1,V]
    return MODEL.writer_step(tok, cond, pos, wcache, COST, SINT)

print("[cortex] warmup compile (planner decode + writer step) ...", flush=True)
_t0 = time.time()
_pr, _ = _predict(jnp.zeros((1, 1, MCFG.d_model), jnp.float32), jnp.int32(0), MODEL.init_chunk_cache(1, MAXC))
_lg, _ = _write(jnp.zeros((1, 1), jnp.int32), _pr, jnp.int32(0), MODEL.init_writer_cache(1, MAXT))
jax.block_until_ready(_lg)
print(f"[cortex] ready: {N_PARAMS/1e6:.0f}M params | warmup {time.time()-_t0:.1f}s", flush=True)

# ----------------------------------------------------------------------- generation
def _no_repeat_banned(out, n=3):
    """Tokens that would complete an n-gram already in `out` (no_repeat_ngram_size=n) — a HARD
    block on exact phrase repetition (the '32 - 32 is 0' loop), keeping the first occurrence."""
    if len(out) < n - 1:
        return ()
    pre = tuple(out[-(n - 1):])
    return {out[i + n - 1] for i in range(len(out) - (n - 1)) if tuple(out[i:i + n - 1]) == pre}

def _sample_row(lg, temperature, top_k, top_p, rng, prev=(), rep_pen=1.3, banned=()):
    lg = np.asarray(lg[:TOKENIZER_VOCAB], dtype=np.float64)
    if banned:                                                # no-repeat n-gram: hard-block phrase loops
        bi = np.fromiter((t for t in banned if t < lg.size), np.int64)
        if bi.size:
            lg[bi] = -1e30
    if rep_pen != 1.0 and prev:                               # repetition penalty (softens token reuse)
        idx = np.fromiter({t for t in prev if t < lg.size}, np.int64)
        if idx.size:
            lg[idx] = np.where(lg[idx] > 0, lg[idx] / rep_pen, lg[idx] * rep_pen)
    if temperature <= 1e-3:
        return int(lg.argmax())
    lg = lg / temperature
    if top_k and 0 < top_k < lg.size:
        kth = np.partition(lg, -top_k)[-top_k]; lg[lg < kth] = -1e30
    p = np.exp(lg - lg.max()); p /= p.sum()
    if 0.0 < top_p < 1.0:
        order = np.argsort(-p); csum = np.cumsum(p[order])
        cut = np.searchsorted(csum, top_p) + 1
        p[order[cut:]] = 0.0; p /= p.sum()
    return int(rng.choice(p.size, p=p))

def _generate(ids, max_new, temperature, top_k, top_p, rep_pen):
    """Token-by-token via the FULL two-level pipeline. The planner predicts the next chunk
    embedding (KV-cached) and the AR writer decodes tokens conditioned on that plan
    (cond[t]=chat[(t+1)//P-1]); the planner advances on every chunk boundary. Yields
    (out_ids, elapsed_s, avg_latent_cos) — latent-cos is the live backbone-vs-realized signal."""
    ids = [int(t) for t in ids]
    maxprompt = MAXT - P                                       # keep the prompt within context; the
    ids = ids[-maxprompt:] if len(ids) > maxprompt else ids    # pos>=MAXT break caps generation length
    if len(ids) % P:
        ids = [BOS_ID] * (P - len(ids) % P) + ids
    L = len(ids); nchunk = L // P
    emb_p = _embed(jnp.asarray([ids], jnp.int32))
    pcache = MODEL.init_chunk_cache(1, MAXC)
    wcache = MODEL.init_writer_cache(1, MAXT)
    chats = []
    for c in range(nchunk):                                    # prime the planner over prompt chunks
        cm = emb_p[:, c * P:(c + 1) * P].mean(axis=1, keepdims=True)
        pc, pcache = _predict(cm, jnp.int32(c), pcache); chats.append(pc)

    def hint_for(t):                                           # the planner plan conditioning token t+1
        si = (t + 1) // P - 1
        return chats[si] if 0 <= si < len(chats) else _ZC

    logits = None
    for t in range(L):                                         # prime the writer over prompt tokens
        logits, wcache = _write(jnp.asarray([[ids[t]]], jnp.int32), hint_for(t), jnp.int32(t), wcache)

    rng = np.random.default_rng()
    out, buf, coss = [], [], []
    pos, cur_c = L, nchunk
    t0 = time.time()
    for _ in range(int(max_new)):
        if pos >= MAXT:
            break
        nxt = _sample_row(np.asarray(logits[0]), temperature, top_k, top_p, rng, prev=out,
                          rep_pen=rep_pen, banned=_no_repeat_banned(out, 3))
        if nxt == EOT_ID:
            break
        out.append(nxt); buf.append(nxt)
        if len(buf) == P:                                      # chunk done -> planner step + diagnostic
            cm = _embed(jnp.asarray([buf], jnp.int32)).mean(axis=1, keepdims=True)
            pr = chats[cur_c - 1]
            coss.append(float(jnp.sum(pr[0, 0] * cm[0, 0]) /
                              (jnp.linalg.norm(pr[0, 0]) * jnp.linalg.norm(cm[0, 0]) + 1e-6)))
            pc, pcache = _predict(cm, jnp.int32(cur_c), pcache); chats.append(pc); cur_c += 1; buf = []
        logits, wcache = _write(jnp.asarray([[nxt]], jnp.int32), hint_for(pos), jnp.int32(pos), wcache)
        pos += 1
        yield out, time.time() - t0, (sum(coss) / len(coss) if coss else 0.0)

def _build_ids(history, message):
    ids = [BOS_ID]
    for m in history:
        if m.get("content"):
            ids += TOK.encode(m["content"], add_eot=True)
    ids += TOK.encode(message, add_eot=False)
    return ids

def respond(message, history, temperature, top_k, top_p, max_new, rep_pen):
    if not message or not message.strip():
        yield history or [], ""
        return
    history = (history or []) + [{"role": "user", "content": message},
                                {"role": "assistant", "content": ""}]
    ids = _build_ids(history[:-2], message)
    n, dt, lc = 0, 0.0, 0.0
    for out, dt, lc in _generate(ids, max_new, temperature, int(top_k), top_p, rep_pen):
        n = len(out)
        history[-1]["content"] = TOK.decode(out)
        yield history, (f"⚡ **{n/max(dt,1e-6):.1f} tok/s** · {n} tokens · {dt:.1f}s &nbsp;|&nbsp; "
                        f"🧠 latent backbone **active** · cos {lc:.3f}")
    if n == 0:
        history[-1]["content"] = "*(model emitted end-of-text immediately — try a different prompt)*"
    yield history, (f"✅ **{n/max(dt,1e-6):.1f} tok/s** · {n} tokens in {dt:.1f}s &nbsp;|&nbsp; "
                    f"🧠 latent backbone **active** · avg cos {lc:.3f}")

# ------------------------------------------------------------------------------- UI
# NOTE: inputs MUST be >=16px or mobile browsers auto-zoom the viewport when the field gains
# focus (e.g. on submit) and never zoom back -- the "screen zooms while generating" bug.
CSS = """
.gradio-container {max-width: 920px !important; margin: auto;}
#hdr {text-align:center; padding: 6px 0 2px;}
#hdr h1 {font-size: 1.9rem; margin: 0; letter-spacing:-.5px;}
#hdr p {color: var(--body-text-color-subdued); margin: 4px 0 0; font-size:.95rem;}
#status {text-align:center; min-height: 1.4em; font-size:.92rem;}
textarea, input, select, .gr-text-input {font-size: 16px !important;}
footer {visibility: hidden;}
"""
# Pin the viewport so nothing can auto-zoom on focus (belt-and-suspenders with the 16px rule).
VIEWPORT = ('<meta name="viewport" content="width=device-width, initial-scale=1, '
            'maximum-scale=1, user-scalable=no">')

def _sync_weights():
    """Pull the freshest Phase-2 checkpoint without restarting the Space."""
    global _P, MODEL
    w = hf_hub_download(WEIGHTS_REPO, "model.safetensors", repo_type="model", token=TOKEN,
                        force_download=True)
    _P = _load_by_path(_P, w)
    MODEL = nnx.merge(_GD, _P)
    jax.clear_caches()                                         # re-jit over the new params
    return "✅ synced to the latest weights"

with gr.Blocks(title="Cortex-A 0.5", theme=gr.themes.Soft(primary_hue="indigo"),
               css=CSS, head=VIEWPORT) as demo:
    gr.HTML(
        f"<div id='hdr'><h1>🧠 Cortex-A 0.5</h1>"
        f"<p>A ~{N_PARAMS/1e6:.0f}M latent-AR language model — a deep planner over "
        f"{P}-token chunks + a shallow AR writer, trained from scratch in JAX/Flax. "
        f"The latent backbone conditions every token.</p></div>")

    chat = gr.Chatbot(height=440, type="messages", show_label=False, show_copy_button=True,
                      avatar_images=(None, None),
                      placeholder="<em>Type the start of a sentence and Cortex-A 0.5 continues it.</em>")
    status = gr.Markdown("", elem_id="status")
    with gr.Row():
        msg = gr.Textbox(placeholder="The capital of France is…", show_label=False, scale=8, lines=1,
                         container=False)
        send = gr.Button("Generate", variant="primary", scale=1, min_width=110)
    with gr.Row():
        clear = gr.Button("🗑 Clear", size="sm")
        stop = gr.Button("⏹ Stop", size="sm", variant="stop")
    gr.Examples(
        ["The capital of France is", "Water is made of hydrogen and", "Once upon a time,",
         "The three primary colors are", "def fibonacci(n):", "In the year 2050,",
         "Q: What is 12 times 8?\nA:"],
        inputs=msg, label="Try a prompt")

    with gr.Accordion("⚙️ Inference settings (tuned for best output)", open=False):
        with gr.Row():
            temp = gr.Slider(0.0, 1.5, 0.5, step=0.05, label="temperature",
                             info="lower = more focused/confident · 0 = greedy")
            mnt = gr.Slider(16, 1024, 200, step=16, label="max new tokens",
                            info="CPU ≈ 7–8 tok/s · 200 ≈ 25s")
        with gr.Row():
            tk_k = gr.Slider(0, 200, 40, step=1, label="top-k", info="0 = off")
            tk_p = gr.Slider(0.1, 1.0, 0.9, step=0.01, label="top-p")
            rpen = gr.Slider(1.0, 2.0, 1.3, step=0.05, label="repetition penalty")
        sync = gr.Button("🔄 Sync latest weights", size="sm")
        sync_msg = gr.Markdown("")

    args = [msg, chat, temp, tk_k, tk_p, mnt, rpen]
    ev1 = send.click(respond, args, [chat, status]).then(lambda: "", None, msg)
    ev2 = msg.submit(respond, args, [chat, status]).then(lambda: "", None, msg)
    stop.click(None, None, None, cancels=[ev1, ev2])
    clear.click(lambda: ([], ""), None, [chat, status])
    sync.click(_sync_weights, None, sync_msg)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))