# -*- coding: utf-8 -*-
"""SamAI Pnet-dMoE 2B (Ponder-MoE) — HF 参考实现对话 Demo (ZeroGPU).

架构: PonderNet 自适应停机(后 8 层思考块) + 动态 MoE(8 专家阈值路由, 仅 ponder 层)
     + MTP(训练期辅助头, 推理自动旁路).
加载: transformers>=5.0 trust_remote_code, 建模代码为仓内单文件 modeling_samai_pnet.py.
诊断: model._ponder_log 每次 forward 追加 {mode, executed, steps_mean, steps_max, moe_active}.
"""
import time
from threading import Thread

import spaces  # 必须在 torch / transformers 之前 import (monkey-patch torch.cuda)
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MODEL_ID = "tchbcb/samai-pnet-dmoe-2b"
N_EXPERTS = 8

tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    dtype="auto",  # 跟随 checkpoint (FP16), 对齐训练方 HF 参考口径
    attn_implementation="sdpa",
).eval()
model.to("cuda")  # ZeroGPU 拦截: 主进程零显存, 调用时权重流式进 VRAM


def _as_text(content):
    """Gradio 6 Chatbot content 归一化: str 或 [{type:'text',text:...}] parts 列表 → str."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(p.get("text", "") for p in content
                       if isinstance(p, dict) and p.get("type") == "text")
    return "" if content is None else str(content)


def _estimate_duration(*args, **kwargs):
    """估算 GPU 时长. 不依赖参数位次 (Gradio 可能把 progress 位置传入):
    取位置参数中第一个落在 [16,512] 的数值 (即 max_new_tokens).
    ZeroGPU 会对估算值加安全余量再校验上限, 故保守封顶 120s."""
    mnt = 128
    for a in args:
        if isinstance(a, (int, float)) and not isinstance(a, bool) and 16 <= a <= 512:
            mnt = int(a)
            break
    return min(120, 30 + int(mnt * 0.5))


@spaces.GPU(duration=_estimate_duration)
def respond(history, max_new_tokens, temperature, top_p, repetition_penalty, seed):
    """流式生成一轮回答; 结束后输出 Ponder / MoE 诊断统计."""
    history = history or []
    if not history or history[-1].get("role") != "user":
        yield history, {"status": "无用户输入"}
        return

    if seed and int(seed) > 0:
        torch.manual_seed(int(seed))

    messages = [{"role": m.get("role"), "content": _as_text(m.get("content"))}
                for m in history if m.get("role") in ("user", "assistant")]
    enc = tok.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=True,
        return_tensors="pt", return_dict=True,
    ).to("cuda")

    streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
    greedy = float(temperature) <= 0.05
    gen_kwargs = dict(
        **enc,
        max_new_tokens=int(max_new_tokens),
        streamer=streamer,
        do_sample=not greedy,
        repetition_penalty=float(repetition_penalty),
        pad_token_id=tok.pad_token_id or 1,
    )
    if not greedy:
        gen_kwargs.update(temperature=float(temperature), top_p=float(top_p))

    log = getattr(model, "_ponder_log", None)   # 注意: 空列表是 falsy, 不能写 `or []`
    if log is None:
        log = getattr(model, "ponder_log", None)
    n0 = len(log) if log is not None else 0
    t0 = time.perf_counter()

    thr = Thread(target=model.generate, kwargs=gen_kwargs)
    thr.start()
    partial = ""
    yield history, {"status": "思考中…"}
    for piece in streamer:
        partial += piece
        yield history + [{"role": "assistant", "content": partial}], {"status": "生成中…"}
    thr.join()
    dt = time.perf_counter() - t0

    entries = log[n0:] if log is not None else []
    pre = [e for e in entries if e.get("mode") == "prefill"]
    dec = [e for e in entries if e.get("mode") == "decode"]
    n_dec = len(dec)
    act = [e["moe_active"] for e in dec
           if e.get("moe_active") is not None] or \
          [e["moe_active"] for e in pre if e.get("moe_active") is not None]
    stats = {
        "生成": {
            "输出 tokens": n_dec,
            "耗时 (秒)": round(dt, 2),
            "速度 (tok/s)": round(n_dec / dt, 1) if dt > 0 else None,
            "采样": ("greedy" if greedy else f"T={temperature}, top_p={top_p}")
                    + f", rep={repetition_penalty}",
        },
        f"Ponder 思考 (max_steps={model.config.max_ponder_steps})": {
            "prefill executed": pre[-1]["executed"] if pre else None,
            "prefill steps_mean": pre[-1].get("steps_mean") if pre else None,
            "decode steps_mean (平均)": round(
                sum(e["steps_mean"] for e in dec) / n_dec, 3) if n_dec else None,
            "decode steps_max (峰值)": max(
                (e.get("steps_max") or 0) for e in dec) if n_dec else None,
        },
        f"MoE 专家 (每层 {N_EXPERTS} 个)": {
            "active experts (平均)": round(
                sum(act) / len(act), 3) if act else None,
        },
    }
    yield history + [{"role": "assistant", "content": partial or "(空输出)"}], stats


HEADER = f"""
# 🧠 SamAI Pnet-dMoE 2B · Ponder-MoE 对话 Demo

[{MODEL_ID}](https://huggingface.co/{MODEL_ID}) · HF 参考实现 · trust_remote_code 单文件建模

**Ponder 思考块**：后 8 层 (L34-41) 循环执行，累计停机质量 1-ε 即停，hard 样本自动"多想"
｜ **动态 MoE**：8 专家阈值路由 (p_i ≥ τ)，每 token 仅首轮路由一次
｜ **MTP**：训练期辅助头，推理自动旁路

> ⚠️ **能力边界（提问前请先读）**：本模型仅用 **1,200 条中文数学 CoT**（MetaMathQA-GSM8K）训练。
> **数学/推理问答**（点下方示例）效果最佳；**闲聊、日常对话、英文提问属训练分布外**，
> 输出会退化为无意义复读（这是数据缺口，非部署故障——HF 参考实现在同类输入上同样退化，
> 修复需训练侧追加对话 SFT 轮）。默认已开启贪心解码 + 防重复惩罚以缓解观感。

生成结束后，下方 JSON 面板给出 **Ponder 步数 / MoE 激活专家数** 诊断（来自 `model._ponder_log`）。

<sub>v1.1 · 防重复缓解 + 能力声明</sub>
"""

FOOTER = """
---
💡 调参建议：数学/推理题用低温度或 greedy 并调大 **max_new_tokens**，观察 decode
steps_mean 是否高于简单题（自适应思考）；首字延迟为 ZeroGPU 冷启动预热，属正常现象。
**防重复惩罚 (repetition_penalty)** 为演示层缓解措施，1.0 = 关闭（严格参考口径）；
对分布外输入只能减轻复读观感，无法让模型真正"会聊天"。

 HF 模型卡 · 本 Space 仅作参考实现演示，生产 CPU 引擎 (lal/GPQ8) 另行交付。
"""

with gr.Blocks(title="SamAI Pnet-dMoE 2B Chat Demo") as demo:
    gr.Markdown(HEADER)
    chat = gr.Chatbot(height=430, label="对话")  # Gradio 6: messages 格式为唯一格式
    with gr.Row():
        msg = gr.Textbox(placeholder="输入消息，Enter 发送…", show_label=False,
                         scale=9, container=False, autofocus=True)
        send = gr.Button("发送 ▶", scale=1, variant="primary")
    with gr.Accordion("生成参数", open=False):
        max_tokens = gr.Slider(16, 512, value=256, step=16, label="max_new_tokens")
        temperature = gr.Slider(0.0, 1.5, value=0.0, step=0.05,
                                label="temperature (≤0.05 = greedy，推荐)")
        top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="top_p")
        rep_pen = gr.Slider(1.0, 1.5, value=1.15, step=0.05,
                            label="repetition_penalty (1.0 = 关闭)")
        seed = gr.Number(value=0, precision=0, label="seed (0 = 随机)")
    stats = gr.JSON(label="Ponder / MoE 诊断 (mode · executed · steps_mean · moe_active)")
    gr.Examples(
        examples=[
            ["一个长方形长 12 米宽 8 米，面积是多少？请简要推理。"],
            ["斐波那契数列 F(1)=1, F(2)=1, F(n)=F(n-1)+F(n-2)。求 F(20) 是多少？"],
            ["用 Python 写一个函数，判断一个字符串是否是回文，并给出示例。"],
            ["小明有 24 个苹果，平均分给 6 个朋友，每人分到几个？请简要推理。"],
        ],
        inputs=[msg],
    )
    with gr.Row():
        clear = gr.Button("🗑 清空对话")

    def add_user(history, m):
        m = (m or "").strip()
        if not m:
            return (history or []), ""
        return (history or []) + [{"role": "user", "content": m}], ""

    common = dict(fn=respond,
                  inputs=[chat, max_tokens, temperature, top_p, rep_pen, seed],
                  outputs=[chat, stats])
    msg.submit(add_user, [chat, msg], [chat, msg]).then(
        api_name="respond", **common)
    send.click(add_user, [chat, msg], [chat, msg]).then(**common)
    clear.click(lambda: ([], None), None, [chat, stats])

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Soft())