# -*- coding: utf-8 -*-
"""
把三个智能体的回放视频拼成一支"总回放"视频（方便上传 Hub / 交作业）。
==============================================================
输入：videos/replay_frozenlake4x4.mp4、replay_frozenlake8x8.mp4、replay_taxi.mp4
输出：videos/replay_value_based_all.mp4（1280x720 / 30fps，开头有总览卡，结尾有分数汇总卡）

用法：
  D:\\study\\robot\\pyenv\\DRL\\python.exe make_combined_video.py
"""

import json
import os

import imageio.v2 as imageio
import numpy as np

import evaluate_and_record as ev   # 复用录像脚本里的画布/字体/配色

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
VIDEO_DIR = os.path.join(BASE_DIR, "videos")
EVAL_DIR = os.path.join(BASE_DIR, "eval_results")

FPS = ev.FPS
RUNS = [
    ("frozenlake4x4", "FrozenLake-v1 · 4x4 不滑", "PPT 第一部分（10,000 回合）"),
    ("frozenlake8x8", "FrozenLake-v1 · 8x8", "进阶：地图变大 + 探索调度"),
    ("taxi", "Taxi-v3", "PPT 第二部分（25,000 回合）"),
]


def load_scores():
    scores = {}
    for run, _, _ in RUNS:
        path = os.path.join(EVAL_DIR, f"eval_{run}.json")
        if os.path.exists(path):
            with open(path, encoding="utf-8") as f:
                scores[run] = json.load(f)
    return scores


def intro_card(scores):
    canvas, draw = ev.base_canvas()
    ev.draw_text(draw, (ev.WIDTH // 2, 110), "value-based 实践 · Q-Learning 智能体回放合集",
                 ev.FONT_TITLE, fill=ev.ACCENT, anchor="mm")
    ev.draw_text(draw, (ev.WIDTH // 2, 165),
                 "从零手写 Q 表 + epsilon-greedy + Bellman 更新（PPT 第 27–40 页）",
                 ev.FONT_BODY, fill=ev.TEXT, anchor="mm")
    draw.line([120, 205, ev.WIDTH - 120, 205], fill=ev.PANEL_EDGE, width=2)

    y = 265
    for i, (run, name, note) in enumerate(RUNS):
        s = scores.get(run)
        ev.draw_text(draw, (110, y), f"{i + 1}. {name}", ev.FONT_H2, fill=ev.TEXT)
        ev.draw_text(draw, (110, y + 34), note, ev.FONT_SMALL, fill=ev.MUTED)
        if s:
            score_text = f"{s['mean_reward']:.2f} ± {s['std_reward']:.2f}"
            ev.draw_text(draw, (ev.WIDTH - 110, y + 16), score_text, ev.FONT_H1,
                         fill=ev.GREEN if s["mean_reward"] > 0 else ev.RED, anchor="rm")
            ev.draw_text(draw, (ev.WIDTH - 110, y + 46),
                         f"成功率 {s['success_rate'] * 100:.0f}%   "
                         f"平均步数 {s['mean_steps']:.1f}",
                         ev.FONT_SMALL, fill=ev.MUTED, anchor="rm")
        y += 95

    ev.draw_text(draw, (ev.WIDTH // 2, 620),
                 "每段回放都标注了步数 / 动作 / Q 值条形图，片尾有该智能体的评估分数",
                 ev.FONT_SMALL, fill=ev.MUTED, anchor="mm")
    return canvas


def end_card(scores):
    canvas, draw = ev.base_canvas()
    ev.draw_text(draw, (ev.WIDTH // 2, 105), "三个智能体评估分数汇总", ev.FONT_TITLE,
                 fill=ev.ACCENT, anchor="mm")
    ev.draw_text(draw, (ev.WIDTH // 2, 155),
                 "独立环境 · 贪婪策略 argmax_a Q(s,a) · 100 局", ev.FONT_BODY,
                 fill=ev.MUTED, anchor="mm")
    draw.line([120, 195, ev.WIDTH - 120, 195], fill=ev.PANEL_EDGE, width=2)

    y = 265
    for run, name, _ in RUNS:
        s = scores.get(run)
        if not s:
            continue
        ev.draw_text(draw, (120, y), name, ev.FONT_H2, fill=ev.TEXT)
        ev.draw_text(draw, (560, y), f"回报 {s['mean_reward']:.2f} ± {s['std_reward']:.2f}",
                     ev.FONT_H2, fill=ev.GREEN if s["mean_reward"] > 0 else ev.RED)
        ev.draw_text(draw, (880, y), f"成功率 {s['success_rate'] * 100:.0f}%",
                     ev.FONT_H2, fill=ev.ACCENT)
        y += 90

    ev.draw_text(draw, (ev.WIDTH // 2, 610),
                 "模型文件（Q 表 pkl）+ 训练代码 + 日志 已打包上传 Hub",
                 ev.FONT_BODY, fill=ev.TEXT, anchor="mm")
    return canvas


def main():
    scores = load_scores()
    clips = []
    for run, _, _ in RUNS:
        path = os.path.join(VIDEO_DIR, f"replay_{run}.mp4")
        if not os.path.exists(path):
            print(f"[跳过] 找不到 {path}，请先跑 evaluate_and_record.py --model models/qtable_{run}.pkl")
            continue
        clips.append(path)
    if not clips:
        raise SystemExit("没有可拼接的视频。")

    out_path = os.path.join(VIDEO_DIR, "replay_value_based_all.mp4")
    writer = imageio.get_writer(out_path, fps=FPS, codec="libx264", quality=8,
                                macro_block_size=2)
    frames = 0

    def push(img, repeat=1):
        nonlocal frames
        arr = np.asarray(img, dtype=np.uint8)
        for _ in range(repeat):
            writer.append_data(arr)
            frames += 1

    push(intro_card(scores), int(3.0 * FPS))
    for path in clips:
        reader = imageio.get_reader(path)
        n = 0
        for frame in reader:
            writer.append_data(frame)
            frames += 1
            n += 1
        reader.close()
        print(f"  [拼接] {os.path.basename(path)}：{n} 帧")
    push(end_card(scores), int(5.0 * FPS))

    writer.close()
    size_mb = os.path.getsize(out_path) / (1024 * 1024)
    print(f"[输出] 总回放视频 = videos/{os.path.basename(out_path)} "
          f"({frames} 帧 / {frames / FPS:.1f} 秒 / {size_mb:.1f} MB)")


if __name__ == "__main__":
    main()