"""Compile the Hershey pen-stroke letters (public domain, US National Bureau of Standards, 1967) and the KanjiVG
stroke-order data (CC BY-SA 3.0) into weights.safetensors.

    python export_weights.py

Nothing here is learned. Each glyph is the centreline a pen follows: a list of strokes, each a list of (x, y) points,
in font units (cap height 21, y down, baseline at y=9). The pen constants in quill.py are written into the same file
so every number the model uses can be inspected in one place.

The safetensors format is an 8-byte little-endian header length, a JSON header, then the raw tensor bytes.
"""
import gzip
import json
import re
import struct
from pathlib import Path

import numpy as np

HERE = Path(__file__).parent
FONTS = ["rowmans", "scripts", "cursive", "futural", "gothiceng"]   # pen centrelines; gothic is multi-stroke
FIRST_CHAR = 32


def parse_jhf(path):
    """-> metrics dict, list of 96 glyphs; a glyph is (left, right, [stroke arrays])."""
    lines = Path(path).read_text().splitlines()
    metrics = json.loads(lines[0][1:])
    glyphs = []
    for line in lines[1:]:
        if not line.strip():
            continue
        data = line[8:]
        val = lambda c: ord(c) - ord("R")
        left, right = val(data[0]), val(data[1])
        strokes, cur = [], []
        for k in range(2, len(data) - 1, 2):
            pair = data[k:k + 2]
            if pair == " R":
                if cur:
                    strokes.append(np.array(cur, np.float32))
                cur = []
            else:
                cur.append((val(pair[0]), val(pair[1])))
        if cur:
            strokes.append(np.array(cur, np.float32))
        glyphs.append((left, right, strokes))
    return metrics, glyphs


# KanjiVG stroke kinds -> how the brush ends the stroke. Hand-set from the Unicode CJK stroke names.
STROKE_KIND = {"tome": 0, "dot": 1, "left_sweep": 2, "right_sweep": 3, "hook": 4, "rise": 5}
KIND_OF = {"㇔": "dot", "㇒": "left_sweep", "㇓": "left_sweep", "㇏": "right_sweep", "㇝": "right_sweep", "㇀": "rise",
           **{c: "hook" for c in "㇚㇆㇟㇂㇉㇈㇖㇙㇌㇋㇢㇡"}}


def svg_path_points(d, per_curve=6):
    """SVG path (M/C/S, absolute or relative) -> polyline; each cubic sampled at per_curve points."""
    toks = re.findall(r"[MmCcSsLl]|-?(?:\d+\.?\d*|\.\d+)(?:e[-+]?\d+)?", d)
    pts, cur, start, last_ctrl, cmd, i = [], np.zeros(2), None, None, None, 0
    t = np.linspace(0, 1, per_curve + 1)[1:, None]
    while i < len(toks):
        if toks[i].isalpha():
            cmd = toks[i]
            i += 1
        nums = lambda k: np.array([float(x) for x in toks[i:i + k]])
        rel = cmd.islower()
        if cmd in "Mm":
            cur = nums(2) + (cur if rel and pts else 0)
            pts.append(cur.copy())
            last_ctrl, i = None, i + 2
            cmd = "l" if rel else "L"
        elif cmd in "Ll":
            cur = nums(2) + (cur if rel else 0)
            pts.append(cur.copy())
            last_ctrl, i = None, i + 2
        else:
            if cmd in "Cc":
                v = nums(6).reshape(3, 2) + (cur if rel else 0)
                c1, c2, p = v
                i += 6
            else:                                                     # S: first control mirrors the last one
                v = nums(4).reshape(2, 2) + (cur if rel else 0)
                c1 = 2 * cur - last_ctrl if last_ctrl is not None else cur.copy()
                c2, p = v
                i += 4
            seg = (1 - t) ** 3 * cur + 3 * (1 - t) ** 2 * t * c1 + 3 * (1 - t) * t ** 2 * c2 + t ** 3 * p
            pts.extend(seg)
            cur, last_ctrl = p, c2
    return np.array(pts, np.float32)


def parse_kanjivg(path):
    """-> list of (codepoint, [(stroke kind, points)])"""
    xml = gzip.open(path, "rt", encoding="utf-8").read()
    out = []
    for m in re.finditer(r'<kanji id="kvg:kanji_([0-9a-f]+)">(.*?)</kanji>', xml, re.S):
        strokes = []
        for p in re.finditer(r"<path\b([^>]*)/>", m.group(2)):
            attrs = dict(re.findall(r'([\w:]+)="([^"]*)"', p.group(1)))
            n = int(re.search(r"-s(\d+)$", attrs["id"]).group(1))
            kind = KIND_OF.get(attrs.get("kvg:type", "").split("/")[0][:1], "tome")
            strokes.append((n, STROKE_KIND[kind], svg_path_points(attrs["d"])))
        out.append((int(m.group(1), 16), [(k, s) for _, k, s in sorted(strokes, key=lambda x: x[0])]))
    return out


def collect():
    from quill import STYLES, INKS, CHAIKIN
    tensors, fonts_meta = {}, {}
    kanji = parse_kanjivg(HERE / "data" / "kanjivg" / "kanjivg-20260714.xml.gz")
    pts, stroke_start, kinds, glyph_start = [], [], [], []
    n_pts = 0
    for _, strokes in kanji:
        glyph_start.append(len(stroke_start))
        for k, s in strokes:
            stroke_start.append(n_pts)
            kinds.append(k)
            pts.append(s)
            n_pts += len(s)
    glyph_start.append(len(stroke_start))
    stroke_start.append(n_pts)
    tensors["glyphs.kanji.points"] = np.concatenate(pts)
    tensors["glyphs.kanji.stroke_start"] = np.array(stroke_start, np.int32)
    tensors["glyphs.kanji.stroke_kind"] = np.array(kinds, np.int32)
    tensors["glyphs.kanji.glyph_start"] = np.array(glyph_start, np.int32)
    tensors["glyphs.kanji.codepoint"] = np.array([c for c, _ in kanji], np.int32)
    for f in FONTS:
        metrics, glyphs = parse_jhf(HERE / "data" / "hershey" / f"{f}.jhf")
        pts, stroke_start, glyph_start = [], [], []
        for _, _, strokes in glyphs:
            glyph_start.append(len(stroke_start))
            for s in strokes:
                stroke_start.append(sum(len(p) for p in pts))
                pts.append(s)
        glyph_start.append(len(stroke_start))
        stroke_start.append(sum(len(p) for p in pts))
        tensors[f"glyphs.{f}.points"] = np.concatenate(pts).astype(np.float32)
        tensors[f"glyphs.{f}.stroke_start"] = np.array(stroke_start, np.int32)
        tensors[f"glyphs.{f}.glyph_start"] = np.array(glyph_start, np.int32)
        tensors[f"glyphs.{f}.advance"] = np.array([(l, r) for l, r, _ in glyphs], np.float32)
        fonts_meta[f] = {k.replace("define_", ""): float(v) for k, v in metrics.items()}
    tensors["smooth.chaikin"] = np.array(CHAIKIN, np.float32)
    for name, rgb in INKS.items():
        tensors[f"ink.{name}"] = np.array(rgb, np.float32) / 255
    meta = {"charset": "".join(chr(FIRST_CHAR + i) for i in range(96)), "fonts": json.dumps(fonts_meta),
            "styles": json.dumps(STYLES), "stroke_kinds": json.dumps(STROKE_KIND),
            "source": "Hershey fonts (public domain), via the MIT-licensed Hershey-Fonts package; "
                      "kanji/kana strokes from KanjiVG (c) Ulrich Apel, CC BY-SA 3.0, http://kanjivg.tagaini.net",
            "units": "Hershey: cap height 21, y down, baseline y=9. KanjiVG: 109 x 109 box, y down"}
    return tensors, meta


def save_safetensors(tensors, meta, path):
    header, blobs, offset = {"__metadata__": meta}, [], 0
    for name, t in tensors.items():
        dt, code = (np.dtype("<i4"), "I32") if t.dtype.kind == "i" else (np.dtype("<f4"), "F32")
        b = np.ascontiguousarray(t, dtype=dt).tobytes()
        header[name] = {"dtype": code, "shape": list(t.shape), "data_offsets": [offset, offset + len(b)]}
        blobs.append(b)
        offset += len(b)
    h = json.dumps(header, separators=(",", ":")).encode()
    h += b" " * (-len(h) % 8)
    Path(path).write_bytes(struct.pack("<Q", len(h)) + h + b"".join(blobs))


def load_safetensors(path):
    raw = Path(path).read_bytes()
    n = struct.unpack("<Q", raw[:8])[0]
    header = json.loads(raw[8:8 + n])
    meta = header.pop("__metadata__", {})
    out = {}
    for name, h in header.items():
        a, b = h["data_offsets"]
        dt = "<i4" if h["dtype"] == "I32" else "<f4"
        out[name] = np.frombuffer(raw[8 + n + a:8 + n + b], dtype=dt).reshape(h["shape"])
    return out, meta


if __name__ == "__main__":
    tensors, meta = collect()
    out = HERE / "weights.safetensors"
    save_safetensors(tensors, meta, out)
    back, _ = load_safetensors(out)
    assert all(np.array_equal(back[k], tensors[k]) for k in tensors)
    n = sum(t.size for t in tensors.values())
    print(f"{len(tensors)} tensors, {n:,} numbers -> {out.name} ({out.stat().st_size / 1024:.0f} KB)")