"""P12 原创教学计算与配图；合成数据，不读取用户业务材料。"""
from pathlib import Path
from fractions import Fraction as F
from itertools import product
import json
import math
import random
import contextlib
import io
import re
import sys
import zlib

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties

ROOT = Path(__file__).resolve().parents[1]
ASSETS = ROOT / "assets" / "P12"
ASSETS.mkdir(parents=True, exist_ok=True)
font = FontProperties(fname="C:/Windows/Fonts/msyh.ttc")
plt.rcParams.update({"font.family": font.get_name(), "font.size": 11,
                     "axes.unicode_minus": False, "svg.fonttype": "none"})
COLORS = ["#24689b", "#27856d", "#cf7a2c", "#9d557d"]

def save(fig, stem):
    fig.savefig(ASSETS / (stem + ".png"), dpi=160, bbox_inches="tight")
    fig.savefig(ASSETS / (stem + ".svg"), bbox_inches="tight")
    plt.close(fig)

code = {"A": "0", "B": "10", "C": "110", "D": "111"}
p = {"A": F(4, 10), "B": F(3, 10), "C": F(2, 10), "D": F(1, 10)}
assert all(not b.startswith(a) for a in code.values() for b in code.values() if a != b)
assert sum(F(1, 2 ** len(c)) for c in code.values()) == 1
def encode(items):
    return "".join(code[x] for x in items)
def decode(bits):
    reverse = {v: k for k, v in code.items()}
    result, pending = [], ""
    for bit in bits:
        pending += bit
        if pending in reverse:
            result.append(reverse[pending])
            pending = ""
    if pending:
        raise ValueError("末尾码字不完整")
    return "".join(result)
checks = 0
for length in range(1, 8):
    for values in product(code, repeat=length):
        source = "".join(values)
        assert decode(encode(source)) == source
        checks += 1
assert encode("ABAD") == "0100111"
assert decode("0100111") == "ABAD"
bad = {"A": "0", "B": "1", "C": "01"}
assert bad["C"] == bad["A"] + bad["B"]
avg = sum(p[k] * len(code[k]) for k in p)
entropy = -sum(float(v) * math.log2(float(v)) for v in p.values())
assert avg == F(19, 10)
assert entropy <= float(avg) < entropy + 1
feasible_lengths = [ls for ls in product(range(1, 7), repeat=4)
                    if sum(F(1, 2 ** l) for l in ls) <= 1]
bounded_minimum = min(sum(prob * l for prob, l in zip(p.values(), ls))
                      for ls in feasible_lengths)
assert bounded_minimum == avg
assert sum(int(p[k] * 40) * len(code[k]) for k in p) == 76
assert 2 * 160 == F(19, 10) * 160 + 16
shifted = {"A": F(1, 10), "B": F(2, 10), "C": F(3, 10), "D": F(4, 10)}
old_shift = sum(shifted[k] * len(code[k]) for k in code)
new_lengths = {"A": 3, "B": 3, "C": 2, "D": 1}
new_shift = sum(shifted[k] * new_lengths[k] for k in code)
assert old_shift == F(26, 10) and new_shift == F(19, 10)
assert sum(2 ** k for k in range(8)) == 255 < 2 ** 8

# 存储单位及可复现压缩实验。
rng = random.Random(20260907)
samples = {
    "short": b"AI",
    "repeated": b"status=done\n" * 1000,
    "pseudo_random": rng.randbytes(12000),
}
compressed = {}
for label, raw in samples.items():
    packed = zlib.compress(raw, level=6)
    assert zlib.decompress(packed) == raw
    compressed[label] = {"original_bytes": len(raw), "compressed_bytes": len(packed),
                         "packed_to_original": len(packed) / len(raw)}
utf8 = {s: len(s.encode("utf-8")) for s in ["A", "叶", "小叶", "AI"]}
assert utf8 == {"A": 1, "叶": 3, "小叶": 6, "AI": 2}

# 在预先规定的 [0, 100) 范围内作等宽中点量化。
xs = [F(98, 10), F(102, 10), F(147, 10), F(153, 10), F(249, 10), F(251, 10)]
def quantize(x, step):
    assert F(0) <= x < F(100)
    return (x // step) * step + step / 2
quant = {}
for step in [F(10), F(5), F(1)]:
    q = [quantize(x, step) for x in xs]
    errors = [b - a for a, b in zip(xs, q)]
    mse = sum(e * e for e in errors) / len(errors)
    mae = sum(abs(e) for e in errors) / len(errors)
    maxerr = max(abs(e) for e in errors)
    levels = int(F(100) / step)
    quant[str(step)] = {"reconstruction": [float(a) for a in q], "mse": float(mse),
                        "mae": float(mae), "max_abs_error": float(maxerr),
                        "levels": levels, "fixed_bits": math.ceil(math.log2(levels)),
                        "bound": float(step / 2)}
    assert maxerr <= step / 2
assert quant["10"]["reconstruction"] == [5, 15, 15, 15, 25, 25]
assert math.isclose(quant["10"]["mse"], float((2 * F(48, 10)**2 + 2 * F(3, 10)**2 + 2 * F(1, 10)**2) / 6))
assert math.isclose(quant["1"]["mse"], 29 / 300)
# 相同区间，按组标记被无损压缩，不会恢复组内原小数。
coarse = [quantize(x, F(10)) for x in xs]
threshold_original = [x >= 25 for x in xs]
threshold_rebuilt = [x >= 25 for x in coarse]
assert threshold_original == [False, False, False, False, False, True]
assert threshold_rebuilt == [False, False, False, False, True, True]
assert [xs[i] >= 25 for i in range(6)] != threshold_rebuilt
assert quant["5"]["mse"] < quant["10"]["mse"]
assert quant["5"]["mae"] > quant["10"]["mae"]
assert math.isclose(quant["5"]["mse"], float((2*F(23,10)**2 + 2*F(22,10)**2 + 2*F(24,10)**2)/6))
# 用四舍五入的整数量化单独演示二次舍入；明确半值向上。
def round_to(x, step):
    return math.floor(x / step + F(1, 2)) * step
assert round_to(F(149, 10), F(1)) == 15
assert round_to(round_to(F(149, 10), F(1)), F(10)) == 20
assert round_to(F(149, 10), F(10)) == 10
assert round_to(round_to(F(149, 10), F(1)), F(1)) == 15

# 充分统计量例：四次独立同分布二元观察，K=2的六个序列。
seqs = [bits for bits in product([0, 1], repeat=4) if sum(bits) == 2]
assert len(seqs) == 6
for theta in [F(1, 5), F(1, 2), F(4, 5)]:
    masses = [theta ** sum(b) * (1 - theta) ** (4 - sum(b)) for b in seqs]
    assert len(set(masses)) == 1
    assert all(m / sum(masses) == F(1, 6) for m in masses)
# 文本字段权重：三个关键字段5分，两个普通字段1分。
weights = [5, 5, 5, 1, 1]
assert sum(weights) == 17
weighted_missing = {"one_key": F(5, 17), "two_ordinary": F(2, 17)}
assert weighted_missing["one_key"] > weighted_missing["two_ordinary"]

# 从已写正文取出实验代码并实际执行，以免脚本正确而教材片段抄错。
chapter_path = next((ROOT / "chapters" / "04-概率论与数理统计").glob("P12-*.md"))
chapter_text = chapter_path.read_text(encoding="utf-8")
code_excerpt = re.search(r"(    import random\n.*?    +print\(len\(original\), len\(packed\), restored == original\)\n)",
                         chapter_text, flags=re.S).group(1)
code_excerpt = "\n".join(line[4:] if line.startswith("    ") else line for line in code_excerpt.splitlines())
captured = io.StringIO()
with contextlib.redirect_stdout(captured):
    exec(compile(code_excerpt, str(chapter_path) + ":excerpt", "exec"), {})
expected_excerpt = "".join(f"{item['original_bytes']} {item['compressed_bytes']} True\n" for item in compressed.values())
assert captured.getvalue() == expected_excerpt

# 图1：实际前缀树，叶子即完整码字。
fig, ax = plt.subplots(figsize=(10, 5.7))
pos = {"root": (0, 0), "A": (-2, -1), "r": (1.5, -1),
       "B": (.3, -2), "rr": (2.5, -2), "C": (1.7, -3), "D": (3.3, -3)}
edges = [("root", "A", "0"), ("root", "r", "1"),
         ("r", "B", "0"), ("r", "rr", "1"),
         ("rr", "C", "0"), ("rr", "D", "1")]
for parent, child, bit in edges:
    x1, y1 = pos[parent]; x2, y2 = pos[child]
    ax.plot([x1, x2], [y1, y2], color="#a3b2bf", lw=2, zorder=1)
    ax.text((x1+x2)/2, (y1+y2)/2 + .08, bit, ha="center",
            bbox=dict(fc="white", ec="none", pad=2), fontsize=12)
for label in ["root", "r", "rr"]:
    x, y = pos[label]
    ax.scatter(x, y, s=230, color="#52677c", zorder=3)
for i, letter in enumerate(code):
    x, y = pos[letter]
    ax.text(x, y, f"{letter}：{code[letter]}\n概率 {float(p[letter]):.1f}",
            ha="center", va="center", color="white", fontsize=12,
            bbox=dict(boxstyle="round,pad=.5", fc=COLORS[i], ec="none"), zorder=4)
ax.text(-2.3, -2.15, "读取 0100111：\n0 | 10 | 0 | 111\nA → B → A → D",
        fontsize=13, linespacing=1.7)
ax.text(-2.3, -3.35, "到达叶子就输出一个符号，再从根节点读取下一个。", fontsize=11)
ax.set(xlim=(-3.2, 4.2), ylim=(-3.8, .6))
ax.set_title("前缀码：完整码字不再通向另一个码字", fontsize=17, pad=14)
ax.axis("off")
save(fig, "01-prefix-tree")

# 图2：附加开销影响整批期望长度，不画成已测真实文件大小。
fig, axes = plt.subplots(1, 2, figsize=(12, 4.8), gridspec_kw={"width_ratios": [1.1, 1]})
ns = np.arange(0, 401, 10)
ax = axes[0]
ax.plot(ns, 2 * ns, label="定长正文：2N", color=COLORS[0], lw=2.3)
ax.plot(ns, 1.9 * ns + 16, label="变长正文与附加开销：1.9N+16", color=COLORS[2], lw=2.3)
ax.scatter([160], [320], color="#303c48", s=40, zorder=5)
ax.annotate("N=160 时相同", (160, 320), xytext=(200, 220),
            arrowprops=dict(arrowstyle="->", color="#687584"), fontsize=11)
ax.set(xlabel="一批记录数量 N", ylabel="期望长度（bit）", title="总体收益取决于批量")
ax.legend(loc="upper left", fontsize=9)
ax.grid(alpha=.18)
ax = axes[1]
vals = [[80, 92], [400, 396]]
x = np.arange(2)
for j, label in enumerate(["定长正文", "变长正文＋16 bit"]):
    heights = [a[j] for a in vals]
    bars = ax.bar(x + (j - .5)*.3, heights, width=.3, color=[COLORS[0], COLORS[2]][j], label=label)
    for bar, h in zip(bars, heights):
        ax.text(bar.get_x()+bar.get_width()/2, h+8, str(h), ha="center", fontsize=11)
ax.set_xticks(x, ["40 条", "200 条"])
ax.set(ylim=(0, 470), ylabel="长度（bit）", title="按 4∶3∶2∶1 构造两批记录")
ax.legend(fontsize=9, loc="upper left")
fig.suptitle("压缩正文更短，完整表示仍需计算附加开销", fontsize=17)
fig.text(.5, .01, "教学模型：双方公共开销省略；16 bit 是人为设定的额外开销，不是真实文件格式测量。", ha="center", fontsize=10)
fig.tight_layout(rect=[0, .06, 1, .93])
save(fig, "02-coding-overhead")

# 图3：在门槛附近看到数值误差与行动差异。
fig, axes = plt.subplots(1, 2, figsize=(12, 4.8))
v = np.array([float(x) for x in xs])
idx = np.arange(1, 7)
for ax, step in zip(axes, ["10", "1"]):
    rebuilt = np.array(quant[step]["reconstruction"])
    for i, a, b in zip(idx, v, rebuilt):
        ax.plot([i, i], [a, b], color="#b9c2cc", lw=2)
    ax.scatter(idx-.055, v, label="原记录", color=COLORS[0], s=50, zorder=3)
    ax.scatter(idx+.055, rebuilt, label="重建值", color=COLORS[2], marker="s", s=45, zorder=3)
    ax.axhline(25, ls="--", color="#9d557d", lw=1.2)
    ax.text(.65, 26.5, "安排加测门槛：25 秒", fontsize=10, color="#9d557d")
    ax.set(xlabel="第几条记录", ylabel="完成时间（秒）", ylim=(0, 31), xticks=idx,
           title=f"步长 {step} 秒；均方误差 {quant[step]['mse']:.4f} 秒²")
    ax.legend(loc="lower right", fontsize=10)
    ax.grid(alpha=.12)
fig.suptitle("量化：保留区间，放弃区间内的原始位置", fontsize=17)
fig.text(.5, .01, "原值 24.9 秒与 25.1 秒在 10 秒量化中都重建为 25 秒，原来分属门槛两侧的信息被合并。", ha="center", fontsize=10)
fig.tight_layout(rect=[0, .06, 1, .93])
save(fig, "03-quantization")

result = {
    "code": code, "decode_sequences_checked": checks,
    "prefix_average_length": float(avg), "entropy": entropy,
    "prefix_minus_entropy": float(avg) - entropy,
    "prefix_length_search_1_to_6": float(bounded_minimum),
    "shifted_old": float(old_shift), "shifted_new": float(new_shift),
    "compressed": compressed, "utf8": utf8, "quantization": quant,
    "python_version": sys.version, "zlib_runtime_version": zlib.ZLIB_RUNTIME_VERSION,
    "sufficient_statistic_K2_sequences": [list(b) for b in seqs],
    "weighted_missing": {k: float(v) for k, v in weighted_missing.items()},
    "manuscript_code_output": captured.getvalue(),
}
(ROOT / "checks" / "P12-results.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
print(json.dumps(result, ensure_ascii=True, indent=2))