"""
=============================================================================
CS324 · 第 11 讲共享模块：codec_tools.py —— 用语言模型做真正的无损压缩（算术编码）
=============================================================================

【这个模块是干什么的】
    第 11 讲的核心论断是"预测就是压缩"：一个给下一个字符打出概率 p 的模型，可以把这个字符编码成约 −log2 p 个 bit。
    这个模块把这句话变成能运行的代码：
        ArithmeticEncoder / ArithmeticDecoder   整数实现的算术编码器（Witten–Neal–Cleary 1987 的经典写法）
        quantize(probs)                         把模型的浮点概率变成"整数频率表"（编码器只能处理整数）
        ngram_distributions(...)                第 1 讲 n-gram 模型给文本每个位置的预测分布
        gpt_distributions(...)                  第 1 讲 base model（TinyGPT）给文本每个位置的预测分布
        encode_text / decode_text               用任意一串预测分布把文本压成比特流、再解压回来
        ideal_bits(...)                         理论码长 Σ −log2 p（交叉熵 × 字符数）
    编码器和解码器必须看到**完全相同**的预测分布——解码器每解出一个字符，就把它交给模型，
    得到下一个位置的分布，再解下一个字符。这里为了简单，先一次算好所有位置的分布（编码、解码共用），
    decode_text 的说明里会解释真正的解码器是怎样一个字符一个字符地"边解边预测"的。

【运行方式】
    python3 codec_tools.py     # 自检：随机分布上的编码、解码往返；码长接近理论值（约 1 秒）
    依赖 PyTorch（只有 gpt_distributions 用到），以及 code/lecture01 的 ngram_lm.py、tiny_gpt.py。

【算术编码的直觉】
    把 [0, 1) 区间看成"所有可能的消息"。第一个字符把区间切成若干段，每个可能的字符占一段，长度等于它的概率；
    选中真实字符那一段，在这一段里再按第二个字符的概率切……最后得到一个很小的区间，长度 = 所有概率的乘积。
    要指明这个区间里的一个点，大约需要 −log2(区间长度) = Σ −log2 p 个 bit。
    整数实现：用 32 位整数表示区间的上下界，每当区间落进 [0, 1/2) 或 [1/2, 1) 时就输出一个确定的 bit 并把区间放大一倍，
    这样不需要无限精度。

【Python 小知识】
    - Python 的整数没有位数上限，所以 32 位的区间运算不会溢出；我们用 & MASK 之类的位运算只是为了表达"32 位"。
    - itertools.accumulate(列表)：前缀和，[3, 1, 2] → [3, 4, 6]。
"""

import itertools
import math
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

PRECISION = 32
FULL = (1 << PRECISION) - 1          # 0xFFFFFFFF
HALF = 1 << (PRECISION - 1)
QUARTER = 1 << (PRECISION - 2)
TOTAL = 1 << 16                      # 频率表的总和：每个分布量化成和为 65536 的整数


# =============================================================================
# 1. 把概率量化成整数频率
# =============================================================================
def quantize(probs, total=TOTAL):
    """
    把一个概率分布（浮点数列表）变成和为 total 的正整数频率表。

    【思路】每个符号先分到 1（保证概率不为 0——否则真的出现时无法编码），剩下的 total − V 按概率比例分配、向下取整，
            最后把凑不齐的零头全部加到概率最大的那个符号上。量化会让码长比理论值略长一点（代价通常 < 0.01 bit/字符）。
    【参数】probs：长度 V 的非负浮点数列表（和应为 1，不严格要求）；total：频率总和
    【返回】list[int]，长度 V，每个 ≥ 1，和为 total
    【边界】V > total 时无法保证每个都 ≥ 1，抛出 ValueError。
    """
    V = len(probs)
    if V > total:
        raise ValueError("符号太多，频率表放不下")
    s = math.fsum(probs)
    spare = total - V
    freqs = [1 + int(p / s * spare) for p in probs]
    freqs[max(range(V), key=lambda i: probs[i])] += total - sum(freqs)
    return freqs


# =============================================================================
# 2. 算术编码器与解码器
# =============================================================================
class ArithmeticEncoder:
    """
    整数算术编码器。用法：enc = ArithmeticEncoder(); enc.encode(freqs, symbol) 若干次; bits = enc.finish()

    【属性】low, high：当前区间（闭区间，32 位整数）；pending：还没决定的"下溢"位数；bits：输出的 0/1 列表
    """

    def __init__(self):
        self.low, self.high, self.pending, self.bits = 0, FULL, 0, []

    def _emit(self, bit):
        self.bits.append(bit)
        self.bits.extend([1 - bit] * self.pending)      # 之前悬而未决的位，现在知道它们和 bit 相反
        self.pending = 0

    def encode(self, freqs, symbol):
        """
        编码一个符号。

        【思路】区间 [low, high] 按频率表切成段，symbol 对应 [cum[symbol], cum[symbol+1]) 那一段；
                把区间缩到这一段，然后只要区间完全落在下半或上半，就输出一个 bit 并放大区间。
                如果区间跨在正中间、但已经很窄（落在 [1/4, 3/4) 内），就把它以中点为中心放大，
                并记下一个"待定位"（pending）——等以后知道往哪边走了再一起输出。
        【参数】freqs：quantize 得到的整数频率表；symbol：符号的下标
        """
        total = sum(freqs)
        cum_low = sum(freqs[:symbol])
        cum_high = cum_low + freqs[symbol]
        rng = self.high - self.low + 1
        self.high = self.low + rng * cum_high // total - 1
        self.low = self.low + rng * cum_low // total
        while True:
            if self.high < HALF:
                self._emit(0)
            elif self.low >= HALF:
                self._emit(1)
                self.low -= HALF
                self.high -= HALF
            elif self.low >= QUARTER and self.high < 3 * QUARTER:
                self.pending += 1
                self.low -= QUARTER
                self.high -= QUARTER
            else:
                break
            self.low = 2 * self.low
            self.high = 2 * self.high + 1

    def finish(self):
        """结束编码：再输出两位，足以让解码器确定最后的区间。返回 bit 列表。"""
        self.pending += 1
        self._emit(0 if self.low < QUARTER else 1)
        return self.bits


class ArithmeticDecoder:
    """
    与 ArithmeticEncoder 对应的解码器。用法：dec = ArithmeticDecoder(bits); s = dec.decode(freqs) 若干次。

    【思路】解码器维护同样的区间，以及一个"当前读到的数" value；每一步看 value 落在哪一段，就解出哪个符号，
            然后做和编码器完全相同的区间缩放。比特流读完之后，后面当作 0。
    """

    def __init__(self, bits):
        self.bits, self.pos = bits, 0
        self.low, self.high, self.value = 0, FULL, 0
        for _ in range(PRECISION):
            self.value = (self.value << 1) | self._next_bit()

    def _next_bit(self):
        b = self.bits[self.pos] if self.pos < len(self.bits) else 0
        self.pos += 1
        return b

    def decode(self, freqs):
        """解出一个符号，返回它的下标。freqs 必须和编码时这个位置用的频率表完全相同。"""
        total = sum(freqs)
        rng = self.high - self.low + 1
        target = ((self.value - self.low + 1) * total - 1) // rng
        cums = list(itertools.accumulate(freqs))
        symbol = next(i for i, c in enumerate(cums) if c > target)
        cum_low = cums[symbol] - freqs[symbol]
        self.high = self.low + rng * cums[symbol] // total - 1
        self.low = self.low + rng * cum_low // total
        while True:
            if self.high < HALF:
                pass
            elif self.low >= HALF:
                self.low -= HALF
                self.high -= HALF
                self.value -= HALF
            elif self.low >= QUARTER and self.high < 3 * QUARTER:
                self.low -= QUARTER
                self.high -= QUARTER
                self.value -= QUARTER
            else:
                break
            self.low = 2 * self.low
            self.high = 2 * self.high + 1
            self.value = (self.value << 1) | self._next_bit()
        return symbol


def encode_text(text, alphabet, dists):
    """
    用一串预测分布把 text 压成比特流。

    【参数】text：字符串；alphabet：字母表；dists：list，第 i 个元素是预测第 i 个字符的概率分布（长度 = |alphabet|）
    【返回】bit 列表（0/1）
    【边界】len(dists) 必须等于 len(text)；text 里不能有字母表之外的字符。
    """
    assert len(dists) == len(text), "每个字符都要有一个预测分布"
    index = {c: i for i, c in enumerate(alphabet)}
    enc = ArithmeticEncoder()
    for ch, p in zip(text, dists):
        enc.encode(quantize(p), index[ch])
    return enc.finish()


def decode_text(bits, alphabet, dists):
    """
    把比特流解压回文本。

    【思路】真正的解码器不能提前知道所有的预测分布（分布依赖于前面的字符，而前面的字符正是要解出来的东西）：
            它每解出一个字符，就把它交给模型、算出下一个位置的分布，再解下一个。只要编码器和解码器用的是
            **同一个模型、同样的计算**，两边的分布就逐位相同，解码就是精确的。
            这里传入的 dists 是编码时算好的那一串（为了让演示快），效果和"边解边预测"相同；
            代码 01 演示 3 对 n-gram 模型真的做了一次"边解边预测"，确认两种方式结果一样。
    【返回】字符串
    """
    dec = ArithmeticDecoder(bits)
    return "".join(alphabet[dec.decode(quantize(p))] for p in dists)


def ideal_bits(text, alphabet, dists):
    """理论码长 Σ −log2 p(真实字符)（bit）。实际码长只比它多一点点（量化和结尾的几个 bit）。"""
    index = {c: i for i, c in enumerate(alphabet)}
    return math.fsum(-math.log2(p[index[ch]]) for ch, p in zip(text, dists))


# =============================================================================
# 3. 模型给出的预测分布
# =============================================================================
def ngram_distributions(lm, text):
    """
    第 1 讲 CharNGramLM 对 text 每个位置的预测分布（顺序和 lm._events 相同：每行开头补 n−1 个换行）。

    【返回】list[list[float]]，长度 = len(text)
    """
    return [lm.distribution(h) for h, _ in lm._events(text)]


def gpt_distributions(model, tok, text, stride=48):
    """
    第 1 讲 TinyGPT 对 text 每个位置的预测分布，上下文最多 block_size − 1 个字符。

    【思路】在 text 前面补一个 "\\n" 当作开头（第一个字符也有一个上下文）。用长 block_size 的窗口滑过整段文本，
            每次前进 stride 个字符：第一个窗口的所有预测都用；之后每个窗口只用最后 stride 个位置的预测
            （它们的上下文至少有 block_size − stride 个字符）。这样每个位置恰好被预测一次。
    【参数】model：TinyGPT；tok：CharTokenizer；text：要压缩的文本；stride：窗口每次前进多少
    【返回】list[list[float]]，长度 = len(text)
    【边界】一次前向就把所有窗口放进一个 batch；文本很长时会占较多内存（几十万字符也只要几百 MB）。
    """
    import torch
    ids = [tok.stoi["\n"]] + tok.encode(text)
    T = model.cfg.block_size
    out = [None] * len(text)
    starts = list(range(0, max(1, len(ids) - T), stride))
    if starts[-1] + T < len(ids):
        starts.append(len(ids) - T)
    model.eval()
    with torch.no_grad():
        for b0 in range(0, len(starts), 256):
            batch = starts[b0:b0 + 256]
            x = torch.tensor([ids[s:s + T] for s in batch if s + T <= len(ids)] or [ids], dtype=torch.long)
            probs = torch.softmax(model(x).double(), dim=-1).tolist()
            for row, s in zip(probs, batch):
                first = 0 if s == 0 else T - stride
                for j in range(first, min(T, len(ids) - 1 - s)):
                    pos = s + j                              # ids[s + j] 预测 ids[s + j + 1]，即 text[s + j]
                    if pos < len(text) and out[pos] is None:
                        out[pos] = row[j]
    assert all(o is not None for o in out)
    return out


if __name__ == "__main__":
    import random
    rng = random.Random(0)
    alphabet = "abcdefgh"
    for skew in (0.5, 3.0):
        dists, text = [], []
        for _ in range(3000):
            w = [rng.random() ** skew for _ in alphabet]
            s = sum(w)
            p = [x / s for x in w]
            dists.append(p)
            text.append(rng.choices(alphabet, weights=p)[0])
        text = "".join(text)
        bits = encode_text(text, alphabet, dists)
        assert decode_text(bits, alphabet, dists) == text
        ideal = ideal_bits(text, alphabet, dists)
        print(f"skew {skew}: 理论 {ideal:.1f} bit，实际 {len(bits)} bit（多 {len(bits) - ideal:.1f} bit），解码一致 ✓")
    print("自检通过。")