"""数据与学习效率：读取筛选实测记录，计算递归生成的分布变化。"""

from argparse import ArgumentParser
from collections import Counter
from hashlib import sha256
import json
from pathlib import Path
import random
import statistics

DATA = Path(__file__).with_name("data")
PLAN = DATA / "data-efficiency-plan.json"
RECORD = DATA / "data-efficiency-training.json"
LOOPS = DATA / "data-efficiency-loops.json"
ARMS = ["all", "random", "easy", "hard", "balanced", "curriculum"]
LABELS = {"all": "全部候选", "random": "随机子集", "easy": "低损失", "hard": "高损失",
          "balanced": "分类低损失", "curriculum": "渐进课程"}
PROTOCOLS = ["replace", "accumulate", "anchor", "accumulate_sharp"]
LOOP_LABELS = {"replace": "只留新生成", "accumulate": "累积全部", "anchor": "原始加本轮",
               "accumulate_sharp": "累积但偏爱高频"}


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


def probability(counts):
    total = sum(counts)
    if total <= 0 or any(n < 0 for n in counts):
        raise ValueError("计数必须非负，总和必须为正")
    return [n/total for n in counts]


def total_variation(p, q):
    return sum(abs(a-b) for a, b in zip(p, q))/2


def next_counts(original, current, generated, protocol):
    if protocol == "replace":
        return generated[:]
    if protocol in {"accumulate", "accumulate_sharp"}:
        return [a+b for a, b in zip(current, generated)]
    if protocol == "anchor":
        return [a+b for a, b in zip(original, generated)]
    raise ValueError(protocol)


def loop_experiment():
    original = [50, 25, 15, 8, 2]
    config = {"original": original, "seeds": list(range(1000, 1200)),
              "generations": 20, "new_samples": 100,
              "protocols": PROTOCOLS, "source_sha256": digest(__file__)}
    runs = []
    for seed in config["seeds"]:
        paths = {}
        for protocol in PROTOCOLS:
            # Shared initial seed, not shared later samples after distributions diverge.
            rng = random.Random(seed)
            current = original[:]
            history = [current[:]]
            for _ in range(config["generations"]):
                p = probability(current)
                if protocol == "accumulate_sharp":
                    p = probability([v*v for v in p])
                samples = rng.choices(range(len(p)), weights=p, k=config["new_samples"])
                counted = Counter(samples)
                generated = [counted[i] for i in range(len(p))]
                current = next_counts(original, current, generated, protocol)
                history.append(current[:])
            paths[protocol] = history
        runs.append({"seed": seed, "paths": paths})
    return {"config": config, "runs": runs}


def loop_summary(record):
    target = probability(record["config"]["original"])
    result = {}
    for protocol in PROTOCOLS:
        rows = []
        for t in range(record["config"]["generations"]+1):
            counts = [run["paths"][protocol][t] for run in record["runs"]]
            distributions = [probability(c) for c in counts]
            rows.append({"generation": t, "n": sum(counts[0]),
                         "tail_zero": sum(c[-1] == 0 for c in counts),
                         "tail_mean": statistics.mean(p[-1] for p in distributions),
                         "tv_mean": statistics.mean(total_variation(p, target) for p in distributions)})
        result[protocol] = rows
    return result


def mean_training_metrics(record, split="test"):
    result = {}
    for arm in ARMS:
        rows = [run["evaluations"][split]["metrics"] for run in record["runs"] if run["arm"] == arm]
        result[arm] = {"accuracy": statistics.mean(x["accuracy"] for x in rows),
                       "tail_accuracy": statistics.mean(x["tail_accuracy"] for x in rows)}
    return result


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--run-loops", type=Path)
    args = parser.parse_args()
    if args.run_loops:
        with args.run_loops.open("x") as stream:
            json.dump(loop_experiment(), stream, ensure_ascii=False, separators=(",", ":"))
            stream.write("\n")
        print(f"已保存 200 个种子、4 种协议、20 代的记录：{args.run_loops}")
        return
    record = json.loads(RECORD.read_text())
    assert record["plan_sha256"] == digest(PLAN)
    print("留出测试均值：总体 / 数字 8 与 9")
    for arm, m in mean_training_metrics(record).items():
        print(f"{LABELS[arm]}：{100*m['accuracy']:.1f}% / {100*m['tail_accuracy']:.1f}%")
    loops = json.loads(LOOPS.read_text())
    print("第 20 代：罕见类概率为零 / 平均总变差")
    for protocol, rows in loop_summary(loops).items():
        last = rows[-1]
        print(f"{LOOP_LABELS[protocol]}：{last['tail_zero']}/200 / {last['tv_mean']:.4f}")
    print(f"概率 0.02 的类别在 100 次中一次未出现：{.98**100:.4f}")


if __name__ == "__main__":
    main()