Tufte-style charts

#!/usr/bin/env python3
"""A slopegraph, Tufte style: each item's value at two points, joined by a line, labelled on both
sides, with a thin leader from each end to its label (labels of tied values are pushed apart).
Colour only what changed.

Usage:
  python3 draw_slope.py --spec SPEC.json --out chart.html [--svg chart.svg] [--paper]

Data (the spec's "data", a file beside it or inline), either:
  {"items": [{"name", "from": 12.5, "to": 17.3, "kind"?: "rise"|"fall"|"steady"}, ...]}
or stacked data as draw_stack.py reads it ({"q", "groups": [{"name", "v"}], ...}); then the spec's
"left" and "right" name the two periods, and "measure" is "value" (the group's v) or "share"
(percent of the period's total, the default).

Spec keys: title, subtitle [lines], left_label, right_label (column heads; {left} and {right} are
the period names), value_format ("{:.1f}%" for shares, "{:.1f}" for values), threshold (0.15:
relative change that counts as rise or fall, used when an item has no kind), footnote [lines],
layout {width, x_left, x_right, row (min px between labels), top, plot_height}.

Background: transparent by default, with light and dark palettes switched by prefers-color-scheme
inside the SVG. --paper (or spec "background": "paper") draws the classic #fffff8 page instead,
one palette, as before this option existed.
"""
import argparse

from tufte import (BLUE, FOOT, INK, PAPER, RED, Axis, add_background_arg, background, esc, filler, fmt,
                    load_spec, spread, write)


def items_of(T):
    D = T["data"]
    if isinstance(D.get("items"), list):
        return [dict(i) for i in D["items"]], {}
    Q, groups = D["q"], D["groups"]
    left, right = T.get("left", D.get("from_q", Q[0])), T.get("right", D.get("to_q", Q[-1]))
    for q in (left, right):
        if q not in Q:
            raise SystemExit(f"period {q!r} is not in the data's q: {', '.join(Q)}")
    i0, i1 = Q.index(left), Q.index(right)
    nf = D.get("nfiles")
    for i, q in ((i0, left), (i1, right)):
        empty = (nf and not nf[i]) or any(g["v"][i] is None for g in groups) or (D.get("tot") and D["tot"][i] is None)
        if empty:
            raise SystemExit(f"the {q} bin is empty, so there is nothing to compare; choose other \"left\"/\"right\" periods")
    tot = D.get("tot") or [sum(g["v"][i] for g in groups) + (D.get("rest") or [0] * len(Q))[i] for i in range(len(Q))]
    share = T.get("measure", "share") == "share"
    out = []
    for g in groups[::-1]:  # top of the stack first
        a, b = g["v"][i0], g["v"][i1]
        if share:
            a, b = 100 * a / tot[i0], 100 * b / tot[i1]
        it = {"name": g["name"], "from": a, "to": b}
        if g.get("kind") in ("rise", "fall", "steady", "too few"):
            it["kind"] = g["kind"]
        out.append(it)
    ax = Axis(Q, 0, 1)
    return out, {"left": ax.labels[i0], "right": ax.labels[i1]}


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--spec", required=True)
    ap.add_argument("--out", required=True)
    add_background_arg(ap)
    ap.add_argument("--svg")
    a = ap.parse_args()
    T = load_spec(a.spec)
    L = T.get("layout", {})
    items, names = items_of(T)
    fill = filler(names)
    thr = T.get("threshold", 0.15)
    vf = T.get("value_format", "{:.1f}%" if not isinstance(T["data"].get("items"), list) and T.get("measure", "share") == "share" else "{:.1f}")
    for it in items:
        if "kind" not in it:
            ch = (it["to"] / it["from"] - 1) if it["from"] else (1.0 if it["to"] > 0 else 0.0)
            it["kind"] = "rise" if ch > thr else "fall" if ch < -thr else "steady"

    W = L.get("width", 1000)
    XL, XR = L.get("x_left", 380), L.get("x_right", 620)
    TOP = L.get("top", 170)
    PH = L.get("plot_height", max(300, 26 * len(items)))
    row = L.get("row", 22)
    vmax = max(max(i["from"], i["to"]) for i in items)
    vmin = min(min(i["from"], i["to"]) for i in items)
    vmin = min(0.0, vmin)
    yv = lambda v: TOP + PH - (v - vmin) / (vmax - vmin or 1) * PH
    ly = spread([yv(i["from"]) for i in items], row, TOP, TOP + PH + 40)
    ry = spread([yv(i["to"]) for i in items], row, TOP, TOP + PH + 40)
    H = int(max(max(ly), max(ry)) + 60 + 18 * len(T.get("footnote", [])))

    o = [f'<svg id="fig" viewBox="0 0 {W} {H}" width="{W}" height="{H}" xmlns="http://www.w3.org/2000/svg">',
         f'<rect width="{W}" height="{H}" fill="{PAPER}"/>']
    if T.get("title"):
        o.append(f'<text x="40" y="52" font-size="31" fill="{INK}">{esc(fill(T["title"]))}</text>')
    for k, line in enumerate(T.get("subtitle", [])):
        o.append(f'<text x="40" y="{84 + 22 * k}" font-size="17" fill="#222">{esc(fill(line))}</text>')
    for x, key, anchor in ((XL, "left_label", "end"), (XR, "right_label", "start")):
        o.append(f'<text x="{x}" y="{TOP - 30}" font-size="15" text-anchor="{anchor}" fill="#333">{esc(fill(T.get(key, "{" + key.split("_")[0] + "}")))}</text>')
    for it, y0, y1 in zip(items, ly, ry):
        col = RED if it["kind"] == "rise" else BLUE if it["kind"] == "fall" else "#999"
        txt = RED if it["kind"] == "rise" else BLUE if it["kind"] == "fall" else "#555"
        w = ' font-weight="600"' if it["kind"] in ("rise", "fall") else ""
        a0, a1 = yv(it["from"]), yv(it["to"])
        o.append(f'<line x1="{XL + 30}" x2="{XR - 30}" y1="{a0:.1f}" y2="{a1:.1f}" stroke="{col}" stroke-width="{1.6 if w else 1}"/>')
        # a thin leader from each end of the line to its label, which may have been pushed apart from a tie
        o.append(f'<path d="M{XL + 2},{y0:.1f} C{XL + 16},{y0:.1f} {XL + 14},{a0:.1f} {XL + 30},{a0:.1f}" fill="none" stroke="{col}" stroke-opacity=".7" stroke-width=".7"/>')
        o.append(f'<path d="M{XR - 30},{a1:.1f} C{XR - 14},{a1:.1f} {XR - 16},{y1:.1f} {XR - 2},{y1:.1f}" fill="none" stroke="{col}" stroke-opacity=".7" stroke-width=".7"/>')
        o.append(f'<text x="{XL - 4}" y="{y0 + 5:.1f}" font-size="14" text-anchor="end" fill="{txt}"{w}>{esc(it["name"])}  {esc(fmt(vf, it["from"]))}</text>')
        o.append(f'<text x="{XR + 4}" y="{y1 + 5:.1f}" font-size="14" fill="{txt}"{w}>{esc(fmt(vf, it["to"]))}  {esc(it["name"])}</text>')
    for k, line in enumerate(T.get("footnote", [])):
        o.append(f'<text x="40" y="{H - 24 - 18 * (len(T["footnote"]) - 1 - k)}" font-size="12" font-style="italic" fill="{FOOT}">{esc(fill(line))}</text>')
    o.append("</svg>")
    write(a.out, a.svg, o, W, svg_font=False, bg=background(T, a))
    print(f"wrote {a.out}: {len(items)} items, " + ", ".join(f"{i['kind']} {sum(1 for x in items if x['kind'] == i['kind'])}" for i in {x['kind']: x for x in items}.values()))


if __name__ == "__main__":
    main()