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

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()