app.py

import os
import orjson
import concurrent.futures
import random
import torch
import threading
import time
import uuid
import glob
import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
from huggingface_hub import snapshot_download, hf_hub_download, HfApi
import secrets

from riichienv import RiichiEnv, GameRule
from model3pLOCAL import load_model

# ==========================================
# 0. 分布式多开与云端持久化配置
# ==========================================
DATA_REPO_ID = "ffzeroHua/mj-eval-results"  # 📊 战绩数据集仓库
MODEL_REPO_ID = "ffzeroHua/Riichi-Model-Repo" # 🧠 模型权重仓库
HF_TOKEN = os.getenv("HF_TOKEN")

# 为当前节点生成唯一的 ID
WORKER_ID = os.getenv("WORKER_ID", str(uuid.uuid4())[:6])
REPORT_FILE_PREFIX = 'Step40800P42998_vs_9070_Fair_eval_report'
REPORT_FILE = f"{REPORT_FILE_PREFIX}_{WORKER_ID}.txt"

api = HfApi()
EVAL_RUNNING = True

# 🚀 设定要从云端拉取并进行对抗的两个模型
TEST_MODEL = "Elite3P_Step40800_P42998.pth"
EXAMINER_MODEL = "Elite4z9070.pth"

def sync_models_from_hub():
    """启动时从指定的模型仓库拉取对战双方的权重文件"""
    if HF_TOKEN and "你的用户名" not in MODEL_REPO_ID:
        print(f"☁️ 正在从模型仓库 [{MODEL_REPO_ID}] 拉取评估模型...")
        try:
            hf_hub_download(repo_id=MODEL_REPO_ID, filename=TEST_MODEL, repo_type="model", local_dir=".", token=HF_TOKEN)
            print(f"✅ 成功拉取测试模型: {TEST_MODEL}")
            
            hf_hub_download(repo_id=MODEL_REPO_ID, filename=EXAMINER_MODEL, repo_type="model", local_dir=".", token=HF_TOKEN)
            print(f"✅ 成功拉取考官模型: {EXAMINER_MODEL}")
            
            print("🎉 模型环境准备完毕！")
        except Exception as e:
            print(f"❌ 拉取模型失败，请检查文件名或仓库权限: {e}")
    else:
        print("⚠️ 未配置有效 HF_TOKEN 或未修改 MODEL_REPO_ID，将尝试使用本地已存在的模型文件。")

def sync_data_from_hub():
    """启动时从数据集下载所有节点的战绩分片文件"""
    if HF_TOKEN and "你的用户名" not in DATA_REPO_ID:
        try:
            print(f"🔄 正在从 Hub 拉取全局历史战绩数据...")
            snapshot_download(
                repo_id=DATA_REPO_ID,
                repo_type="dataset",
                local_dir=".",
                allow_patterns=REPORT_FILE_PREFIX + "_*.txt",
                token=HF_TOKEN
            )
            print("✅ 历史数据拉取完成。")
        except Exception as e:
            print(f"⚠️ 拉取历史战绩失败: {e}")

def sync_data_to_hub():
    """将当前节点的战绩文件备份到数据集"""
    if HF_TOKEN and "你的用户名" not in DATA_REPO_ID:
        try:
            api.upload_file(
                path_or_fileobj=REPORT_FILE,
                path_in_repo=REPORT_FILE,
                repo_id=DATA_REPO_ID,
                repo_type="dataset",
                token=HF_TOKEN
            )
            print(f"☁️ 节点 {WORKER_ID} 战绩已同步至 Hub: {time.strftime('%H:%M:%S')}")
        except Exception as e:
            print(f"❌ 同步失败: {e}")

# ==========================================
# 1. 高频及模型加载逻辑
# ==========================================
def patch_event_fast(event_str):
    if '"kita"' in event_str:
        event_str = event_str.replace('"kita"', '"nukidora"')
    
    if '"start_kyoku"' in event_str or '"deltas"' in event_str:
        event = orjson.loads(event_str)
        if event.get('type') == 'start_kyoku':
            scores = event.setdefault('scores', [])
            while len(scores) < 4: scores.append(0)
            tehais = event.setdefault('tehais', [])
            while len(tehais) < 4: tehais.append(["?" for _ in range(13)])
        if 'deltas' in event:
            deltas = event['deltas']
            while len(deltas) < 4: deltas.append(0)
        return orjson.dumps(event).decode('utf-8')
    return event_str

def patch_resp_fast(resp_str):
    if not resp_str: return resp_str
    return resp_str.replace('"nukidora"', '"kita"')

_MODEL_CACHE = {}

def get_cached_model(player_id: int, model_file: str):
    key = (player_id, model_file)
    if key not in _MODEL_CACHE:
        torch.set_num_threads(1) 
        _MODEL_CACHE[key] = load_model(player_id, model_file)
    return _MODEL_CACHE[key]

class MortalAgent:
    def __init__(self, player_id: int, model_file: str):
        self.player_id = player_id
        self.model = get_cached_model(player_id, model_file)

    def act(self, obs):
        resp = None
        for event in obs.new_events():
            event_patched = patch_event_fast(event)
            resp = patch_resp_fast(self.model.react(event_patched))
        action = obs.select_action_from_mjai(resp)
        assert action is not None, "Mortal must return a legal action"
        return action

# ==========================================
# 2. 核心对局任务 (复式评测：同种子 3 局轮换)
# ==========================================
def play_three_games_with_seed(seed):
    """
    接收一个全局唯一的随机种子，让测试模型在 0, 1, 2 三个位置各打一局。
    返回 3 局的成绩列表。
    """
    results = []
    
    # 轮换测试模型的位置
    for test_seat in range(3):
        env = RiichiEnv(game_mode="3p-red-half", rule=GameRule.default_tenhou())
        agents = {}
        for i in range(3):
            # 测试模型坐在 test_seat，其他位置是考官模型
            model_file = TEST_MODEL if i == test_seat else EXAMINER_MODEL
            agents[i] = MortalAgent(i, model_file)
        
        # 🚀 关键：注入指定的种子！这样三次循环生成的牌山是完全一致的
        obs_dict = env.reset(seed=seed)
        
        while not env.done():
            actions = {pid: agents[pid].act(obs) for pid, obs in obs_dict.items()}
            obs_dict = env.step(actions)
            
        scores = env.scores()
        ranks = env.ranks()
        
        # 记录测试模型在这一局的表现
        results.append((ranks[test_seat], scores[test_seat]))
        
    return results

# ==========================================
# 3. 后台独立评估线程
# ==========================================
def background_eval_loop():
    sync_models_from_hub() # 启动时从 Riichi-Model-Repo 拉取对战模型
    sync_data_from_hub()   # 启动时从战绩仓库拉取历史战绩
    
    NUM_WORKERS = 1 
    print(f"🚀 节点 [{WORKER_ID}] 后台对战线程已启动: {TEST_MODEL} 挑战双 {EXAMINER_MODEL} (复式打法)")
    
    if not os.path.exists(REPORT_FILE):
        open(REPORT_FILE, 'w').close()

    games_since_last_sync = 0

    with concurrent.futures.ProcessPoolExecutor(max_workers=NUM_WORKERS) as executor:
        # 🚀 使用 secrets.randbits(31) 生成跨节点绝对不冲突的 31 位正整数种子
        futures = {executor.submit(play_three_games_with_seed, secrets.randbits(31)) for _ in range(NUM_WORKERS * 2)}
        games_completed = 0
        
        while EVAL_RUNNING and futures:
            done, futures = concurrent.futures.wait(
                futures, return_when=concurrent.futures.FIRST_COMPLETED
            )
            
            with open(REPORT_FILE, "a") as f:
                for future in done:
                    try:
                        # 现在 future.result() 返回的是包含 3 局结果的列表
                        batch_results = future.result()
                        for rank, score in batch_results:
                            f.write(f"{rank} {score}\n")
                            games_completed += 1
                            games_since_last_sync += 1
                            print(f"[节点 {WORKER_ID}] 完成复式对局: 顺位 {rank}, 得点 {score}")
                        f.flush()
                    except Exception as e:
                        print(f"对局异常: {e}")
                    
                    if EVAL_RUNNING:
                        # 完成了一轮复式，再投递一个全新的安全种子
                        futures.add(executor.submit(play_three_games_with_seed, secrets.randbits(31)))
            
            # 由于一次产出 3 局，同步频率可以保持在累计完成 50 局左右时触发
            if games_since_last_sync >= 60:
                sync_data_to_hub()
                sync_data_from_hub()
                games_since_last_sync = 0

# ==========================================
# 4. 前端 Gradio 实时展示面板 (全局汇总)
# ==========================================
def read_and_analyze():
    all_files = glob.glob(f"{REPORT_FILE_PREFIX}_*.txt")
    if not all_files:
        return f"⏳ 正在拉取模型 `{TEST_MODEL}` 和 `{EXAMINER_MODEL}`，等待第一局完成...", None
    
    ranks, scores = [], []
    try:
        for file in all_files:
            with open(file, "r") as f:
                lines = f.readlines()
                for line in lines:
                    parts = line.strip().split()
                    if len(parts) == 2:
                        ranks.append(int(float(parts[0])))
                        scores.append(float(parts[1]))
        total = len(ranks)
        if total == 0:
            return f"⏳ 模型已就绪，正在进行第一局对抗...", None

        avg_rank = sum(ranks) / total
        avg_score = sum(scores) / total
        rank1_rate = ranks.count(1) / total * 100
        rank2_rate = ranks.count(2) / total * 100
        rank3_rate = ranks.count(3) / total * 100

        last_update = time.strftime('%Y-%m-%d %H:%M:%S')

        md_text = f"""
        ### 📊 对战简报
        - ⚔️ **对抗阵容:** 1只 `{TEST_MODEL}` **VS** 2只 `{EXAMINER_MODEL}`
        - 🧮 **总对局数:** {total} 局 (跨节点全局汇集)
        - 🏆 **平均顺位:** {avg_rank:.3f}
        - 💰 **平均得点:** {avg_score:.0f}
        ---
        - 🥇 **一位率:** {rank1_rate:.1f}%
        - 🥈 **二位率:** {rank2_rate:.1f}%
        - 🥉 **三位率:** {rank3_rate:.1f}%
        ---
        - 🌐 **当前节点 ID:** `{WORKER_ID}`
        - 🕒 **刷新时间:** {last_update}
        """

        fig = plt.figure(figsize=(10, 4))
        
        ax1 = fig.add_subplot(121)
        ax1.bar(['1st', '2nd', '3rd'], [rank1_rate, rank2_rate, rank3_rate], color=['#FFD700', '#C0C0C0', '#CD7F32'])
        ax1.set_title(f'Rank Distribution for {TEST_MODEL}')
        ax1.set_ylim(0, max(100, max([rank1_rate, rank2_rate, rank3_rate] + [0]) + 10))
        for i, v in enumerate([rank1_rate, rank2_rate, rank3_rate]):
            ax1.text(i, v + 2, f"{v:.1f}%", ha='center')

        ax2 = fig.add_subplot(122)
        df = pd.DataFrame({'score': scores})
        df['ma'] = df['score'].rolling(window=min(10, max(1, len(df))), min_periods=1).mean()
        ax2.plot(df['score'], alpha=0.3, color='gray', label='Raw Score')
        ax2.plot(df['ma'], color='crimson', linewidth=2, label='Moving Avg (10)')
        ax2.set_title('Score Trend')
        ax2.legend()

        plt.tight_layout()
        return md_text, fig
        
    except Exception as e:
        return f"❌ 数据解析出错: {e}", None

# ==========================================
# 5. 启动 Gradio 应用
# ==========================================
with gr.Blocks() as demo:
    gr.Markdown("# 🀄 Mahjong AI 基准评估舱")
    gr.Markdown(f"当前正在评估: **{TEST_MODEL}** 单挑两名 **{EXAMINER_MODEL}**。启动时会自动从 `Riichi-Model-Repo` 拉取权重。")
    
    with gr.Row():
        with gr.Column(scale=1):
            stats_output = gr.Markdown("🚀 正在初始化基准环境并连接模型仓库...")
            refresh_btn = gr.Button("🔄 手动刷新全局战绩")
        with gr.Column(scale=2):
            plot_output = gr.Plot()

    demo.load(fn=read_and_analyze, inputs=None, outputs=[stats_output, plot_output])
    
    timer = gr.Timer(15)
    timer.tick(fn=read_and_analyze, inputs=None, outputs=[stats_output, plot_output])
    
    refresh_btn.click(fn=read_and_analyze, inputs=None, outputs=[stats_output, plot_output])

if __name__ == "__main__":
    t = threading.Thread(target=background_eval_loop, daemon=True)
    t.start()
    
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft())