"""持续学习：训练记录汇总、线性关联编辑与先预测后更新的数值实验。"""

from argparse import ArgumentParser
from hashlib import sha256
import json
from pathlib import Path
import statistics

DATA = Path(__file__).with_name("data")
PLAN = DATA/"continual-learning-plan.json"
RECORD = DATA/"continual-learning-training.json"
SMALL = DATA/"continual-learning-small.json"
ARMS = ["sequential", "replay50", "replay500", "ewc100", "ewc1000"]
LABELS = {"sequential": "仅学新类", "replay50": "回放 50 张", "replay500": "回放全部 500 张",
          "ewc100": "EWC λ=100", "ewc1000": "EWC λ=1000"}


def digest(path):
    return sha256(Path(path).read_bytes()).hexdigest()


def save_new(path, value):
    with Path(path).open("x") as f:
        json.dump(value, f, ensure_ascii=False, separators=(",", ":"))
        f.write("\n")


def matvec(matrix, vector):
    return [sum(a*b for a, b in zip(row, vector)) for row in matrix]


def edit(matrix, key, desired):
    """满足 W_new @ key = desired 的最小 Frobenius 范数单次改动。"""
    norm2 = sum(x*x for x in key)
    if norm2 == 0:
        raise ValueError("编辑键不能为零向量")
    residual = [v-old for v, old in zip(desired, matvec(matrix, key))]
    return [[weight+error*k/norm2 for weight, k in zip(row, key)]
            for row, error in zip(matrix, residual)]


def editing_experiment():
    original = [[1., 0.], [0., 1.]]
    first = edit(original, [1., 1.], [2., 0.])
    second = edit(first, [1., 0.], [0., 1.])
    keys = {"edited": [1., 1.], "orthogonal": [1., -1.],
            "overlapping": [1., 0.], "nearby": [1., .8]}
    return dict(keys=keys, matrices=[original, first, second],
                values=[{name: matvec(matrix, key) for name, key in keys.items()}
                        for matrix in [original, first, second]])


def online_experiment():
    noise = [-1., 1., 0., 1., -1.]*8
    observations = [(10. if t < 20 else 14.)+n for t, n in enumerate(noise)]
    result = {}
    for name, alpha in [("frozen", 0.), ("slow", .1), ("fast", .5)]:
        estimate, rows = 10., []
        for t, observed in enumerate(observations):
            prediction = estimate  # 先记预测；此时不使用 observed。
            error = (prediction-observed)**2
            estimate += alpha*(observed-estimate)  # 得到反馈后，才更新下一步的状态。
            rows.append(dict(t=t, prediction=prediction, observed=observed,
                             squared_error=error, next_estimate=estimate))
        result[name] = dict(alpha=alpha, rows=rows,
                            before_mse=statistics.mean(r["squared_error"] for r in rows[:20]),
                            after_mse=statistics.mean(r["squared_error"] for r in rows[20:]))
    return result


def small_experiments():
    return dict(source_sha256=digest(__file__), editing=editing_experiment(), online=online_experiment())


def summary(record):
    result = {}
    for arm in ARMS:
        rows = [r for r in record["runs"] if r["arm"] == arm]
        result[arm] = {key: statistics.mean(r["evaluations"]["test"]["metrics"][key] for r in rows)
                       for key in ["old_accuracy", "new_accuracy", "accuracy"]}
        result[arm]["forgetting"] = statistics.mean(
            r["old_test_before"]-r["evaluations"]["test"]["metrics"]["old_accuracy"] for r in rows)
    return result


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--small", type=Path)
    args = parser.parse_args()
    if args.small:
        save_new(args.small, small_experiments())
        print(f"已保存线性编辑与 40 步在线预测记录：{args.small}")
        return
    record = json.loads(RECORD.read_text())
    assert record["complete"] and record["plan_sha256"] == digest(PLAN)
    print("第二阶段结束：旧类 / 新类 / 总体测试正确率")
    for arm, row in summary(record).items():
        print(f"{LABELS[arm]}：{100*row['old_accuracy']:.1f}% / "
              f"{100*row['new_accuracy']:.1f}% / {100*row['accuracy']:.1f}%")
    small = small_experiments()
    print("第一次编辑后的目标值：", small["editing"]["values"][1]["edited"])
    print("第二次编辑后，旧目标变成：", small["editing"]["values"][2]["edited"])
    print("在线预测：变化前 / 变化后 MSE")
    for name, row in small["online"].items():
        print(f"{name}：{row['before_mse']:.3f} / {row['after_mse']:.3f}")


if __name__ == "__main__":
    main()

"""真实类别增量学习：冻结方案后训练，保留全部种子、支路和错误。"""

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

import numpy as np
import sklearn
from sklearn.datasets import load_digits
import torch
from torch import nn
from torch.nn import functional as F

import continual_learning as lab


def data():
    x, y = load_digits(return_X_y=True)
    raw = json.dumps({"pixels": x.astype(int).tolist(), "labels": y.tolist()},
                     separators=(",", ":")).encode()
    return torch.tensor(x/16, dtype=torch.float32), torch.tensor(y), lab.sha256(raw).hexdigest()


def make_plan():
    _, y, digest = data()
    rng = random.Random(20261020)
    groups = {f"{split}_{age}": [] for split in ["train", "dev", "test"] for age in ["old", "new"]}
    memory = []
    for label in range(10):
        ids = [i for i, value in enumerate(y.tolist()) if value == label]
        rng.shuffle(ids)
        age = "old" if label < 5 else "new"
        groups[f"train_{age}"].extend(ids[:100])
        groups[f"dev_{age}"].extend(ids[100:120])
        groups[f"test_{age}"].extend(ids[120:150])
        if label < 5:
            memory.extend(ids[:10])
    return dict(version=1, date="2026-10-20", dataset_sha256=digest,
                ids={k: sorted(v) for k, v in groups.items()}, memory=sorted(memory),
                seeds=[11, 22, 33], arms=lab.ARMS, hidden=48, lr=.003,
                old_steps=240, new_steps=240, batch=32, replay_old_per_batch=8,
                dev_every=20, fisher="predictive expectation over 10 labels on 50 old inputs",
                ewc_lambda={"ewc100": 100., "ewc1000": 1000.},
                selection="all arms, all seeds, last checkpoint; no dev/test parameter search",
                evaluation="one shared 10-class head, no task label or output mask")


def model(plan):
    return nn.Sequential(nn.Linear(64, plan["hidden"]), nn.Tanh(), nn.Linear(plan["hidden"], 10))


def metrics(logits, labels):
    z, y = torch.tensor(logits), torch.tensor(labels)
    correct = z.argmax(-1) == y
    old = y < 5
    return dict(n=len(y), correct=int(correct.sum()), accuracy=correct.float().mean().item(),
                old_n=int(old.sum()), old_correct=int(correct[old].sum()),
                old_accuracy=correct[old].float().mean().item(),
                new_n=int((~old).sum()), new_correct=int(correct[~old].sum()),
                new_accuracy=correct[~old].float().mean().item(),
                nll=F.cross_entropy(z, y).item())


@torch.no_grad()
def evaluate(net, x, y, ids):
    logits, labels = net(x[ids]).tolist(), y[ids].tolist()
    return dict(ids=ids, labels=labels, logits=logits, metrics=metrics(logits, labels))


def fisher_diagonal(net, x):
    """对输入均匀平均，对标签按旧模型分布精确求和；不是批梯度的平方。"""
    params = list(net.parameters())
    importance = [torch.zeros_like(p) for p in params]
    for example in x:
        logp = net(example.unsqueeze(0)).log_softmax(-1)[0]
        probabilities = logp.detach().exp()
        for label in range(len(probabilities)):
            gradients = torch.autograd.grad(logp[label], params, retain_graph=label+1 < len(probabilities))
            for total, gradient in zip(importance, gradients):
                total.add_(probabilities[label]*gradient.square()/len(x))
    return importance


def penalty(net, anchor, importance):
    return sum((f*(p-old).square()).sum() for p, old, f in zip(net.parameters(), anchor, importance))/2


def train(net, x, y, pool, plan, seed, steps, dev_ids, replay=None, importance=None, strength=0.):
    generator = torch.Generator().manual_seed(seed+10000)
    old_generator = torch.Generator().manual_seed(seed+20000)
    anchor = [p.detach().clone() for p in net.parameters()]
    optimizer = torch.optim.Adam(net.parameters(), lr=plan["lr"])  # 每个阶段重新建立优化器。
    history = []
    curve = [dict(step=0, **evaluate(net, x, y, dev_ids)["metrics"])]
    for step in range(steps):
        # 先为各支路抽同样的 32 个新样本；回放支路只取前 24 个。
        positions = torch.randint(len(pool), (plan["batch"],), generator=generator)
        ids = [pool[i] for i in positions.tolist()]
        if replay is not None:
            n_old = plan["replay_old_per_batch"]
            old_positions = torch.randint(len(replay), (n_old,), generator=old_generator)
            ids = ids[:-n_old]+[replay[i] for i in old_positions.tolist()]
        batch = torch.tensor(ids)
        ce = F.cross_entropy(net(x[batch]), y[batch])
        regularizer = penalty(net, anchor, importance) if importance is not None else torch.tensor(0.)
        loss = ce+strength*regularizer
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        history.append(dict(step=step+1, ids=ids, cross_entropy=ce.item(),
                            penalty=regularizer.item(), loss=loss.item()))
        if (step+1) % plan["dev_every"] == 0:
            curve.append(dict(step=step+1, **evaluate(net, x, y, dev_ids)["metrics"]))
    return history, curve


def run(plan_path, out):
    out.mkdir(parents=True, exist_ok=False)
    plan = json.loads(plan_path.read_text())
    assert plan == make_plan()
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    x, y, raw_hash = data()
    assert raw_hash == plan["dataset_sha256"]
    split_ids = {s: sorted(plan["ids"][f"{s}_old"]+plan["ids"][f"{s}_new"]) for s in ["dev", "test"]}
    record = dict(complete=False, plan_sha256=lab.digest(plan_path),
                  source_sha256={p.name: lab.digest(p) for p in [Path(__file__), Path(lab.__file__)]},
                  environment=dict(python=platform.python_version(), torch=torch.__version__,
                                   numpy=np.__version__, sklearn=sklearn.__version__,
                                   platform=platform.platform(), threads=1, device="cpu"),
                  parameters=sum(p.numel() for p in model(plan).parameters()), bases=[], runs=[])
    started = time.monotonic()

    def checkpoint():
        temporary = out/"record.pending.json"
        temporary.write_text(json.dumps(record, ensure_ascii=False, separators=(",", ":"))+"\n")
        temporary.replace(out/"record.json")

    def save_model(net, name, evaluations):
        path = out/f"{name}.pt"
        torch.save({"plan": plan, "state_dict": net.state_dict()}, path)
        restored = model(plan)
        restored.load_state_dict(torch.load(path, weights_only=True)["state_dict"])
        assert all(evaluate(restored, x, y, split_ids[s]) == ev for s, ev in evaluations.items())
        return dict(file=path.name, sha256=lab.digest(path), reload_equal=True)

    for seed in plan["seeds"]:
        torch.manual_seed(seed)
        old_model = model(plan)
        history, curve = train(old_model, x, y, plan["ids"]["train_old"], plan, seed,
                               plan["old_steps"], split_ids["dev"])
        evaluations = {s: evaluate(old_model, x, y, ids) for s, ids in split_ids.items()}
        fisher_started = time.monotonic()
        importance = fisher_diagonal(old_model, x[plan["memory"]])
        fisher_seconds = time.monotonic()-fisher_started
        base = dict(seed=seed, history=history, curve=curve, evaluations=evaluations,
                    fisher=[f.tolist() for f in importance], fisher_seconds=fisher_seconds,
                    checkpoint=save_model(old_model, f"seed-{seed}-old", evaluations))
        record["bases"].append(base)
        checkpoint()
        for arm in plan["arms"]:
            net = deepcopy(old_model)
            replay = plan["memory"] if arm == "replay50" else plan["ids"]["train_old"] if arm == "replay500" else None
            strength = plan["ewc_lambda"].get(arm, 0.)
            history, curve = train(net, x, y, plan["ids"]["train_new"], plan, seed+30000,
                                   plan["new_steps"], split_ids["dev"], replay,
                                   importance if strength else None, strength)
            evaluations = {s: evaluate(net, x, y, ids) for s, ids in split_ids.items()}
            record["runs"].append(dict(seed=seed, arm=arm, history=history, curve=curve,
                                       old_test_before=base["evaluations"]["test"]["metrics"]["old_accuracy"],
                                       evaluations=evaluations,
                                       checkpoint=save_model(net, f"seed-{seed}-{arm}", evaluations)))
            checkpoint()
            m = evaluations["test"]["metrics"]
            print(f"seed={seed} arm={arm} old={m['old_correct']}/{m['old_n']} "
                  f"new={m['new_correct']}/{m['new_n']}", flush=True)
    record["wall_seconds"] = time.monotonic()-started
    record["complete"] = True
    checkpoint()
    print(f"完成 3 次旧类训练与 15 次后续训练：{out}")


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--freeze", type=Path)
    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 args.freeze:
        lab.save_new(args.freeze, make_plan())
        print(f"已冻结划分、回放样本与全部训练配置：{args.freeze}")
        return
    if not args.run:
        print("默认不训练。使用 --run --out-dir 新目录，才运行全部 18 次训练。")
        return
    if args.out_dir is None:
        parser.error("需要尚不存在的 --out-dir")
    run(args.plan, args.out_dir)


if __name__ == "__main__":
    main()

1. 这等于额外提供任务身份并屏蔽一半输出，改变了测试问题。模型可能仍能在旧类内部区分，却已经把旧图更倾向地判为新类；本实验要检查的正是统一十分类表现。
2. 前后都是 49%。新能力替换了旧能力，平均数看不出来。
3. 旧样本 1,920 次，新样本 5,760 次。若在 50 张缓存中均匀有放回抽样，每张期望出现 $1920/50=38.4$ 次，实际次数会波动。
4. 惩罚为 $[2(0.2)^2+0.5(0.4)^2]/2=0.08$，梯度为 $(0.4,0.2)$。乘以 $\lambda$ 后，两者相应缩放。
5. 不能。它只反映当前模型、选定旧输入与对角近似下的估计，可能漏掉其他输入或参数之间的共同作用，也可能受到数值精度影响。
6. 点积为 1，变化为 $(0.5,-0.5)^T$；输出从 $(0,1)^T$ 变成 $(0.5,0.5)^T$。两个坐标方向都与目标键有重叠。
7. 希望得到 $(2,0)^T$，实际为 $(0.5,1.5)^T$，误差长度为 $\sqrt{(-1.5)^2+1.5^2}=\sqrt{4.5}\approx2.121$。只查最新目标会漏掉先前编辑被覆盖。
8. 前两次预测为 10、9.5，平方误差为 1、2.25；看到第二次结果后，第三次预测为 10.25。不能先用第二次观测更新，再把 10.25 当作第二次预测。
9. 模型可能继续把过期答案用于“现在”的问题。可显式保留有效时间、问题的时间条件与来源，并分别测试历史查询和当前查询；不必把旧记录一律删除。
10. 第二阶段新样本与旧样本各 3,840 次，新类学习量进一步减少。旧类可能保留更多，新类可能受影响，但方向不能由公式保证。新方案要保留全部种子和未改善的结果。