"""Multi-turn bounded SP-evict chat in MLX (the 'long-CoT rally' the project wanted).
Persistent KV across turns but bounded: each chunk trims cache to the system prompt, prepends
32 SP (compressing the distant conversation) + recent raw window. User msg tokens are force-fed,
then the assistant samples its reply. Tests whether reasoning is carried across turns under bounded memory."""
import sys, os, time
import numpy as np
import mlx.core as mx
from mlx_lm.models.cache import make_prompt_cache
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import sp_mlx

import sys as _s
RW = int(_s.argv[1]) if len(_s.argv)>1 else 128
C, MAXD, TEMP, MAX_RESP = 64, 4096, 0.6, 500
M = sp_mlx.get()
model, tok, pooler, H, eos, embT, mdtype = (M["model"], M["tok"], M["pooler"], M["H"], M["eos"], M["embT"], M["mdtype"])


def emb_ids(ids):
    return mx.zeros((1, 0, H), dtype=mdtype) if len(ids) == 0 else embT(mx.array([ids]))


def evict(kept):
    if not MAXD or len(kept) <= MAXD:
        return kept
    _, mass = pooler.forward_with_mass(emb_ids(kept).astype(mx.float32)); mx.eval(mass)
    m = np.array(mass)[0]; idx = np.sort(np.argsort(m)[-MAXD:])
    return [kept[i] for i in idx]


SYSTEM = "You are a helpful math tutor. Reason briefly step by step, then give the final answer."
cache = make_prompt_cache(model)
sys_ids = tok.encode(SYSTEM)
model(mx.array([sys_ids]), cache=cache); MQ = cache[0].offset
gen, kept, absorbed = [], [], 0
print(f"SYSTEM seeded MQ={MQ}\n", flush=True)


def turn(user_msg):
    global gen, kept, absorbed
    feed = tok.encode(("<｜end▁of▁sentence｜>" if gen else "") + f"<｜User｜>{user_msg}<｜Assistant｜>", add_special_tokens=False)
    fi, sampled, start, done = 0, 0, len(gen), False
    t0 = time.time()
    while not done:
        c0 = len(gen); R = min(c0, RW); nd_end = c0 - R
        if nd_end > absorbed:
            kept.extend(gen[absorbed:nd_end]); absorbed = nd_end; kept = evict(kept)
        sp = pooler.forward(emb_ids(kept).astype(mx.float32)).astype(mdtype)
        parts = [sp] + ([emb_ids(gen[c0 - R:c0])] if R > 0 else [])
        block = mx.concatenate(parts, axis=1)
        for c in cache:
            c.trim(c.offset - MQ)
        logits = model(mx.zeros((1, block.shape[1]), dtype=mx.int32), cache=cache, input_embeddings=block)
        last = logits[:, -1, :]
        for _ in range(C):
            if fi < len(feed):
                t = feed[fi]; fi += 1
            elif sampled < MAX_RESP:
                t = int(mx.random.categorical(last * (1.0 / TEMP)).item()); sampled += 1
                if t == eos:
                    done = True; break
            else:
                done = True; break
            gen.append(t)
            logits = model(mx.zeros((1, 1), dtype=mx.int32), cache=cache, input_embeddings=emb_ids([t]))
            last = logits[:, -1, :]
        if done:
            break
    reply = tok.decode(gen[start + len(feed):])
    return reply, time.time() - t0


CONV = [
    "I have 3 apples and then I buy 5 more. How many apples do I have?",
    "If I now give away 2 apples, how many do I have left?",
    "Now I double the number of apples I have. How many is that?",
    "I eat 4 of them. How many remain?",
    "Finally, I split the remaining apples equally between 2 friends. How many does each friend get?",
]
for i, u in enumerate(CONV, 1):
    r, dt = turn(u)
    print(f"--- Turn {i}  total_ctx={len(gen)}  survivors={len(kept)}  {dt:.1f}s ---", flush=True)
    print(f"USER: {u}", flush=True)
    print(f"ASSISTANT: {r.strip()[:500]}", flush=True)
    print("", flush=True)
print("MT_MLX_DONE", flush=True)