"""Devanagari for Lipi (Hindi, Marathi, Sanskrit): a small rule-based shaper over hand-drawn pen strokes.

Devanagari is not written letter by letter. Text is grouped into syllables (aksharas), and inside a syllable:

  i-matra    ि is written BEFORE the consonant it follows in the text, and its loop reaches over to that consonant
  half forms a consonant with a vertical bar on its right drops the bar before another consonant (स्त, न्य, प्य)
  ligatures  क्ष, त्र, ज्ञ and श्र have their own shapes
  reph       र at the start of a cluster becomes a hook above the end of the syllable (धर्म, कार्य)
  rakar      र at the end of a cluster becomes a slanted tick off the bar (प्र, क्र), or a caret below (ट्र)
  stacking   a consonant with no bar (द ट ड ह ...) carries the next one small underneath it, the Sanskrit way (द्ध, द्व,
             ट्ट, ह्न)
  marks      other vowel signs and ं ँ ः ् ़ hang off the last consonant's bar
  headline   the shirorekha joins a whole word; it is written last, after the word's letters
  fallback   anything else keeps a visible virama (ङ्क्ष -> ङ्‌क्ष), which is still correct text

Units: headline y=0.4, baseline y=10 (see data/devanagari/strokes.txt). Non-Devanagari characters in the text are
drawn with the Hershey futural letters at the same body height.
"""
import math
import re
import unicodedata

import numpy as np

HEAD_Y, BASE_Y = 0.4, 10.0
VIRAMA, NUKTA, ZWJ, ZWNJ = "्", "़", "‍", "‌"
CONSONANTS = set(chr(c) for c in range(0x915, 0x93a)) | set("ॹॺॻॼॽॾॿ")
MATRAS = set("ािीुूृॄॅॆेैॉॊोौ")
MODS = set("ँंः")
INDEPENDENT = {"आ": ["अ", "ा"], "ओ": ["अ", "ा", "@े"], "औ": ["अ", "ा", "@ै"], "ऑ": ["अ", "ा", "@ॅ"],
               "ॲ": ["अ", "@ॅ"], "ई": ["इ", "@reph:6.2"], "ऐ": ["ए", "@े:6.6"]}
LIGATURES = ("क्ष", "त्र", "ज्ञ", "श्र")


def is_deva(text):
    return any(0x900 <= ord(c) <= 0x97F for c in text)


class Deva:
    """The hand-drawn table from weights.safetensors: key -> (advance, bar, head0, head1, span, strokes)."""

    def __init__(self, tensors, keys):
        pts, ss, gs = tensors["glyphs.deva.points"], tensors["glyphs.deva.stroke_start"], tensors["glyphs.deva.glyph_start"]
        self.glyphs, self.marks = {}, {}
        for i, (key, kind) in enumerate(keys):
            adv, bar, h0, h1, span = (float(v) for v in tensors["glyphs.deva.metrics"][i])
            strokes = [pts[ss[k]:ss[k + 1]].astype(np.float64) for k in range(gs[i], gs[i + 1])]
            nan = lambda v: None if math.isnan(v) else v
            (self.glyphs if kind == "glyph" else self.marks)[key] = dict(
                adv=adv, bar=nan(bar), head=None if math.isnan(h0) else (h0, h1), span=nan(span), strokes=strokes)

    def halfable(self, key):
        g = self.glyphs.get(key)
        return g is not None and g["bar"] is not None and g["bar"] >= g["adv"] - 1.85   # float32 metrics: keep slack

    def stackable(self, key, nxt):
        return (key in self.glyphs and self.glyphs[key]["bar"] is None and len(key) == 1 and nxt in self.glyphs
                and len(nxt) == 1)


# ---- shaping: text -> syllables ------------------------------------------------------------------------------------

def syllables(word):
    """-> list of dicts: cons [(consonant, nukta)], matra, mods, final_virama; or other=character."""
    s = unicodedata.normalize("NFD", word).replace(ZWJ, "").replace(ZWNJ, "")
    out, i = [], 0
    while i < len(s):
        c = s[i]
        if c in CONSONANTS:
            cons, final_virama = [], False
            while i < len(s) and s[i] in CONSONANTS:
                nukta = i + 1 < len(s) and s[i + 1] == NUKTA
                cons.append((s[i], nukta))
                i += 1 + nukta
                if i < len(s) and s[i] == VIRAMA:
                    if i + 1 < len(s) and s[i + 1] in CONSONANTS:
                        i += 1
                        continue
                    final_virama, i = True, i + 1
                break
            matra = ""
            if i < len(s) and s[i] in MATRAS:
                matra, i = s[i], i + 1
            mods = ""
            while i < len(s) and s[i] in MODS:
                mods, i = mods + s[i], i + 1
            out.append(dict(cons=cons, matra=matra, mods=mods, final_virama=final_virama))
        else:
            mods = ""
            i += 1
            while c not in MODS and i < len(s) and s[i] in MODS:
                mods, i = mods + s[i], i + 1
            out.append(dict(other=c, mods=mods))
    return out


# ---- layout: syllables -> placed strokes ---------------------------------------------------------------------------

class Word:
    """Strokes of one word in writing order, in units, plus the headline spans."""

    def __init__(self, deva):
        self.d, self.x, self.items, self.heads = deva, 0.0, [], []

    def stroke(self, P, layer, pen=1.0, steady=False):
        self.items.append((P, layer, pen, steady))

    def glyph(self, key, layer, half=False):
        g = self.d.glyphs[key]
        x0 = self.x
        for P in g["strokes"]:
            self.stroke(P + [x0, 0], layer)
        adv = g["adv"]
        if g["bar"] is not None and not half:
            self.stroke(np.array([[x0 + g["bar"], HEAD_Y], [x0 + g["bar"], BASE_Y]]), layer)
        if half:                                                    # the next letter tucks in where the body ends
            adv = max(2.0, max(P[:, 0].max() for P in g["strokes"]) - 1.0)
        if g["head"]:
            h0, h1 = g["head"]
            self.heads.append((x0 + h0, x0 + min(h1, adv), layer))
        self.x += adv
        return dict(x=x0, adv=adv, bar=None if g["bar"] is None or half else x0 + g["bar"])

    def stacked(self, key, over, layer, k=0.58):
        """Draw key small, hanging under the letter `over` (no headline of its own, no advance)."""
        g = self.d.glyphs[key]
        off = np.array([over["x"] + over["adv"] / 2 - g["adv"] * k / 2, 9.4 - HEAD_Y * k])
        for P in g["strokes"]:
            self.stroke(P * k + off, layer, pen=0.75)                # a smaller letter, a finer pen
        if g["bar"] is not None:
            self.stroke(np.array([[g["bar"], HEAD_Y], [g["bar"], BASE_Y]]) * k + off, layer, pen=0.75)

    def mark(self, key, ax, layer, span=None):
        m = self.d.marks[key]
        k = span / m["span"] if span and m["span"] else 1.0
        for P in m["strokes"]:
            self.stroke(P * [k, 1] + [ax, 0], layer)


def anchor(placed):
    """Where marks hang: the consonant's bar, or a little left of its right edge when it has none."""
    return placed["bar"] if placed["bar"] is not None else placed["x"] + placed["adv"] * 0.55


def right_anchor(placed):
    """Where ि's loop lands: a right-hand bar, or near the right edge of a letter with a centre bar (क, फ)."""
    if placed["bar"] is not None and placed["bar"] >= placed["x"] + placed["adv"] - 2:
        return placed["bar"]
    return placed["bar"] + 2.6 if placed["bar"] is not None else placed["x"] + placed["adv"] - 2.4


def fallback_glyph(w, ch, layer, hershey):
    """A character Devanagari has no stroke for (Latin letters, digits, punctuation): Hershey futural, same body."""
    left, right, strokes = hershey(ch)
    k = (BASE_Y - HEAD_Y) / 21.0
    for s in strokes:
        P = s.copy()
        P[:, 0] = (P[:, 0] - left) * k + w.x
        P[:, 1] = (P[:, 1] - 9.0) * k + BASE_Y
        w.stroke(P, layer)
    w.x += (right - left) * k


def lay_syllable(w, syl, layer, hershey):
    d = w.d
    if "other" in syl:
        c = syl["other"]
        seq = INDEPENDENT.get(c, [c])
        placed = None
        for part in seq:
            if part.startswith("@"):
                key, _, at = part[1:].partition(":")
                w.mark(key, placed["x"] + float(at) if at else anchor(placed), layer)
            elif part in d.glyphs:
                placed = w.glyph(part, layer)
            else:
                fallback_glyph(w, part, layer, hershey)
        if placed is not None:
            for m in syl["mods"]:
                if m == "ः":
                    w.glyph("ः", layer)
                else:
                    w.mark(m, anchor(placed), layer)
        return

    cons = [c for c, _ in syl["cons"]]
    nuktas = [n for _, n in syl["cons"]]
    reph = len(cons) > 1 and cons[0] == "र"
    if reph:
        cons, nuktas = cons[1:], nuktas[1:]
    keys = []                                                       # merge ligatures left to right
    i = 0
    while i < len(cons):
        if i + 1 < len(cons) and cons[i] + VIRAMA + cons[i + 1] in LIGATURES:
            keys.append(cons[i] + VIRAMA + cons[i + 1])
            i += 2
        else:
            keys.append(cons[i])
            i += 1
    rakar = len(keys) > 1 and keys[-1] == "र"
    if rakar:
        keys = keys[:-1]
    matra = syl["matra"]
    if keys[-1] == "र" and matra in "ुू" and matra:                 # रु and रू have their own shapes
        keys[-1], matra = "र" + matra, ""

    stem = w.glyph("ि", layer) if matra == "ि" else None
    placed, j = [], 0
    while j < len(keys):
        key, last = keys[j], j == len(keys) - 1
        j += 1
        if key not in d.glyphs:
            continue
        if j == len(keys) - 1 and d.stackable(key, keys[j]):        # last two: द्ध, द्व, ट्ट, ह्न
            p = w.glyph(key, layer)
            w.stacked(keys[j], p, layer)
            placed.append(p)
            j += 1
            continue
        if not last and key + VIRAMA in d.glyphs:                   # क् and फ् are drawn as their own half shapes
            p = w.glyph(key + VIRAMA, layer)
        elif not last and d.halfable(key):
            p = w.glyph(key, layer, half=True)
        else:
            p = w.glyph(key, layer)
            if not last:                                            # no half form: keep the virama visible
                w.mark(VIRAMA, anchor(p), layer)
        if j - 1 < len(nuktas) and nuktas[j - 1]:
            w.mark(NUKTA, p["x"] + p["adv"] * 0.4, layer)
        placed.append(p)
    if not placed:
        return
    base = placed[-1]
    if rakar:
        if base["bar"] is not None:
            w.mark("rakar", base["bar"], layer)
        else:
            w.mark("rakar^", base["x"] + base["adv"] * 0.5, layer)
    if stem is not None:
        w.mark("ि", stem["bar"], layer, span=right_anchor(base) - stem["bar"])
    top = base                                                      # the letter whose bar carries reph and ं
    if matra in "ाीोौॉ" and matra:
        top = w.glyph("ी" if matra == "ी" else "ा", layer)
        if matra in "ोौॉ":
            w.mark({"ो": "े", "ौ": "ै", "ॉ": "ॅ"}[matra], anchor(top), layer)
    elif matra in ("े", "ै", "ॅ"):
        w.mark(matra, anchor(base), layer)
    elif matra in ("ु", "ू", "ृ"):
        w.mark(matra, anchor(base), layer)
    if syl["final_virama"]:
        w.mark(VIRAMA, anchor(base), layer)
    if reph:
        w.mark("reph", anchor(top), layer)
    upper = bool(matra) and matra in "िीेैोौॉॅ"
    for m in syl["mods"]:
        if m == "ः":
            w.glyph("ः", layer)
        elif m == "ँ" and upper:                                    # no room above: candrabindu is written as ं
            w.mark("ं", anchor(top), layer)
        else:
            w.mark(m, anchor(top) + (1.6 if reph else 0), layer)


def lay_word(deva, word, hershey, initial_layer=None):
    """-> Word with strokes in writing order: letters, then the headline joining them."""
    w = Word(deva)
    for i, syl in enumerate(syllables(word)):
        lay_syllable(w, syl, initial_layer if (i == 0 and initial_layer is not None) else 0, hershey)
    spans = sorted(w.heads)
    merged = []
    for a, b, layer in spans:
        if merged and a <= merged[-1][1] + 0.3 and layer == merged[-1][2]:
            merged[-1][1] = max(merged[-1][1], b)
        else:
            merged.append([a, b, layer])
    for a, b, layer in merged:
        w.stroke(np.array([[a - 0.2, HEAD_Y], [b + 0.2, HEAD_Y]]), layer, steady=True)
    return w


def layout_deva(g, text, st, size, max_width, initial, chaikin):
    """Text -> marks (one per stroke, in writing order, pixels), line widths, line height. Same contract as
    a Latin layout. With initial, the first syllable is drawn larger in red, hanging from the shared headline."""
    scale = size / (BASE_Y - HEAD_Y)
    hershey = lambda c: g.glyph("futural", c)
    space = 3.6
    line_h = 2.1 * (BASE_Y - HEAD_Y)
    lines, first = [], initial
    for para in text.split("\n"):
        cur, cur_w = [], 0.0
        for word in para.split():
            w = lay_word(g.deva, word, hershey, 1 if first else None)
            k = 1.0
            if first:
                k, first = st.get("initial_scale", 1.9), False
                w = scale_initial(w, k)
            ww = w.x
            if cur and (cur_w + space + ww) * scale > max_width:
                lines.append((cur, cur_w))
                cur, cur_w = [], 0.0
            cur.append((w, cur_w + (space if cur else 0)))
            cur_w += (space if len(cur) > 1 else 0) + ww
        lines.append((cur, cur_w))
    widths = [lw * scale for _, lw in lines]
    marks = []
    for li, (words, lw) in enumerate(lines):
        dx = (max(widths) / scale - lw) / 2
        for w, x in words:
            for P, layer, pen, steady in w.items:
                P = chaikin(P.copy(), st["smooth"]) + [dx + x, li * line_h]
                marks.append(dict(P=P * scale, size=size * pen * (st.get("initial_scale", 1.9) if layer else 1.0),
                                  layer=layer, steady=steady))
    return marks, widths, line_h * scale


def scale_initial(w, k):
    """Scale the red first syllable about the headline so it hangs below it, and push the rest of the word right."""
    red = [i for i, item in enumerate(w.items) if item[1] == 1]
    if not red:
        return w
    x_end = max(w.items[i][0][:, 0].max() for i in red)
    shift = x_end * (k - 1)
    items = []
    for P, layer, pen, steady in w.items:
        P = P.copy()
        if layer == 1:
            P[:, 0] *= k
            P[:, 1] = HEAD_Y + (P[:, 1] - HEAD_Y) * k
        else:
            P[:, 0] += shift
        items.append((P, layer, pen, steady))
    w.items = items
    w.x += shift
    return w


def seal_strokes(g, text):
    """Seal text (e.g. ॐ, श्री) -> strokes in units, one word, headline included."""
    w = lay_word(g.deva, re.sub(r"\s+", "", text), lambda c: g.glyph("futural", c))
    return [item[0] for item in w.items]