import gradio as gr
import torch
import sys, os, time, hashlib

sys.path.insert(0, "src")
from mulattice.protocol import SynthCfg, build_model, force_fast_scan, logits_at, synth_batch

DEVICE = "cpu"
CFG = SynthCfg()

print("Loading models...", flush=True)
MODELS = {}
for name in ["capacity8k", "forgetting", "realtext"]:
    path = f"checkpoints/{name}.pt"
    if os.path.exists(path):
        ck = torch.load(path, map_location="cpu", weights_only=False)
        c = SynthCfg()
        if ck.get("proto_cfg", {}).get("addr_scale"):
            c.addr_scale = float(ck["proto_cfg"]["addr_scale"])
        m = build_model(c)
        m.load_state_dict(ck["model"])
        m = m.to(DEVICE).eval()
        force_fast_scan(m)
        MODELS[name] = m
        print(f"  {name}.pt loaded", flush=True)
print("All models ready", flush=True)

INK = "#1c1e26"; MUT = "#6b7280"; HAIR = "#e5e7eb"; BG = "#fcfcfa"
ACC = "#2547d0"; OK = "#0e8a5f"; BAD = "#c2413b"; SOFT = "#f4f4f0"; UPD = "#b07807"

with open("demo.css") as f:
    DEMO_CSS = f.read()


def run_custom(text, n1, v1, nv1, n2, v2, nv2):
    m = MODELS.get("forgetting", MODELS.get("capacity8k"))
    seed = int(time.time()) % 100000
    pairs = []
    for n, v, nv in [(n1, v1, nv1), (n2, v2, nv2)]:
        if n.strip() and v.strip():
            pairs.append({"name": n.strip(), "value": v.strip(),
                          "new_value": nv.strip() if nv.strip() else None})
    if not pairs:
        return "Enter at least one fact (a name and a value).", ""

    used = set()
    def h(s, span, off):
        x = int(hashlib.md5(f"{seed}:{s}".encode()).hexdigest(), 16) % span
        while (off + x) in used:
            x = (x + 1) % span
        used.add(off + x)
        return off + x

    kmap, vmap = {}, {}
    for p in pairs:
        kmap[p["name"]] = h("k" + p["name"], CFG.NK, 0)
    for p in pairs:
        vmap[p["value"]] = h("v" + p["value"], CFG.NV, CFG.NK)
        if p["new_value"]:
            vmap[p["new_value"]] = h("v2" + p["new_value"], CFG.NV, CFG.NK)

    words = [w for w in text.split() if w.strip()][:500] or ["..."]
    fid = {w: h("f" + w, CFG.NK, 0) for w in set(words)}
    fill = [fid[w] for w in words]

    store_rows, upd_rows = [], []
    for p in pairs:
        vid = vmap[p["value"]]
        row = dict(p); row["kid"] = kmap[p["name"]]; row["vid"] = vid
        if p["new_value"]:
            row["nid"] = vmap[p["new_value"]]
            upd_rows.append(row)
        store_rows.append(row)

    F = 6; block = F + 2
    n_upd = len(upd_rows); Q = len(store_rows)
    L = (len(store_rows) + n_upd) * block + 2 * Q
    toks = torch.zeros(1, L, dtype=torch.long)

    def take_f(n):
        out = []
        while len(out) < n:
            i = int(torch.randint(0, len(fill), (1,)))
            out.extend(fill[i:i+n-len(out)])
        return out[:n]

    for i, r in enumerate(store_rows):
        s = i * block
        toks[0, s:s+F] = torch.tensor(take_f(F))
        toks[0, s+F], toks[0, s+F+1] = r["kid"], r["vid"]
    for j, r in enumerate(upd_rows):
        s = (len(store_rows) + j) * block
        toks[0, s:s+F] = torch.tensor(take_f(F))
        toks[0, s+F], toks[0, s+F+1] = r["kid"], r["nid"]
    S2 = (len(store_rows) + n_upd) * block
    ans_l = []
    for q, r in enumerate(store_rows):
        toks[0, S2+2*q] = r["kid"]
        toks[0, S2+2*q+1] = CFG.SEP
        ans_l.append(S2 + 2*q + 1)

    t0 = time.time()
    with torch.no_grad():
        pred = logits_at(m, toks, torch.tensor([ans_l])).argmax(-1).cpu()
    dt = time.time() - t0
    cells = m.blocks[0].mixer.scan.last_n_cells

    inv = {v: k for k, v in vmap.items()}
    correct = 0
    rows_html = '<table><tr><th>Asked</th><th>Answered</th><th>Expected</th><th></th></tr>'
    for q, r in enumerate(store_rows):
        got = inv.get(int(pred[0, q]), "?")
        exp = r["new_value"] or r["value"]
        ok = got == exp
        correct += ok
        old = f' <span class="pill">was {r["value"]}</span>' if r["new_value"] else ""
        cls = "y" if ok else "n"
        icon = "✅" if ok else "❌"
        rows_html += f'<tr><td><b>{r["name"]}</b></td><td class="{cls}">{got}</td><td>{exp}{old}</td><td>{icon}</td></tr>'
    rows_html += '</table>'

    stats_html = ('<div class="stats" style="display:grid">'
        f'<div class="stat"><b>{correct}/{len(store_rows)}</b><span>correct</span></div>'
        f'<div class="stat"><b>{L:,}</b><span>words read</span></div>'
        f'<div class="stat"><b>{dt:.1f}s</b><span>time</span></div>'
        f'<div class="stat"><b>{cells:,}</b><span>memory slots</span></div>'
        '</div>')
    return rows_html, stats_html


def run_scale(n_pairs, gap, pct_changed):
    m = MODELS.get("forgetting", MODELS.get("capacity8k"))
    g = torch.Generator().manual_seed(42)
    n = int(n_pairs)
    b = synth_batch(1, n, min(32, n), CFG,
                    filler=int(gap), overwrite_frac=float(pct_changed)/100, generator=g)
    t0 = time.time()
    with torch.no_grad():
        pred = logits_at(m, b["tokens"], b["ans_pos"]).argmax(-1).cpu()
    dt = time.time() - t0
    ok = (pred == b["gold"]).squeeze(0)
    up = b["was_updated"].squeeze(0)
    cells = m.blocks[0].mixer.scan.last_n_cells
    acc = ok.float().mean().item()
    acc_up = ok[up].float().mean().item() if bool(up.any()) else 1.0
    acc_keep = ok[~up].float().mean().item() if bool((~up).any()) else 1.0

    result_html = f'''<div style="text-align:center;padding:28px 0">
        <div style="font-size:48px;font-weight:600;color:{OK};margin-bottom:8px">{acc:.1%}</div>
        <div style="font-size:15px;color:{MUT};margin:6px 0 16px">{int(ok.sum())} of {ok.numel()} recalled correctly</div>
        <div style="font-size:14px;margin:4px 0">Changed facts: <b style="color:{UPD}">{acc_up:.1%}</b> &middot; Unchanged: <b style="color:{OK}">{acc_keep:.1%}</b></div>
    </div>
    <div class="stats" style="display:grid">
        <div class="stat"><b>{b["L"]:,}</b><span>words read</span></div>
        <div class="stat"><b>{cells:,}</b><span>memory slots</span></div>
        <div class="stat"><b>{dt:.1f}s</b><span>time</span></div>
    </div>'''
    return result_html


with gr.Blocks(title="μLattice — AI Memory That Grows Without the Cost Growing",
               css=DEMO_CSS + f"""
.gradio-container {{ max-width: 980px !important; margin: auto !important;
    background: {BG} !important; padding: 0 24px 60px !important;
    font-family: -apple-system, BlinkMacSystemFont, Inter, sans-serif !important; }}
footer {{ display: none !important; }}
""") as demo:

    gr.HTML(f"""
    <header>
     <div class="kicker">Tensorbend Research &middot; Eric E. &middot; Live proof &middot; Real weights</div>
     <h1 class="serif">&mu;Lattice &mdash; the growing memory, running</h1>
     <p class="dek serif">A model that reads text and keeps every fact in it &mdash; stores
     associations, replaces them on sight, and answers from memory. Every word costs the
     same, whether the model holds ten facts or eight thousand.</p>
     <div class="meta">3 trained models loaded: capacity8k &middot; forgetting &middot; realtext &middot; 4.2M parameters</div>
    </header>
    <main>
    """)

    # ── Section 01 ──
    gr.HTML("""
    <section>
     <div class="sn"><span class="n">01</span><h2 class="serif">Try it with your own facts</h2></div>
     <p class="lead">Paste any text. Define facts as <em>name &rarr; value</em> pairs.
     Mark any to be <strong>changed mid-stream</strong>. The model reads your words,
     stores the facts, and must return the <em>latest</em> value.</p>
     <div class="fig">
    """)

    with gr.Row():
        with gr.Column(scale=1):
            doc = gr.Textbox(label="Your text (the model reads this)", lines=5,
                value="The harbour was quiet that morning. Gulls circled the masts while the tide went out. A clerk stamped papers without looking up. Somewhere a bell rang twice, then stopped. The road east had been washed out again. Dust settled over the market stalls.")
        with gr.Column(scale=1):
            n1 = gr.Textbox(label="Name", value="the vault password")
            v1 = gr.Textbox(label="Value", value="swordfish")
            nv1 = gr.Textbox(label="Changed to (optional)", value="bluedragon")
            n2 = gr.Textbox(label="Name", value="the meeting time")
            v2 = gr.Textbox(label="Value", value="3pm")
            nv2 = gr.Textbox(label="Changed to (optional)", value="")

    btn1 = gr.Button("Store, overwrite, and answer", variant="primary")
    out1 = gr.HTML()
    stats1 = gr.HTML()

    gr.HTML(f"""
     <div class="figcap">The model reads your words as the stream, stores your facts as it
     passes them, replaces marked ones, then answers from memory. Uses the
     <b>forgetting</b> checkpoint (best at overwrite tasks).</div>
    </div>
    </section>
    """)

    # ── Section 02 ──
    gr.HTML(f"""
    <section>
     <div class="sn"><span class="n">02</span><h2 class="serif">The store at scale</h2></div>
     <p class="lead">The same mechanism under load. Store facts, change some mid-stream,
     then ask for them all back. The cost per word stays flat no matter how many are stored.</p>
     <div class="fig">
    """)

    with gr.Row():
        np_s = gr.Slider(64, 512, value=128, step=64, label="How many facts",
            info="Capped at 512 for CPU (GPU handles 8,192+)")
        gap_s = gr.Slider(0, 14, value=8, step=1, label="Words between each fact")
        ow_s = gr.Slider(0, 100, value=50, step=10, label="% of facts changed mid-stream")

    btn2 = gr.Button("Run the stress test", variant="primary")
    out2 = gr.HTML()

    gr.HTML(f"""
     <div class="figcap">Uses the <b>forgetting</b> checkpoint. On GPU this scales to
     8,192 facts in one 133K-token pass. On free-tier CPU, capped at 512.</div>
    </div>
    </section>
    </main>
    <footer style="max-width:980px;margin:60px auto 0;padding:24px 0;color:{MUT};font-size:13px;border-top:1px solid {HAIR};text-align:center">
     <p>Every answer is produced by a <b>live forward pass</b> of the trained model.</p>
     <p style="margin-top:12px"><a href="https://github.com/eelbaz/mLattice/blob/main/PAPER.md" style="color:{ACC};text-decoration:none">Full paper</a> &middot;
        <a href="https://github.com/eelbaz/mLattice" style="color:{ACC};text-decoration:none">Code &amp; tests</a></p>
     <p style="margin-top:12px;font-size:12px">Tensorbend Research &middot; Eric E. &middot; October 2026</p>
    </footer>
    """)

    btn1.click(run_custom, inputs=[doc, n1, v1, nv1, n2, v2, nv2], outputs=[out1, stats1])
    btn2.click(run_scale, inputs=[np_s, gap_s, ow_s], outputs=[out2])

demo.launch()