#!/usr/bin/env python3
"""Measure per-block / per-expert expert-routing activity on live frank-primary.

Method: AMD IBS op sampling (data addresses, `perf record -e ibs_op// --data`)
on the llama-server process while concurrent generations of mixed content run.
Sampled load addresses are mapped through /proc/<pid>/maps (mmap VMAs of the
6 GGUF parts) back to individual tensors, then aggregated per expert block
and per expert WITHIN each block (512 experts, row-major expert-major layout).

Tensor extents are DERIVED from consecutive GGUF header offsets — no quant
type-size table (those are easy to get subtly wrong; a wrong per-element size
silently under-covers each tensor and dumps tail samples into a junk bucket).

Fleet law this instrument established (2026-09-26, 75 s window, 4 mixed
streams, 2.7M expert-tensor samples): per-BLOCK read volume is FLAT
(0.93x-1.08x of mean across all 40 CPU-resident blocks) because MoE top-k
routing fires every block every token — there are no "most active layers" at
`-ot` granularity. The real skew is per-EXPERT (top 10% of (block,expert)
cells carry ~40% of samples), which tensor-level overrides cannot reach.
Any 8 GPU-resident blocks are therefore equivalent; COUNT is the only lever.

Usage (needs passwordless sudo for perf only; model files must be readable):
  fp-expert-activity.py <pid> [duration_s]        # dump + sample + analyze
  fp-expert-activity.py <pid> --analyze-only     # re-analyze saved samples
Artifacts: /tmp/fp-tensors.tsv, /tmp/fp-ibs.perf, /tmp/fp-addrs.txt
"""
import bisect, json, os, re, struct, subprocess, sys, threading, time
import urllib.request

MODEL_DIR = "/home/dank/models/Qwen3.8-Flash-Next-GGUF/Q8_0"
N_PARTS, N_BLOCKS, N_EXPERTS = 6, 48, 512
TSV, PERF_OUT, ADDRS = "/tmp/fp-tensors.tsv", "/tmp/fp-ibs.perf", "/tmp/fp-addrs.txt"
PORT = 7770

# ---------------------------------------------------------------- GGUF header
class P:
    def __init__(self, f): self.f, self.ver = f, 2
    def rd(self, n):
        b = self.f.read(n)
        assert len(b) == n, "eof"
        return b
    def u8(self): return self.rd(1)[0]
    def u32(self): return struct.unpack("<I", self.rd(4))[0]
    def u64(self): return struct.unpack("<Q", self.rd(8))[0]
    def s(self):
        n = self.u64() if self.ver >= 3 else self.u32()
        return self.rd(n).decode("utf-8", "replace")
    def val(self, t):
        if t == 8: return self.s()
        if t == 9:
            et, n = self.u32(), self.u64()
            for _ in range(n): self.val(et)
            return None
        if t in (0, 7): return self.u8()
        if t == 1: return struct.unpack("<b", self.rd(1))[0]
        if t == 2: return struct.unpack("<H", self.rd(2))[0]
        if t == 3: return struct.unpack("<h", self.rd(2))[0]
        if t == 4: return self.u32()
        if t == 5: return struct.unpack("<i", self.rd(4))[0]
        if t == 6: return struct.unpack("<f", self.rd(4))[0]
        if t == 10: return self.u64()
        if t == 11: return struct.unpack("<q", self.rd(8))[0]
        if t == 12: return struct.unpack("<d", self.rd(8))[0]
        raise ValueError(t)

def dump_part(path):
    """[(name, start, end)] with extents derived from consecutive offsets."""
    with open(path, "rb") as f:
        p = P(f)
        assert p.rd(4) == b"GGUF"
        p.ver = p.u32()
        ntens, nkv = p.u64(), p.u64()
        for _ in range(nkv):
            p.s(); p.val(p.u32())
        infos = []
        for _ in range(ntens):
            name = p.s()
            nd = p.u32()
            for _ in range(nd): p.u64()
            p.u32()
            infos.append((name, p.u64()))
        data_start = (f.tell() + 31) & ~31
    fsz = os.path.getsize(path)
    out = []
    for i, (name, off) in enumerate(infos):
        start = data_start + off
        end = (data_start + infos[i + 1][1]) if i + 1 < len(infos) else fsz
        out.append((name, start, end))
    return out

# ------------------------------------------------------------------- workload
PROMPTS = [
    [{"role": "user", "content": "Write a Rust function that parses a simple binary log format (u32 length-prefixed records) with proper error handling, then explain the unsafe choices you avoided."}],
    [{"role": "user", "content": "Compare the strategic logic of Soviet deep-battle doctrine and Blitzkrieg: where did they agree, where did they diverge, and what did each assume about logistics that proved wrong?"}],
    [{"role": "system", "content": "You have tools: get_weather(city), book_flight(from,to,date). Plan a 3-city trip for next week; call tools as needed and explain each decision."},
     {"role": "user", "content": "Seattle, Denver, Austin. Cheapest reasonable plan."}],
    [{"role": "user", "content": "Explain how DNS recursive resolution works end to end, then write a small Python stub that walks a referral chain."}],
]

def fire(idx, results, dur):
    body = json.dumps({"messages": PROMPTS[idx], "max_tokens": 260,
                       "temperature": 0.7, "cache_prompt": False}).encode()
    req = urllib.request.Request(f"http://127.0.0.1:{PORT}/v1/chat/completions",
                                 data=body, headers={"Content-Type": "application/json"})
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=dur + 60) as r:
            d = json.load(r)
        results[idx] = (r.status, round(time.time() - t0),
                        len(d.get("choices", [{}])[0].get("message", {}).get("content", "")))
    except Exception as e:
        results[idx] = ("ERR", str(e)[:120])

# --------------------------------------------------------------------- main
def main():
    pid = int(sys.argv[1])
    analyze_only = "--analyze-only" in sys.argv
    dur = 75.0
    for a in sys.argv[2:]:
        if a.isdigit(): dur = float(a)
    parts = {i: f"{MODEL_DIR}/Qwen3.8-Flash-Next-Q8_0-0000{i}-of-0000{N_PARTS}.gguf"
             for i in range(1, N_PARTS + 1)}

    vmas = []
    with open(f"/proc/{pid}/maps") as f:
        for line in f:
            m = re.match(r"([0-9a-f]+)-([0-9a-f]+) (\S+) ([0-9a-f]+) (\S+) (\d+)\s*(.*)", line)
            if not m: continue
            start, end, pgoff, path = int(m.group(1), 16), int(m.group(2), 16), int(m.group(4), 16), m.group(7)
            for i, p in parts.items():
                if path == p: vmas.append((start, end, pgoff, i))
    vmas.sort()
    assert vmas, "no GGUF VMAs in maps — wrong pid?"
    print(f"[rig] {len(vmas)} GGUF VMAs across {len({v[3] for v in vmas})} parts")

    tensors = {}
    with open(TSV, "w") as tsv:
        for i, p in parts.items():
            for name, start, end in dump_part(p):
                tsv.write(f"{p}\t{name}\t{start}\t{end}\n")
                tensors.setdefault(i, ([], [], []))
                tensors[i][0].append(start); tensors[i][1].append(end); tensors[i][2].append(name)
    n_exp = sum(1 for i in tensors for nm in tensors[i][2] if "exps" in nm)
    print(f"[rig] {n_exp} exps tensors, extents offset-derived")

    if not analyze_only:
        results = {}
        perf = subprocess.Popen(["sudo", "perf", "record", "-e", "ibs_op//", "-c", "50000",
                                 "--data", "-o", PERF_OUT, "-p", str(pid), "--", "sleep", str(dur)],
                                stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
        time.sleep(2)
        threads = [threading.Thread(target=fire, args=(i, results, dur)) for i in range(4)]
        for t in threads: t.start()
        for t in threads: t.join(dur + 120)
        err = perf.communicate(timeout=60)[1].decode()[-200:]
        for i, r in sorted(results.items()):
            print(f"[work] slot{i}: {r}")
        print("[perf]", err.strip())
        with open(ADDRS, "w") as f:
            f.write(subprocess.run(["sudo", "perf", "script", "-i", PERF_OUT, "-F", "addr"],
                                    capture_output=True, text=True).stdout)

    block, expert = {}, {}
    n_exps = n_ple = n_named = n_gap = n_out = 0
    for line in open(ADDRS):
        tok = line.split()
        if not tok: continue
        try: a = int(tok[0], 16)
        except ValueError: continue
        if a == 0 or a > 0xffff000000000000: continue
        vi = bisect.bisect_right(vmas, (a, 1 << 62, 0, 0)) - 1
        if vi < 0: n_out += 1; continue
        start, end, pgoff, pi = vmas[vi]
        if not start <= a < end: n_out += 1; continue
        foff = pgoff + (a - start)
        st, en, nm = tensors[pi]
        ti = bisect.bisect_right(st, foff) - 1
        if ti < 0 or foff >= en[ti]: n_gap += 1; continue
        name, lo, span = nm[ti], foff - st[ti], en[ti] - st[ti]
        if "exps" in name:
            n_exps += 1
            b = int(re.match(r"blk\.(\d+)\.", name).group(1))
            e = int(lo / (span / N_EXPERTS))
            block[b] = block.get(b, 0) + 1
            expert[(b, e)] = expert.get((b, e), 0) + 1
        elif "per_layer" in name: n_ple += 1
        else: n_named += 1
    print(f"\n[rig] attribution: exps={n_exps} ple={n_ple} named={n_named} "
          f"gap={n_gap} outside={n_out}")
    tot = sum(block.values())
    print(f"== BLOCK histogram ({tot} samples, {len(block)} CPU-resident blocks):")
    for b, c in sorted(block.items(), key=lambda kv: -kv[1]):
        print(f"  blk{b:2d}: {c:7d} {100*c/tot:5.2f}%")
    mean = tot / max(len(block), 1)
    print(f"block spread: max={max(block.values())/mean:.3f}x mean, min={min(block.values())/mean:.3f}x mean")
    ev = sorted(expert.values(), reverse=True); et = sum(ev)
    print(f"== EXPERT skew: {len(ev)} distinct (block,expert) cells hit of {N_BLOCKS*N_EXPERTS}")
    print(f"   top-half of cells carry {100*sum(ev[:len(ev)//2])/et:.1f}%; "
          f"top-10% {100*sum(ev[:len(ev)//10])/et:.1f}%; hottest cell {100*ev[0]/et:.3f}%")

if __name__ == "__main__":
    main()