"""Record which experts a MoE onw model picks, per layer and per step, over a set of realistic chats (text, code,
image, multi-turn). Writes a JSON trace that sim_expert_cache.py replays against expert caches of different sizes.
usage: python trace_experts.py MODEL_DIR OUT.json [IMAGE]"""
import json, sys, time
from PIL import Image
from onw.chat import ChatEngine
from onw.runtime import SegmentedModel

TRACE = []                                   # one entry per _bind_experts call: [layer, S, n, [experts...]]
_orig = SegmentedModel._bind_experts


def _spy(self, req, m, idx, w, S, n):
    TRACE.append([m["consume_moe"], S, n, sorted({int(e) for e in idx[:n].reshape(-1)})])
    return _orig(self, req, m, idx, w, S, n)


SegmentedModel._bind_experts = _spy

CODE = '''def load_config(path):
    with open(path) as f:
        data = json.load(f)
    if "name" not in data:
        raise ValueError("missing name")
    data["port"] = int(data.get("port", 8080))
    return data'''
CHATS = [
    ("chat-ja", [["NPUとGPUの違いを、身近なたとえを使って説明してください。"],
                 ["では、ノートPCでLLMを動かすならどちらが向いていますか？理由も。"]]),
    ("code", [[f"次のPython関数に型ヒントとdocstringを付けて、関数全体を書き直してください。\n\n```python\n{CODE}\n```"]]),
    ("summary-ja", [["次の文章を3文で要約してください。\n\n" + "東京は日本の首都であり、政治・経済・文化の中心地である。"
                     "江戸時代には徳川幕府が置かれ、1868年の明治維新で東京と改称された。現在は約1400万人が暮らす世界有数の大都市で、"
                     "多くの企業や大学が集まっている。一方で、人口集中による住宅価格の高騰や通勤の混雑、災害への備えといった課題も抱えている。" * 2]]),
    ("chat-en", [["Explain the difference between a process and a thread, with a short Python example."]]),
    ("math-ja", [["ある商品を定価の2割引きで買うと960円でした。定価はいくらですか？途中の式も書いてください。"]]),
]


def main():
    d, out = sys.argv[1], sys.argv[2]
    img = Image.open(sys.argv[3]) if len(sys.argv) > 3 else None
    e = ChatEngine(d, "NPU", pld=False)
    runs = []
    chats = list(CHATS)
    if img is not None:
        chats.insert(1, ("image", [[{"type": "image", "image": img}, {"type": "text", "text": "この画像について詳しく説明してください。"}],
                                   ["その動物の気持ちを想像して、短い物語にしてください。"]]))
    for name, turns in chats:
        e.checkpoint = None
        msgs = []
        for t in turns:
            content = t[0] if len(t) == 1 and isinstance(t[0], str) else t
            msgs.append({"role": "user", "content": content})
            start = len(TRACE)
            parts, st, t0 = [], None, time.time()
            for x in e.stream_chat([dict(m) for m in msgs], 200):
                if isinstance(x, dict):
                    st = x
                else:
                    parts.append(x)
            text = "".join(parts)
            msgs.append({"role": "assistant", "content": text})
            runs.append({"name": name, "start": start, "end": len(TRACE), "prompt_tokens": st["prompt_tokens"],
                         "cached": st["cached_tokens"], "completion": st["completion_tokens"]})
            print(f"{name}: prompt {st['prompt_tokens']} (reused {st['cached_tokens']}), {st['completion_tokens']} tokens, "
                  f"{time.time() - t0:.1f}s | {text[:50]!r}", flush=True)
    meta = {k: e.model.meta.get(k) for k in ("arch",)}
    json.dump({"E": e.model.E, "K": e.model.K, "meta": meta, "runs": runs, "trace": TRACE}, open(out, "w"))
    print(f"wrote {len(TRACE)} binds to {out}")


if __name__ == "__main__":
    main()