"""基线：单问题 / 多问题 延迟拆解（tokenize vs GPU forward），以及 profiler top kernels。"""
import os, sys, time, torch
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "apps"))
from common import get_agent
from laya.common import build_sequence, collate_items, QTYPES, render_options

agent = get_agent(os.environ.get("LAYA_VARIANT", "multilingual"))
state = {"subject": "Duplicate charge on invoice 4411", "body": "We were billed twice for March. Please refund the duplicate or we're moving to a competitor. " * 3}
Q = {"department": {"type": "choice", "instructions": "Which team should handle this?", "criteria": {"billing": "invoices, refunds", "technical": "bugs, outages", "sales": "pricing"}},
     "urgency": {"type": "score", "instructions": "How urgent is this?", "criteria": ["not urgent", "soon", "blocking"]},
     "churn": {"type": "noul", "instructions": "Does the user threaten to cancel?"}}
def qs(n): return {f"{k}{i}": v for i in range(n) for k, v in Q.items()}

def timeit(fn, iters=20):
    for _ in range(3): fn()
    torch.cuda.synchronize(); t = time.perf_counter()
    for _ in range(iters): fn()
    torch.cuda.synchronize(); return (time.perf_counter() - t) / iters * 1000

for n in (1, 3, 10, 30):
    q = qs(n) if n != 1 else {"department": Q["department"]}
    ms = timeit(lambda: agent.predict(state, q))
    # tokenization-only cost
    def tok_only():
        items = []
        for qid in q:
            qq = agent._to_internal(q[qid]); seq, m = build_sequence(agent.tok, state, qq, agent.cfg["max_len"], agent.cfg["head_max_len"])
            items.append({"ids": seq, "markers": m, "qtype": QTYPES[qq["t"]]})
        return collate_items([items], agent.tok.pad_token_id)
    tms = timeit(tok_only)
    b = tok_only()
    print(f"questions={len(q):3d}  seq_len={b['input_ids'].shape[1]:4d}  predict={ms:7.2f} ms  tokenize={tms:6.2f} ms  gpu+post={ms-tms:7.2f} ms  per-q={ms/len(q):5.2f} ms")

# profile the model forward at n=10
q = qs(10); b = None
items = []
for qid in q:
    qq = agent._to_internal(q[qid]); seq, m = build_sequence(agent.tok, state, qq, agent.cfg["max_len"], agent.cfg["head_max_len"])
    items.append({"ids": seq, "markers": m, "qtype": QTYPES[qq["t"]]})
b = {k: v.to("cuda") for k, v in collate_items([items], agent.tok.pad_token_id).items() if torch.is_tensor(v)}
def fwd():
    with torch.no_grad(), torch.autocast("cuda", dtype=agent.dtype):
        return agent.model(b["input_ids"], b["attention_mask"], b["marker_pos"], b["marker_mask"], b["qtype"])
print(f"\nmodel forward only (30 q, L={b['input_ids'].shape[1]}): {timeit(fwd):.2f} ms  dtype={agent.dtype}")
from torch.profiler import profile, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for _ in range(5): fwd()
    torch.cuda.synchronize()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=25))