"""真实字符 GRU 后训练；显式 --run --out-dir 才训练，CPU 可运行。"""

from argparse import ArgumentParser
from copy import deepcopy
import json
from pathlib import Path
import platform
import random
import time

import torch
from torch import nn
from torch.nn import functional as F

import post_training_frontier as lab


class Adder(nn.Module):
    def __init__(self, plan):
        super().__init__()
        self.vocab = plan["vocab"]
        self.eos = self.vocab.index(plan["eos"])
        self.pad = self.vocab.index(plan["pad"])
        self.embed = nn.Embedding(len(self.vocab), plan["embedding"])
        self.rnn = nn.GRU(plan["embedding"], plan["hidden"], batch_first=True)
        self.head = nn.Linear(plan["hidden"], len(self.vocab))
        allowed = [c in "0123456789~" for c in self.vocab]
        self.register_buffer("allowed", torch.tensor(allowed))

    def forward(self, tokens):
        hidden, _ = self.rnn(self.embed(tokens))
        return self.head(hidden).masked_fill(~self.allowed, -1e9)

    def prompts(self, pairs):
        return torch.tensor([[self.vocab.index(c) for c in f"{a}+{b}="]
                             for a, b in pairs])


def supervised(model, plan, seed):
    rng = random.Random(seed + 1000)
    opt = torch.optim.Adam(model.parameters(), lr=plan["sft_lr"])
    losses = []
    for _ in range(plan["sft_steps"]):
        pairs = rng.choices(plan["split"]["train"], k=plan["sft_batch"])
        targets = [[model.vocab.index(c) for c in str(a+b)+"~"] for a, b in pairs]
        width = max(map(len, targets))
        target = torch.tensor([t+[model.pad]*(width-len(t)) for t in targets])
        inputs = torch.cat([model.prompts(pairs), target[:, :-1]], dim=1)
        logits = model(inputs)[:, 3:]
        loss = F.cross_entropy(logits.reshape(-1, len(model.vocab)), target.flatten(),
                               ignore_index=model.pad)
        opt.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), plan["gradient_clip"])
        opt.step()
        losses.append(loss.item())
    return losses


@torch.no_grad()
def generate(model, pairs, plan, generator, greedy=False):
    prefix = model.prompts(pairs)
    active = torch.ones(len(pairs), dtype=torch.bool)
    tokens, masks = [], []
    for _ in range(plan["max_new_tokens"]):
        logits = model(prefix)[:, -1] / plan["temperature"]
        if greedy:
            nxt = logits.argmax(-1)
        else:
            nxt = torch.multinomial(logits.softmax(-1), 1, generator=generator).squeeze(1)
        nxt = torch.where(active, nxt, model.pad)
        tokens.append(nxt)
        masks.append(active.clone())
        active &= nxt != model.eos
        prefix = torch.cat([prefix, nxt[:, None]], dim=1)
        if not active.any():
            break
    return torch.stack(tokens, dim=1), torch.stack(masks, dim=1)


def to_rows(model, pairs, tokens, mask):
    rows = []
    for pair, seq, valid in zip(pairs, tokens.tolist(), mask.tolist()):
        ids = [tok for tok, use in zip(seq, valid) if use]
        terminated = bool(ids and ids[-1] == model.eos)
        visible = ids[:-1] if terminated else ids
        rows.append([*pair, "".join(model.vocab[t] for t in visible), terminated])
    return rows


def distribution(model, pairs, tokens):
    inputs = torch.cat([model.prompts(pairs), tokens[:, :-1]], dim=1)
    return model(inputs)[:, 3:].log_softmax(-1)


def reinforce(initial, reference, plan, seed, arm):
    model = deepcopy(initial)
    rng = random.Random(seed+2000)
    generator = torch.Generator().manual_seed(seed+3000)
    opt = torch.optim.Adam(model.parameters(), lr=plan["rl_lr"])
    history = []
    for step in range(plan["rl_steps"]):
        prompts = rng.choices(plan["split"]["train"], k=plan["prompts_per_step"])
        pairs = [pair for pair in prompts for _ in range(plan["group_size"])]
        tokens, mask = generate(model, pairs, plan, generator)
        rows = to_rows(model, pairs, tokens, mask)
        rewards = torch.tensor([lab.reward(r[:2], r[2], r[3], arm) for r in rows],
                               dtype=torch.float32).reshape(-1, plan["group_size"])
        std = rewards.std(-1, unbiased=False, keepdim=True)
        advantage = ((rewards-rewards.mean(-1, keepdim=True))/(std+1e-8)).flatten()
        with torch.no_grad():
            old_logp = distribution(model, pairs, tokens).gather(-1, tokens[..., None]).squeeze(-1)
            ref_logdist = distribution(reference, pairs, tokens)
        updates = []
        for _ in range(plan["updates_per_rollout"]):
            logdist = distribution(model, pairs, tokens)
            logp = logdist.gather(-1, tokens[..., None]).squeeze(-1)
            ratio = (logp-old_logp).exp()
            gain = torch.minimum(ratio*advantage[:, None],
                                 ratio.clamp(1-plan["clip"], 1+plan["clip"])*advantage[:, None])
            kl = (logdist.exp()*(logdist-ref_logdist)).sum(-1)
            loss = ((-gain+plan["kl_beta"]*kl)*mask).sum()/mask.sum()
            entropy = -(logdist.exp()*logdist).sum(-1)
            opt.zero_grad()
            loss.backward()
            grad = nn.utils.clip_grad_norm_(model.parameters(), plan["gradient_clip"])
            opt.step()
            updates.append({"loss": loss.item(), "gradient_norm_before_clip": grad.item(),
                            "kl": (kl*mask).sum().item()/mask.sum().item(),
                            "entropy": (entropy*mask).sum().item()/mask.sum().item(),
                            "ratio_outside_clip": (((ratio-1).abs()>plan["clip"])*mask).sum().item()/mask.sum().item()})
        history.append({"step": step+1, "rows": rows,
                        "zero_variance_groups": int((std == 0).sum()), "updates": updates})
        if (step+1) % 80 == 0:
            metrics = lab.metrics(rows)
            print(f"seed={seed} arm={arm} step={step+1} "
                  f"strict={metrics['strict_rate']:.3f} weak={metrics['substring_rate']:.3f}", flush=True)
    return model, history


@torch.no_grad()
def evaluate_model(model, plan, seed):
    result = {}
    for split_index, (split, pairs) in enumerate(plan["split"].items()):
        generator = torch.Generator().manual_seed(seed+4000+split_index)
        greedy, mask = generate(model, pairs, plan, generator, greedy=True)
        sampled_pairs = [pair for pair in pairs for _ in range(plan["eval_samples"])]
        sampled, smask = generate(model, sampled_pairs, plan, generator)
        result[split] = {"greedy": to_rows(model, pairs, greedy, mask),
                         "sampled": to_rows(model, sampled_pairs, sampled, smask)}
    return result


def run(plan_path, out):
    out.mkdir(parents=True, exist_ok=False)
    plan = json.loads(plan_path.read_text())
    assert plan == lab.make_plan(), "当前实现只支持已固定的教学计划"
    torch.set_num_threads(plan["threads"])
    torch.use_deterministic_algorithms(True)
    result = {"plan_sha256": lab.digest(plan_path),
              "source_sha256": {p.name: lab.digest(p) for p in [Path(__file__), Path(lab.__file__)]},
              "python": platform.python_version(), "torch": torch.__version__,
              "platform": platform.platform(), "device": "cpu", "dtype": "float32", "runs": []}
    start = time.monotonic()
    for seed in plan["seeds"]:
        torch.manual_seed(seed)
        initial = Adder(plan)
        result["parameters"] = sum(p.numel() for p in initial.parameters())
        losses = supervised(initial, plan, seed)
        reference = deepcopy(initial).requires_grad_(False)
        row = {"seed": seed, "sft_loss": losses, "evaluations": {}, "training": {}, "checkpoints": {}}
        row["evaluations"]["initial"] = evaluate_model(initial, plan, seed)
        models = {"initial": initial}
        for arm in plan["arms"]:
            model, history = reinforce(initial, reference, plan, seed, arm)
            row["training"][arm] = history
            row["evaluations"][arm] = evaluate_model(model, plan, seed)
            models[arm] = model
        for stage, model in models.items():
            path = out / f"seed-{seed}-{stage}.pt"
            torch.save({"plan": plan, "state_dict": model.state_dict()}, path)
            restored = Adder(plan)
            restored.load_state_dict(torch.load(path, weights_only=True)["state_dict"])
            assert evaluate_model(restored, plan, seed) == row["evaluations"][stage]
            row["checkpoints"][stage] = {"file": path.name, "sha256": lab.digest(path), "reload_equal": True}
        result["runs"].append(row)
        (out/"record.json").write_text(json.dumps(result, ensure_ascii=False)+"\n")
    result["wall_seconds"] = time.monotonic()-start
    (out/"record.json").write_text(json.dumps(result, ensure_ascii=False)+"\n")
    (out/"summary.json").write_text(json.dumps(lab.summarize(result, plan), indent=2)+"\n")
    print(f"完成：{out}，{result['wall_seconds']:.1f} 秒。", flush=True)


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--plan", type=Path, default=lab.PLAN)
    parser.add_argument("--out-dir", type=Path)
    args = parser.parse_args()
    if not args.run:
        print("默认不训练。请使用 --run --out-dir runs/post-training-new。")
        return
    if args.out_dir is None:
        parser.error("训练需要 --out-dir，且目录不能已存在")
    run(args.plan, args.out_dir)


if __name__ == "__main__":
    main()