"""Compile Lipi's hand-drawn Devanagari pen strokes and the Hershey futural letters (public domain, US National Bureau
of Standards, 1967; used for Latin letters and digits inside Devanagari text) 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. The pen constants in lipi.py are written into the same file
so every number the model uses can be inspected in one place. The type hands' outlines are not in this file: they
are read from the font files in fonts/, which keep their own licence (SIL Open Font License 1.1).

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

import numpy as np

HERE = Path(__file__).parent
FONTS = ["futural"]                     # Hershey pen centrelines for Latin letters and digits
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


def catmull_rom(P, per_seg=6):
    """Curve through every point; the tangent at each point is the chord between its neighbours."""
    if len(P) < 3:
        return P
    Q = np.vstack([2 * P[0] - P[1], P, 2 * P[-1] - P[-2]])
    t = np.linspace(0, 1, per_seg + 1)[1:, None]
    out = [P[:1]]
    for i in range(1, len(Q) - 2):
        p0, p1, p2, p3 = Q[i - 1], Q[i], Q[i + 1], Q[i + 2]
        out.append(0.5 * (2 * p1 + (p2 - p0) * t + (2 * p0 - 5 * p1 + 4 * p2 - p3) * t ** 2
                          + (3 * p1 - p0 - 3 * p2 + p3) * t ** 3))
    return np.vstack(out)


def parse_deva(path):
    """Hand-drawn Devanagari (data/devanagari/strokes.txt) -> list of (key, kind, [adv, bar, head0, head1, span],
    [stroke arrays]). Missing numbers are NaN."""
    out = []
    for line in Path(path).read_text(encoding="utf-8").splitlines():
        if not line.strip() or line.lstrip().startswith("#"):
            continue
        if not line[0].isspace():
            kind, key, *rest = line.split()
            m = {"bar": np.nan, "head": None, "span": np.nan}
            adv = float(rest.pop(0)) if kind == "glyph" else 0.0
            for r in rest:
                k, v = r.split("=")
                m[k] = v if k == "head" else float(v)
            h0, h1 = (0.0, adv) if m["head"] is None else (np.nan, np.nan) if m["head"] == "-" else \
                tuple(map(float, m["head"].split(":")))
            out.append((key, kind, [adv, m["bar"], h0, h1, m["span"]], []))
            continue
        pts, runs = [], []
        for tok in line.split():
            corner = tok.endswith("!")
            pts.append([float(v) for v in tok.rstrip("!").split(",")])
            if corner:                                                # a corner ends one smooth run, starts the next
                runs.append(pts)
                pts = [pts[-1]]
        runs.append(pts)
        curve = [catmull_rom(np.array(r)) for r in runs if len(r) > 1] or [np.array(runs[0])]
        out[-1][3].append(np.vstack([curve[0]] + [c[1:] for c in curve[1:]]).astype(np.float32))
    return out


def collect():
    from lipi import STYLES, INKS, CHAIKIN
    tensors, fonts_meta = {}, {}
    deva = parse_deva(HERE / "data" / "devanagari" / "strokes.txt")
    pts, stroke_start, glyph_start = [], [], []
    for _, _, _, strokes in deva:
        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["glyphs.deva.points"] = np.concatenate(pts)
    tensors["glyphs.deva.stroke_start"] = np.array(stroke_start, np.int32)
    tensors["glyphs.deva.glyph_start"] = np.array(glyph_start, np.int32)
    tensors["glyphs.deva.metrics"] = np.array([m for _, _, m, _ in deva], np.float32)
    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, ensure_ascii=False),
            "deva_keys": json.dumps([[k, kind] for k, kind, _, _ in deva], ensure_ascii=False),
            "deva_metrics": "advance, bar x, headline start, headline end, span (NaN = none)",
            "source": "Devanagari strokes drawn by hand for Lipi (CC BY 4.0); Hershey futural letters (public domain), "
                      "via the MIT-licensed Hershey-Fonts package",
            "units": "Devanagari: headline y=0.4, baseline y=10, y down. Hershey: cap height 21, y down, baseline y=9"}
    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], equal_nan=back[k].dtype.kind == "f") 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)")