"""
harness.py —— 一个最小的"实验框架"：配置、日志、检查点、断点续训
=====================================================================

【这个文件要解决的问题】
    在云上做实验，最常见的事故不是代码写错，而是：
      - 竞价（spot）实例被回收、机器重启、SSH 断开，训练跑到一半没了；
      - 跑了十几个实验，分不清哪个结果对应哪组超参数；
      - "断点续训"之后的结果和一口气跑完的不一样，却不知道为什么。
    这里实现几样最基本的工具：
      1. config_hash：配置的指纹。同一组超参数永远得到同一个运行目录名，不同的不会互相覆盖。
      2. JsonlLogger：每一步的指标写成一行 JSON，追加写入、立即落盘；训练中断也不会丢已有记录。
      3. save_checkpoint / load_checkpoint：保存**完整的训练状态**——不只是模型参数，
         还有优化器状态（Adam 的 m、v）、学习率调度器的进度、步数、以及数据抽样的随机数状态。
         只有全部恢复，续训才和一口气跑完逐位相同。
      4. 原子写入：先写到临时文件，再用 os.replace 一步改名。即使正好在写检查点时被杀掉，
         磁盘上也只会有"旧的完整检查点"或"新的完整检查点"，不会留下写了一半的坏文件。

【LLM 联系】
    大模型训练动辄几周、几千块卡，硬件故障是常态（Llama 3 的技术报告专门统计了训练中断的次数和原因）。
    "频繁保存检查点 + 能精确续训"是大规模训练的基本功；你在云上做课程项目时也一样需要它。
"""

import hashlib
import json
import os

import torch


def config_hash(cfg):
    """
    配置的指纹：把配置字典按键排序后转成 JSON，再取 SHA-1 的前 10 位。

    【参数】cfg：只含基本类型（数字、字符串、列表、字典）的字典。
    【返回】10 个字符的十六进制字符串。键的顺序不同、但内容相同的配置得到同一个指纹。
    """
    return hashlib.sha1(json.dumps(cfg, sort_keys=True).encode()).hexdigest()[:10]


class JsonlLogger:
    """
    把每条记录追加写成一行 JSON（JSON Lines 格式）。

    【为什么用 JSONL】每行独立、可以边写边读；文件被截断时，前面的行仍然完整可用；
            之后用几行代码就能读进 pandas 画曲线。
    """

    def __init__(self, path):
        self.path = path

    def log(self, **record):
        # 【Python 小知识】**record 把所有关键字参数收集成一个字典：log(step=1, loss=0.5)
        with open(self.path, "a", encoding="utf-8") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")
            f.flush()                       # 立即写到磁盘，不留在缓冲区里

    def read(self):
        if not os.path.exists(self.path):
            return []
        with open(self.path, encoding="utf-8") as f:
            return [json.loads(line) for line in f if line.strip()]


def save_checkpoint(path, step, model, opt, sched, data_gen):
    """
    原子地保存完整的训练状态。

    【保存的内容】步数、模型参数、优化器状态、调度器状态、数据抽样生成器的状态、全局 torch 随机数状态。
    【原子写入】先 torch.save 到 path + ".tmp"，再 os.replace 改名（同一个文件系统上改名是原子操作）。
    """
    state = {
        "step": step,
        "model": model.state_dict(),
        "opt": opt.state_dict(),
        "sched": sched.state_dict(),
        "data_gen": data_gen.get_state(),
        "torch_rng": torch.get_rng_state(),
    }
    tmp = path + ".tmp"
    torch.save(state, tmp)
    os.replace(tmp, path)


def latest_checkpoint(run_dir):
    """运行目录里步数最大的检查点文件路径；没有则返回 None。文件名形如 ckpt_000200.pt。"""
    files = sorted(f for f in os.listdir(run_dir) if f.startswith("ckpt_") and f.endswith(".pt"))
    return os.path.join(run_dir, files[-1]) if files else None


def load_checkpoint(path, model, opt, sched, data_gen, restore=("opt", "sched", "rng")):
    """
    从检查点恢复训练状态。

    【参数】restore：要恢复哪些部分（模型参数总是恢复）。故意少恢复某一项，可以观察它对结果的影响。
    【返回】检查点里记录的步数（续训从下一步开始）。
    """
    state = torch.load(path, weights_only=False)
    model.load_state_dict(state["model"])
    if "opt" in restore:
        opt.load_state_dict(state["opt"])
    if "sched" in restore:
        sched.load_state_dict(state["sched"])
    if "rng" in restore:
        data_gen.set_state(state["data_gen"])
        torch.set_rng_state(state["torch_rng"])
    return state["step"]


def budget_check(gpu_hours, price_per_hour, budget):
    """
    开跑之前的预算检查。

    【参数】gpu_hours：估计的 GPU 小时；price_per_hour：你查到的单价；budget：愿意花的上限。
    【返回】(估计花费, 是否在预算内)。建议再留 30%～50% 的余量给调试和重跑。
    """
    cost = gpu_hours * price_per_hour
    return cost, cost <= budget