#!/usr/bin/env python3
"""suggest_ot.py — deterministic -ot placement search for dual-GPU llama.cpp.

Principle: scripts do measurement and search-space generation; the local
model only interprets the compact output and picks/policy-decides. This tool
never emits free-form -ot regexes — only candidates derived from the actual
model tensor manifest.

Subcommands:
  manifest   GGUF [--out PATH]
      Build (and cache) a tensor manifest: per-layer type (attn/ssm/nextn),
      per-tensor MiB, per-layer totals, and precomputed tensor families.

  candidates --manifest M --split "0.6,0.4" [--window W]
      Print the constrained candidate set around the layer boundary:
      whole blocks near the boundary + coherent tensor families, with the
      exact -ot regex, bytes moved, and direction for each.

  evaluate   --history H --manifest M [--key-json K] [--safety N]
      Take the latest history run, compute per-device free / per-1000-token
      cost / context capacity, simulate every candidate, and print a ranked
      table with a recommendation and a stop/balance verdict.

  summary    --history H [--key-json K]
      Compact run table (for the supervisor model), like:
        run  ot                ctx     free0   free1  load
        1    none              110848  1417    1001   ok

  record     --history H --run-json R [--key-json K]
      Append one run record (usually called by run_probe.sh).
"""
import argparse
import json
import math
import os
import re
import subprocess
import sys

MiB = 1024 * 1024
CTX_TRAIN_DEFAULT = 262144

# ggml type -> (block_size, bytes_per_block)
GGML_SIZES = {
    0: (1, 4),    # F32
    1: (1, 2),    # F16
    2: (32, 34),  # Q8_0
    6: (32, 18),  # Q4_0
    7: (1, 1),    # I8
    8: (1, 2),    # I16
    9: (1, 4),    # I32
    10: (1, 8),   # F64
    11: (256, 146),  # Q4_K
    12: (256, 178),  # Q5_K
    13: (256, 210),  # Q6_K
    16: (1, 2),    # BF16
    22: (32, 18),  # IQ4_NL
}

GGUF_VAL_TYPES = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1,
                  8: None, 9: None, 10: 8, 11: 8, 12: 8}


# ---------------------------------------------------------------- manifest

def gguf_bytes_from_dump(path):
    """Tensor sizes via gguf-dump (preferred)."""
    out = subprocess.run(["gguf-dump", path], capture_output=True, text=True, timeout=300).stdout
    rows = []
    for line in out.splitlines():
        m = re.match(r"\s*\d+:\s+(\d+)\s*\|\s*(.*?)\|\s*(\S+)\s*\|\s*(\S+)\s*$", line)
        if m:
            rows.append((int(m.group(1)), m.group(3), m.group(4)))
    return rows


def _gguf_read_kv(buf, off):
    import struct
    t = buf[off]
    off += 1
    if t in GGUF_VAL_TYPES and GGUF_VAL_TYPES[t] is not None:
        return off + GGUF_VAL_TYPES[t]
    if t == 8:  # string
        n = struct.unpack_from("<Q", buf, off)[0]
        off += 8 + n
        off += (8 - off % 8) % 8
        return off
    if t == 9:  # array
        et = buf[off]
        cnt = struct.unpack_from("<Q", buf, off + 1)[0]
        off += 9
        for _ in range(cnt):
            off = _gguf_read_kv(buf, off)
        return off
    raise ValueError(f"bad gguf value type {t}")


def gguf_bytes_native(path):
    """Pure-python fallback: read only header + tensor table (mmap, fast)."""
    import mmap
    import struct
    with open(path, "rb") as f:
        mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    if mm[:4] != b"GGUF":
        raise ValueError("not a GGUF file")
    off = 4 + 4  # magic + version
    n_tensors = struct.unpack_from("<Q", mm, off)[0]
    off += 8
    n_kv = struct.unpack_from("<Q", mm, off)[0]
    off += 8
    for _ in range(n_kv):
        off = _gguf_read_kv(mm, off)
    rows = []
    for _ in range(n_tensors):
        nlen = struct.unpack_from("<Q", mm, off)[0]
        off += 8
        name = mm[off:off + nlen].decode()
        off += nlen + (8 - off % 8) % 8
        ndim = struct.unpack_from("<I", mm, off)[0]
        off += 4
        dims = struct.unpack_from(f"<{ndim}Q", mm, off)
        off += 8 * ndim
        t = struct.unpack_from("<I", mm, off)[0]
        off += 4
        numel = 1
        for d in dims:
            numel *= d
        bs, bb = GGML_SIZES.get(t, (1, 1))
        rows.append((math.ceil(numel / bs) * bb, f"t{t}", name))
    mm.close()
    return rows


def build_manifest(path, out=None):
    model_file = os.path.basename(path)
    try:
        rows = gguf_bytes_from_dump(path)
        source = "gguf-dump"
    except Exception as e:
        print(f"warning: gguf-dump failed ({e}); using native parser", file=sys.stderr)
        rows = gguf_bytes_native(path)
        source = "native"

    layers = {}
    other = {}
    for size, typ, name in rows:
        m = re.match(r"blk\.(\d+)\.(.*)$", name)
        if m:
            L, rest = int(m.group(1)), m.group(2)
            d = layers.setdefault(L, {"tensors": {}})
            d["tensors"][name] = round(size / MiB, 3)
        else:
            other[name] = round(size / MiB, 3)

    for L, d in layers.items():
        d["total_mib"] = round(sum(d["tensors"].values()), 2)
        rests = [t.split(".", 2)[2] for t in d["tensors"] if t.count(".") >= 2]
        d["type"] = ("nextn" if any(r.startswith("nextn") for r in rests) else
                     "attn" if any(r.startswith("attn_q.") for r in rests)
                     else "ssm")

    n_layer = max(L for L in layers if layers[L]["type"] != "nextn") if layers else 0
    man = {
        "schema": 1,
        "model_file": model_file,
        "source": source,
        "n_layer": n_layer,
        "attn_layers": sorted(L for L in layers if layers[L]["type"] == "attn"),
        "ssm_layers": sorted(L for L in layers if layers[L]["type"] == "ssm"),
        "layers": layers,
        "other": other,
        "total_mib": round(sum(d["total_mib"] for d in layers.values()) + sum(other.values()), 2),
    }
    if out:
        os.makedirs(os.path.dirname(out), exist_ok=True)
        with open(out, "w") as f:
            json.dump(man, f, indent=1)
        print(f"manifest written: {out} ({man['total_mib']} MiB, {man['n_layer']} layers, "
              f"{len(man['attn_layers'])} attn / {len(man['ssm_layers'])} ssm)")
    return man


def default_manifest_path(model_file, data_dir):
    return os.path.join(data_dir, "manifests", model_file + ".manifest.json")


def load_manifest(args):
    if args.manifest and os.path.exists(args.manifest):
        man = json.load(open(args.manifest))
        man["layers"] = {int(k): v for k, v in man["layers"].items()}
        return man
    return None


# ------------------------------------------------------------- candidates

def boundary_layer(man, split):
    """Layer index of the last device-0 layer under --split-mode layer."""
    parts = [float(x) for x in split.split(",")]
    p0 = parts[0] / sum(parts)
    # GPU positions = repeating layers + nextn (if present) + output layer
    has_nextn = any(d["type"] == "nextn" for d in man["layers"].values())
    N = (man["n_layer"] + 1) + (1 if has_nextn else 0) + 1
    b = math.ceil(N * p0)
    return max(0, min(b - 1, man["n_layer"] - 1)), N, b


def layer_families(man, L):
    """Coherent tensor groups for one layer (regex, bytes, moves_kv=False)."""
    d = man["layers"][L]
    t = d["tensors"]
    fams = []

    def fam(name, pat):
        names = [k for k in t
                 if k.count(".") >= 2 and re.match(pat, k.split(".", 2)[2])]
        if names:
            fams.append({"family": name, "regex": f"^blk\\.{L}\\." + pat + "$",
                         "mib": round(sum(t[k] for k in names), 2), "tensors": len(names),
                         "moves_kv": False})

    fam("ffn_all", r"ffn_(gate|up|down)\.weight")
    fam("ffn_gate_up", r"ffn_(gate|up)\.weight")
    fam("ffn_down", r"ffn_down\.weight")
    if d["type"] == "attn":
        fam("attn_all", r"attn_(q|k|v|output)\.weight")
        fam("attn_out_kv", r"attn_(output|k|v)\.weight")
        fam("attn_out", r"attn_output\.weight")
        fam("attn_q", r"attn_q\.weight")
    elif d["type"] == "ssm":
        fam("ssm_core", r"ssm_(out\.weight|alpha\.weight|beta\.weight|conv1d\.weight)")
        fam("ssm_out", r"ssm_out\.weight")
        fam("attn_qkv", r"attn_qkv\.weight")
        fam("attn_gate", r"attn_gate\.weight")
    return fams


def gen_candidates(man, split, window=3, devs=("ROCm0", "ROCm1")):
    b, N, bpos = boundary_layer(man, split)
    cands = []
    for L in range(max(0, b - window), min(man["n_layer"], b + window + 1) + 1):
        d = man["layers"][L]
        if d["type"] == "nextn":
            continue
        dev = devs[0] if L <= b else devs[1]
        dest = devs[1] if L <= b else devs[0]
        cands.append({
            "kind": "block",
            "layer": L,
            "family": f"blk.{L} whole",
            "regex": f"^blk\\.{L}\\.\\w[\\w.]*$",
            "ot": f"^blk\\.{L}\\.\\w[\\w.]*$={dest}",
            "mib": d["total_mib"],
            "tensors": len(d["tensors"]),
            "moves_kv": False,  # verified: KV/RS follow layer split, not -ot
            "attn_layers": 1 if d["type"] == "attn" else 0,
            "from": dev,
            "to": dest,
        })
        for f in layer_families(man, L):
            cands.append({
                "kind": "family",
                "layer": L,
                "family": f"blk.{L}.{f['family']}",
                "regex": f["regex"],
                "ot": f["regex"] + f"={dest}",
                "mib": f["mib"],
                "tensors": f["tensors"],
                "moves_kv": False,
                "attn_layers": 0,
                "from": dev,
                "to": dest,
            })
    return {"boundary_layer": b, "N_positions": N, "boundary_positions": bpos,
            "candidates": cands}


# --------------------------------------------------------------- evaluate

def history_records(history_path, key_json=None):
    recs = []
    with open(history_path) as f:
        for line in f:
            line = line.strip()
            if line:
                recs.append(json.loads(line))
    if key_json:
        key = json.loads(key_json)
        for k, v in key.items():
            recs = [r for r in recs if r.get("key", {}).get(k) == v]
    return recs


def latest_run(history_path, key_json=None):
    recs = history_records(history_path, key_json)
    return recs[-1] if recs else None


def evaluate(history_path, man, safety=150.0, key_json=None, ctx_train=CTX_TRAIN_DEFAULT,
             baseline="best"):
    """Fitter-internal capacity model (calibrated on this fork's -fit).

    The fitter reduces ctx until, per device d:
        used_d(ctx) = proj_d - c_d * (ctx_init - ctx) <= limit_d
    where proj_d = projected use at ctx_init (from the fit log),
    limit_d = total_d - target_free_d (fitter's internal per-device target),
    c_d = fitter's per-1000-token slope (>= measured KV/1k; includes hidden
    ctx-scaling costs like MTP/ngram/checkpoint reserves).

    The limiting device is the one whose implied c_d (from the fit itself)
    is >= its measured KV/1k; the other device's implied c_d is a lower
    bound (it still has headroom at the fitted ctx).

    Weight-only -ot moves change F_d = proj_d - c_d*ctx_init by the moved
    model bytes (RS/compute/spec stay put — RS placement follows the layer
    split, NOT -ot; verified empirically).
    """
    rec = latest_run(history_path, key_json)
    if not rec:
        print("error: no matching run in history", file=sys.stderr)
        return 1
    if baseline == "best":
        recs = history_records(history_path, key_json)
        recs = [r for r in recs if r["run"].get("ctx_fitted")]
        if recs:
            rec = max(recs, key=lambda r: r["run"]["ctx_fitted"])
    run = rec["run"]
    base_ots = [s for s in (run.get("ot") or "").split() if s]
    tested = set()
    for r in history_records(history_path, key_json):
        for s in ((r["run"].get("ot") or "")).split():
            tested.add(s)
    ctx = run["ctx_fitted"]
    devs = [d for d in run.get("devices", {}) if d != "CPU"]
    if len(devs) != 2:
        print(f"error: expected 2 GPU devices in run, got: {devs}", file=sys.stderr)
        return 1
    proj = run.get("fit_proj_mib", {})
    tfree = run.get("fit_target_free_mib", {})
    if not proj or not tfree:
        print("error: run lacks fit projection data; re-parse the log with the current parser",
              file=sys.stderr)
        return 1

    ctx_init_k = run.get("ctx_reduced_from", 262144) / 1000.0
    ctx_k = ctx / 1000.0
    limit = {d: run["devices"][d]["total_mib"] - tfree[d] for d in devs}

    # measured KV/1k per device (incl. draft on ROCm1) and per attn layer
    kv_measured = {d: run["kv_mib"].get(d, 0.0) / ctx_k for d in devs}
    kv_attn_per1k = (sum(run["kv_mib"].values()) - run.get("kv_draft_mib", 0.0)) \
        / ctx_k / max(1, len(man["attn_layers"]))

    # implied c_d from the fit; limiting device = implied >= measured
    implied = {d: (proj[d] - limit[d]) / max(1e-9, ctx_init_k - ctx_k) for d in devs}
    limiting = [d for d in devs if implied[d] >= kv_measured[d] - 0.05]
    if not limiting:
        limiting = [max(devs, key=lambda d: implied[d])]
    c = {d: implied[d] if d in limiting else kv_measured[d] for d in devs}
    # fitter fixed cost at this placement
    F = {d: proj[d] - c[d] * ctx_init_k for d in devs}

    split = (rec.get("run", {}).get("tensor_split")
             or rec.get("key", {}).get("tensor_split") or "0.6,0.4")
    d0, d1 = devs
    cg = gen_candidates(man, split, devs=devs)
    cands = cg["candidates"]

    def simulate(moves):
        c0n, c1n = c[d0], c[d1]
        m0 = m1 = 0.0
        for mv in moves:
            if mv["to"] == d1:
                m1 += mv["mib"]
                if mv["moves_kv"]:
                    c0n -= mv["attn_layers"] * kv_attn_per1k
                    c1n += mv["attn_layers"] * kv_attn_per1k
            else:
                m0 += mv["mib"]
                if mv["moves_kv"]:
                    c0n += mv["attn_layers"] * kv_attn_per1k
                    c1n -= mv["attn_layers"] * kv_attn_per1k
        cap0 = (limit[d0] - (F[d0] + m0 - m1)) / c0n if c0n > 0 else 0.0
        cap1 = (limit[d1] - (F[d1] + m1 - m0)) / c1n if c1n > 0 else 0.0
        feas = min(cap0, cap1) > 0
        ctxp = int(min(ctx_train, min(cap0, cap1) * 1000) // 256 * 256) if feas else 0
        return feas, ctxp, round(cap0, 1), round(cap1, 1)

    def brief(mv):
        return {"ot": mv["ot"], "what": mv["family"], "mib": mv["mib"],
                "kind": mv["kind"], "layer": mv["layer"], "from": mv["from"], "to": mv["to"],
                "tested": mv["ot"] in tested}

    def full_command(ots):
        parts = " ".join(f"'{o}'" for o in ots)
        return f"run_probe.sh <launcher> {parts}".strip()

    # Candidates are generated from the manifest, so detect overlap using the
    # actual tensor names rather than comparing regex strings.  In particular,
    # a whole-block override covers every family in that block; adding a child
    # family would otherwise double-count its bytes in simulate().
    def covered_by_base(cv):
        try:
            cand_re = re.compile(cv["regex"])
            cand_names = [name for name in man["layers"][cv["layer"]]["tensors"]
                          if cand_re.fullmatch(name)]
            if not cand_names:
                return False
            base_rules = []
            for raw in base_ots:
                pattern, sep, dest = raw.rpartition("=")
                if not sep:
                    continue
                try:
                    base_rules.append((re.compile(pattern), dest))
                except re.error:
                    continue
            # Exclude any overlap.  Partial overlap is unsafe because this
            # evaluator cannot represent a residual move with one aggregate
            # MiB value; overlap to a different device would also make the
            # command's effective placement ambiguous.
            return any(any(rule.fullmatch(name) for rule, _ in base_rules)
                       for name in cand_names)
        except (KeyError, re.error):
            return False

    rows = []
    for cv in cands:
        if cv["ot"] in base_ots or cv["ot"] in tested or covered_by_base(cv):
            continue
        feas, ctxp, cap0, cap1 = simulate([cv])
        full = base_ots + [cv["ot"]]
        rows.append({"cand": brief(cv), "cap0_k": cap0, "cap1_k": cap1,
                     "ctx_pred": ctxp, "delta": ctxp - ctx, "feasible": feas,
                     "full_ot": full, "command": full_command(full)})
    rows.sort(key=lambda r: (-r["delta"], r["cand"]["mib"]))

    # greedy multi-move (one move per layer; families of one layer conflict)
    best_seqs = {1: None, 2: None, 3: None}
    for depth in (1, 2, 3):
        seq, used_layers, base_ctx = [], set(), ctx
        for _ in range(depth):
            best = None
            for cv in cands:
                if cv["layer"] in used_layers:
                    continue
                if cv["ot"] in base_ots or cv["ot"] in tested or covered_by_base(cv):
                    continue
                feas, ctxp, cap0, cap1 = simulate(seq + [cv])
                if not feas or ctxp - base_ctx < 256:
                    continue
                if best is None or ctxp > best[0]:
                    best = (ctxp, cv, cap0, cap1)
            if best is None:
                break
            base_ctx, _, _, _ = best
            seq.append(best[1])
            used_layers.add(best[1]["layer"])
        if len(seq) == depth:
            feas, ctxp, cap0, cap1 = simulate(seq)
            full = base_ots + [m["ot"] for m in seq]
            best_seqs[depth] = {"moves": [brief(m) for m in seq],
                                "cap0_k": cap0, "cap1_k": cap1,
                                "ctx_pred": ctxp, "delta": ctxp - ctx, "feasible": feas,
                                "full_ot": full, "command": full_command(full)}

    best = best_seqs[3] or best_seqs[2] or best_seqs[1] or (rows[0] if rows else None)
    stop = (best is None or best.get("delta", 0) < 256)

    out = {
        "schema": 2,
        "model": {
            "fit": {
                "ctx": ctx, "ctx_init": int(ctx_init_k * 1000),
                "proj_mib": proj, "limit_mib": {d: limit[d] for d in devs},
                "c_per1k": {d: round(c[d], 3) for d in devs},
                "kv_measured_per1k": {d: round(kv_measured[d], 3) for d in devs},
                "F_fixed_mib": {d: round(F[d], 1) for d in devs},
                "cap_k": {d: round((limit[d] - F[d]) / c[d], 1) for d in devs},
                "limiting": limiting,
            },
            "rocm_smi_free_mib": run.get("free_mib", {}),
        },
        "safety_margin_mib": safety,
        "stop": stop,
        "stop_reason": ("best predicted gain < 256 tokens" if best and best.get("delta", 0) < 256
                        else None),
        "best_sequences": best_seqs,
        "ranked_single": rows[:8],
    }
    print(json.dumps(out, indent=1))
    return 0


def summary(history_path, key_json=None):
    recs = []
    with open(history_path) as f:
        for line in f:
            line = line.strip()
            if line:
                recs.append(json.loads(line))
    if key_json:
        key = json.loads(key_json)
        for k, v in key.items():
            recs = [r for r in recs if r.get("key", {}).get(k) == v]
    print(f"{'run':>3}  {'ot':<52} {'ctx':>7} {'free0':>8} {'free1':>8} {'load':>5} {'tg':>6}")
    for i, r in enumerate(recs, 1):
        run = r["run"]
        ot = (run.get("ot") or "none")[:52]
        d = [x for x in run.get("devices", {}) if x != "CPU"] or ["ROCm0", "ROCm1"]
        f0 = run.get("free_mib", {}).get(d[0])
        f1 = run.get("free_mib", {}).get(d[1] if len(d) > 1 else "")
        tg = run.get("tg_tps")
        print(f"{i:>3}  {ot:<52} {run.get('ctx_fitted') or '-':>7} "
              f"{f0 if f0 is not None else '-':>8} {f1 if f1 is not None else '-':>8} "
              f"{('ok' if run.get('load_ok') else 'FAIL'):>5} {tg if tg is not None else '-':>6}")


# -------------------------------------------------------------------- main

def main():
    ap = argparse.ArgumentParser(description=__doc__)
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("manifest")
    p.add_argument("gguf")
    p.add_argument("--out", default=None)

    p = sub.add_parser("candidates")
    p.add_argument("--manifest", default=None)
    p.add_argument("--gguf", default=None)
    p.add_argument("--split", default="0.6,0.4")
    p.add_argument("--window", type=int, default=4)
    p.add_argument("--devices", default="ROCm0,ROCm1",
                   help="device names (comma-separated), e.g. Vulkan0,Vulkan1")
    p.add_argument("--data-dir", default=None)

    p = sub.add_parser("evaluate")
    p.add_argument("--history", required=True)
    p.add_argument("--manifest", default=None)
    p.add_argument("--gguf", default=None)
    p.add_argument("--key-json", default=None)
    p.add_argument("--safety", type=float, default=150.0)
    p.add_argument("--baseline", choices=["best", "latest"], default="best")
    p.add_argument("--data-dir", default=None)

    p = sub.add_parser("summary")
    p.add_argument("--history", required=True)
    p.add_argument("--key-json", default=None)

    p = sub.add_parser("record")
    p.add_argument("--history", required=True)
    p.add_argument("--run-json", required=True)
    p.add_argument("--key-json", default=None)

    a = ap.parse_args()
    here = os.path.dirname(os.path.abspath(__file__))
    data_dir = getattr(a, "data_dir", None) or os.path.normpath(os.path.join(here, "..", "data"))

    if a.cmd == "manifest":
        out = a.out or default_manifest_path(os.path.basename(a.gguf), data_dir)
        build_manifest(a.gguf, out)
        return 0

    man = None
    if a.cmd in ("candidates", "evaluate"):
        man = load_manifest(a) if hasattr(a, "manifest") else None
        if man is None and getattr(a, "gguf", None):
            man = build_manifest(a.gguf, None)
        if man is None:
            print("error: --manifest or --gguf required", file=sys.stderr)
            return 2

    if a.cmd == "candidates":
        devs = tuple(x.strip() for x in a.devices.split(",") if x.strip())
        print(json.dumps(gen_candidates(man, a.split, a.window, devs=devs), indent=1))
    elif a.cmd == "evaluate":
        return evaluate(a.history, man, a.safety, a.key_json, ctx_train=CTX_TRAIN_DEFAULT,
                        baseline=a.baseline)
    elif a.cmd == "summary":
        summary(a.history, a.key_json)
    elif a.cmd == "record":
        run = json.load(open(a.run_json))
        key = json.loads(a.key_json) if a.key_json else {}
        with open(a.history, "a") as f:
            f.write(json.dumps({"schema": 1, "key": key, "run": run}) + "\n")
        print(f"recorded run {run.get('timestamp')} ctx={run.get('ctx_fitted')}")
    return 0


if __name__ == "__main__":
    sys.exit(main())