"""
=============================================================================
CS525 · 第 4 讲模块：一个 200 行的"迷你 Parquet"（minicol.py）
=============================================================================

对应讲义：lectures/lecture04_data_tools.md §3.2

-----------------------------------------------------------------------------
【Parquet 的几个核心想法，这里全部手写一遍】
-----------------------------------------------------------------------------
    1) 列式存储（columnar）：同一列的值存在一起。只要两列，就只读两列——slides："读数据常常比算数据还贵"。
    2) 类型（schema）：每列有固定的类型（int64 / float64 / string），写入时检查，读出来不用再解析。
    3) 编码 + 压缩：低基数的字符串列用字典编码（"en" 只存一次，每行只存一个小整数）；
       每一列块再用 zlib 压缩——同一列的值彼此相似，压缩率比按行存高得多。
    4) 行组（row group）+ 统计信息：每 R 行切成一组，每组每列记下最小值和最大值。
       查询 "quality > 0.9" 时，max(quality) ≤ 0.9 的整组可以直接跳过（谓词下推，predicate pushdown）。
    5) 文件末尾的元数据（footer）：记录 schema、每个列块在文件里的位置和统计信息。读文件时先读 footer，
       再只 seek 到需要的列块。

真实的 Parquet 还有很多细节（页、嵌套类型、多种编码、比 zlib 更快的压缩算法……），slides 也提到它的不足：
文件不可修改（immutable）、行组大小绑定、元数据格式老旧。这里只保留最核心的部分，方便读懂。

【文件格式】
    [列块 1][列块 2]……[列块 N][footer 的 JSON][footer 长度（8 字节）]

【运行自检】
    python3 minicol.py        # 写一个小文件再读回来，检查内容一致
"""

import json
import struct
import zlib

TYPES = ("int64", "float64", "string")


def _check_type(name, typ, values):
    """
    类型检查：每个值都必须是声明的类型，否则抛出 TypeError（指出列名和第一个坏值）。

    【为什么】JSONL 不检查类型：quality 列里混进一个 "N/A"，要等到下游计算时才炸。
            有类型的格式在写入时就拒绝它——slides：Parquet "enforces types"。
    【边界】bool 在 Python 里是 int 的子类，这里把它当作不合法的 int64（避免 True 被当成 1 悄悄写进去）。
    """
    for i, v in enumerate(values):
        ok = ((typ == "int64" and isinstance(v, int) and not isinstance(v, bool))
              or (typ == "float64" and isinstance(v, float))
              or (typ == "string" and isinstance(v, str)))
        if not ok:
            raise TypeError(f"列 {name}（{typ}）的第 {i} 个值类型不对：{v!r}")


def _encode(typ, values, dict_threshold):
    """
    把一个列块编码成字节串（压缩前）。

    【返回】(字节串, 编码方式) —— 编码方式是 "plain" 或 "dict"
    【思路】int64 / float64：用 struct 按小端格式紧密排列（每个值 8 字节）。
            string：如果不同的值不超过 dict_threshold 个，就做字典编码：
                字典（JSON 列表）+ 每行一个 uint16 下标；
            否则按"长度前缀 + UTF-8 字节"逐个排列。
    【Python 小知识】struct.pack("<3q", 1, 2, 3) 把 3 个整数按小端 8 字节整数打包；"<d" 是 8 字节浮点数，
                    "<H" 是 2 字节无符号整数，"<I" 是 4 字节无符号整数。
    """
    if typ == "int64":
        return struct.pack(f"<{len(values)}q", *values), "plain"
    if typ == "float64":
        return struct.pack(f"<{len(values)}d", *values), "plain"
    uniq = sorted(set(values))
    if len(uniq) <= dict_threshold:
        index = {v: i for i, v in enumerate(uniq)}
        head = json.dumps(uniq).encode("utf-8")
        return struct.pack("<I", len(head)) + head + struct.pack(f"<{len(values)}H", *(index[v] for v in values)), "dict"
    parts = []
    for v in values:
        b = v.encode("utf-8")
        parts.append(struct.pack("<I", len(b)))
        parts.append(b)
    return b"".join(parts), "plain"


def _decode(typ, enc, raw, n):
    """_encode 的逆过程：字节串 → Python 值的列表（n 个）。"""
    if typ == "int64":
        return list(struct.unpack(f"<{n}q", raw))
    if typ == "float64":
        return list(struct.unpack(f"<{n}d", raw))
    if enc == "dict":
        (hl,) = struct.unpack_from("<I", raw, 0)
        uniq = json.loads(raw[4:4 + hl].decode("utf-8"))
        idx = struct.unpack_from(f"<{n}H", raw, 4 + hl)
        return [uniq[i] for i in idx]
    out, pos = [], 0
    for _ in range(n):
        (L,) = struct.unpack_from("<I", raw, pos)
        out.append(raw[pos + 4:pos + 4 + L].decode("utf-8"))
        pos += 4 + L
    return out


def write_table(path, cols, schema, row_group_size=8192, dict_threshold=1000, compress=True):
    """
    把列字典写成一个迷你 Parquet 文件。

    【参数】path：文件路径；cols：{列名: 值列表}；schema：{列名: 类型}（顺序决定列的顺序）；
            row_group_size：每个行组多少行；dict_threshold：字符串列"不同的值不超过这么多"就用字典编码；
            compress：是否用 zlib 压缩每个列块
    【返回】footer（字典）：包含每个列块的位置、大小、编码方式和统计信息
    【边界】所有列的长度必须相同；字典编码最多 65535 个不同的值（uint16 下标）。
    """
    names = list(schema)
    n = len(cols[names[0]])
    for name in names:
        if len(cols[name]) != n:
            raise ValueError(f"列 {name} 的长度 {len(cols[name])} 和第一列 {n} 不同")
        _check_type(name, schema[name], cols[name])
    footer = {"schema": schema, "n_rows": n, "compress": compress, "row_groups": []}
    with open(path, "wb") as f:
        for start in range(0, n, row_group_size):
            stop = min(start + row_group_size, n)
            rg = {"start": start, "n": stop - start, "columns": {}}
            for name in names:
                vals = cols[name][start:stop]
                raw, enc = _encode(schema[name], vals, dict_threshold)
                data = zlib.compress(raw, 6) if compress else raw
                meta = {"offset": f.tell(), "size": len(data), "raw_size": len(raw), "encoding": enc}
                if schema[name] != "string":
                    meta["min"], meta["max"] = min(vals), max(vals)
                f.write(data)
                rg["columns"][name] = meta
            footer["row_groups"].append(rg)
        head = json.dumps(footer).encode("utf-8")
        f.write(head)
        f.write(struct.pack("<Q", len(head)))
    return footer


def read_footer(path):
    """读文件末尾的 footer：先读最后 8 个字节得到 footer 的长度，再往前读 footer。"""
    with open(path, "rb") as f:
        f.seek(-8, 2)
        (hl,) = struct.unpack("<Q", f.read(8))
        f.seek(-8 - hl, 2)
        return json.loads(f.read(hl).decode("utf-8"))


class Reader:
    """
    迷你 Parquet 的读取器：只读需要的列、可以按统计信息跳过行组，并记录一共读了多少字节。

    【用法】
        r = Reader(path)
        vals = r.read_columns(["quality"])                       # 读一列
        vals = r.read_columns(["url"], where=("quality", ">", 0.9))   # 带谓词下推
        r.bytes_read                                             # 读了多少字节（不含 footer）
    """

    def __init__(self, path):
        self.path = path
        self.footer = read_footer(path)
        self.bytes_read = 0
        self.groups_skipped = 0

    def _chunk(self, f, rg, name):
        meta = rg["columns"][name]
        f.seek(meta["offset"])
        data = f.read(meta["size"])
        self.bytes_read += meta["size"]
        raw = zlib.decompress(data) if self.footer["compress"] else data
        return _decode(self.footer["schema"][name], meta["encoding"], raw, rg["n"])

    def read_columns(self, names, where=None):
        """
        读若干列，返回 {列名: 值列表}。

        【参数】names：要返回的列；where：None 或 (列名, ">" 或 "<" 或 "==", 值) 形式的简单条件
        【思路】有 where 时：先看每个行组的统计信息，不可能有满足条件的行的整组跳过；
                对剩下的组，先读条件列算出哪些行满足，再读其他列、只保留这些行。
        【边界】字符串列没有记录 min / max，只能逐组读（不能跳过）。
        """
        out = {k: [] for k in names}
        schema = self.footer["schema"]
        with open(self.path, "rb") as f:
            for rg in self.footer["row_groups"]:
                keep = None
                if where is not None:
                    col, op, val = where
                    meta = rg["columns"][col]
                    if schema[col] != "string":
                        if (op == ">" and meta["max"] <= val) or (op == "<" and meta["min"] >= val) or \
                           (op == "==" and not (meta["min"] <= val <= meta["max"])):
                            self.groups_skipped += 1
                            continue
                    cv = self._chunk(f, rg, col)
                    if op == ">":
                        keep = [i for i, v in enumerate(cv) if v > val]
                    elif op == "<":
                        keep = [i for i, v in enumerate(cv) if v < val]
                    else:
                        keep = [i for i, v in enumerate(cv) if v == val]
                    if not keep:
                        continue
                for name in names:
                    vals = cv if (where is not None and name == where[0]) else self._chunk(f, rg, name)
                    out[name].extend(vals if keep is None else [vals[i] for i in keep])
        return out


if __name__ == "__main__":
    import os
    import tempfile

    cols = {"id": [1, 2, 3, 4, 5], "q": [0.1, 0.9, 0.5, 0.95, 0.2], "lang": ["en", "de", "en", "en", "zh"]}
    schema = {"id": "int64", "q": "float64", "lang": "string"}
    with tempfile.TemporaryDirectory() as d:
        p = os.path.join(d, "t.mcol")
        write_table(p, cols, schema, row_group_size=2)
        r = Reader(p)
        print("全部读回：", r.read_columns(["id", "q", "lang"]) == cols)
        r2 = Reader(p)
        print("q > 0.8 的 id：", r2.read_columns(["id"], where=("q", ">", 0.8))["id"], "跳过的行组：", r2.groups_skipped)
        try:
            write_table(p, {"id": [1, "2"]}, {"id": "int64"})
        except TypeError as e:
            print("类型检查：", e)