Tufte-style charts

#!/usr/bin/env python3
"""Several series over periods, Tufte style: one line each, every point labelled with its value,
each series named at its right end, all on one y scale.

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

Spec keys (all text comes from here):
  data          a JSON file beside the spec, or an object. Whatever shape; series point into it.
  periods       the periods to draw, e.g. ["2025Q2", ...] (2026Q1, 2026-03, 2026, or any labels)
  series        [{label, color, values | from, counts?}] in drawing order. Values may be null:
                the line breaks there. A series gets its values from:
                  "values": [...]                       one per period
                  "from": {"sample": S, "key": K}       data[S][K], matched to periods by data[S]["Q"]
                  "from": {"key": K}                    data[K] at the top level, matched by data["Q"]
                        or data["q"] (e.g. "tot" or "med" of text-rule-trends' pooled.json)
                  "from": {"sample": S, "share_of": A, "among": [A, B, ...]}
                        100 * A / (A + B + ...) from the counts data[S][A], data[S][B] ...;
                        the counts behind each point become the series' counts
                "color" is a palette name (ink, red, blue, grey) or a palette hex, which follow dark
                mode, or any other CSS colour, kept as given in both modes.
                "counts" (optional) are printed under the axis, joined across the series that
                have them with " · "; series without counts are left out of that line.
  y             {min: 0, max: 100, step: 25, tick: "{:.0f}", reference: null}. tick is the
                format of the y-axis numbers: "{:.0f}" (default) for counts, "{:.0f}%" for shares.
                reference draws a faint dashed line at that value (e.g. 50 for shares).
  value_format  how each point's value is printed, default "{:.0f}" ("{:.0f}%" for shares)
  title, subtitle [lines], ylabel, count_label, quote {text, date, mark}, events [{date, label}],
  event_anchor {label: [anchor, dx]}, footnote [lines], layout {width, height, x0, x1, y_bottom, y_top}
  {placeholders} in subtitle and footnote: {a[i]}, {b[i]}, ... are each series' values (series in
  order, i over the periods drawn), {a_last} the last value that is not null, and {data[...]}
  reaches into the data (e.g. {data[random][history_files]:,}). Negative numbers print with a
  true minus.

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 datetime import date

from tufte import (FOOT, INK, PAPER, RED, Axis, add_background_arg, background, color, esc, events, filler, fmt,
                    load_spec, parse_events, write)


def series_values(s, D, periods):
    if "values" in s:
        v = list(s["values"])
        return v, s.get("counts")
    f = s["from"]
    d = D[f["sample"]] if "sample" in f else D           # no sample: the data's top level (pooled.json)
    dq = d.get("Q") or d.get("q")
    if dq is None:
        raise SystemExit("series data needs its periods as \"Q\" or \"q\"")
    idx = {q: i for i, q in enumerate(dq)}
    if "key" in f:
        v = [d[f["key"]][idx[q]] if q in idx else None for q in periods]
        return v, s.get("counts")
    cnt = {k: [d[k][idx[q]] if q in idx else 0 for q in periods] for k in f["among"]}
    tot = [sum(cnt[k][i] for k in f["among"]) for i in range(len(periods))]
    v = [cnt[f["share_of"]][i] / tot[i] * 100 if tot[i] else None for i in range(len(periods))]
    return v, tot


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)
    D = T["data"]
    L = T.get("layout", {})

    P = T.get("periods") or D.get("q") or D.get("Q")
    series = []
    for s in T["series"]:
        v, n = series_values(s, D, P)
        series.append({**s, "color": color(s["color"]), "v": v, "n": n})
    names = "abcdefghijklmnop"
    ph = {}
    for i, s in enumerate(series):
        ph[names[i]] = s["v"]
        known = [x for x in s["v"] if x is not None]
        ph[names[i] + "_last"] = known[-1] if known else float("nan")
    ph["data"] = D
    fill = filler(ph)
    Y = {"min": 0, "max": 100, "step": 25, "tick": "{:.0f}", **T.get("y", {})}
    vfmt = T.get("value_format", "{:.0f}")

    W, H = L.get("width", 1200), L.get("height", 720)
    X0, X1 = L.get("x0", 70), L.get("x1", 790)
    YB, YT = L.get("y_bottom", 560), L.get("y_top", 250)
    ax = Axis(P, X0, X1)
    xs = ax.mid
    lo, hi = Y["min"], Y["max"]
    yv = lambda v: YB - (v - lo) / (hi - lo) * (YB - YT)

    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="{X0}" 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="{X0}" y="{84 + 22 * k}" font-size="17" fill="#222">{esc(fill(line))}</text>')
    header_bottom = max([60] + [84 + 22 * k for k in range(len(T.get("subtitle", [])))])
    hdr_end = len(o)
    quote = T.get("quote")
    if quote:
        bx = ax.date_x(date.fromisoformat(quote["date"]))
        mark = f'<tspan dy="-8" font-size="12">{esc(quote["mark"])}</tspan>' if quote.get("mark") else ""
        o.append(f'<text x="{bx + 10:.1f}" y="150" font-size="18" fill="{RED}">{esc(quote["text"])}{mark}</text>')
    if T.get("events"):
        events(o, ax.date_x, parse_events(T), 214, YB, {k: tuple(v) for k, v in T.get("event_anchor", {}).items()})
    if quote:
        o.append(f'<line x1="{bx:.1f}" x2="{bx:.1f}" y1="136" y2="{YB}" stroke="{RED}" stroke-width="1.8"/>')

    # y axis: a range frame, and an optional faint reference line
    o.append(f'<line x1="{X0 - 8}" x2="{X0 - 8}" y1="{yv(hi):.1f}" y2="{YB}" stroke="#111" stroke-width=".6"/>')
    v = lo
    while v <= hi + 1e-9:
        o.append(f'<line x1="{X0 - 12}" x2="{X0 - 8}" y1="{yv(v):.1f}" y2="{yv(v):.1f}" stroke="#111" stroke-width=".6"/>')
        o.append(f'<text x="{X0 - 16}" y="{yv(v) + 4:.1f}" font-size="12.5" text-anchor="end" fill="#555">{esc(fmt(Y["tick"], v))}</text>')
        v += Y["step"]
    if T.get("ylabel"):
        o.append(f'<text x="{X0 - 40}" y="{yv(hi) - 12:.1f}" font-size="13.5" fill="#333">{esc(T["ylabel"])}</text>')
    if Y.get("reference") is not None:
        r = yv(Y["reference"])
        o.append(f'<line x1="{X0}" x2="{X1}" y1="{r:.1f}" y2="{r:.1f}" stroke="#000" stroke-opacity=".25" stroke-width=".7" stroke-dasharray="4 3"/>')

    for s in series:
        # one polyline per run of known values: a missing period breaks the line
        runs, cur = [], []
        for i, val in enumerate(s["v"]):
            if val is None:
                if cur:
                    runs.append(cur)
                cur = []
            else:
                cur.append(i)
        if cur:
            runs.append(cur)
        for run in runs:
            if len(run) > 1:
                o.append(f'<polyline points="{" ".join(f"{xs[i]:.1f},{yv(s["v"][i]):.1f}" for i in run)}" fill="none" stroke="{s["color"]}" stroke-width="2"/>')
        for i, val in enumerate(s["v"]):
            if val is None:
                continue
            x, y = xs[i], yv(val)
            others = [t["v"][i] for t in series if t is not s and t["v"][i] is not None]
            dy = -11 if all(val >= m for m in others) else 20      # value above the higher line, below the lower
            if y + dy > YB - 4:
                dy = -11                                   # never on the x axis: a value at the floor goes above
            o.append(f'<circle cx="{x:.1f}" cy="{y:.1f}" r="4.5" fill="{PAPER}"/><circle cx="{x:.1f}" cy="{y:.1f}" r="2.8" fill="{s["color"]}"/>')
            o.append(f'<text x="{x:.1f}" y="{y + dy:.1f}" font-size="13" text-anchor="middle" fill="{s["color"]}" stroke="{PAPER}" stroke-width="5" paint-order="stroke">{esc(fmt(vfmt, val))}</text>')
        last = max((i for i, val in enumerate(s["v"]) if val is not None), default=None)
        if last is not None:
            o.append(f'<text x="{xs[last] + 16:.1f}" y="{yv(s["v"][last]) + 5:.1f}" font-size="16" fill="{s["color"]}">{esc(s["label"])}</text>')

    o.append(f'<line x1="{X0}" x2="{X1}" y1="{YB + 8}" y2="{YB + 8}" stroke="#111" stroke-width=".7"/>')
    for i, q in enumerate(P):
        l = ax.left[i]
        o.append(f'<line x1="{l:.1f}" x2="{l:.1f}" y1="{YB + 8}" y2="{YB + 14}" stroke="#111" stroke-width=".6"/>')
        o.append(f'<text x="{xs[i]:.1f}" y="{YB + 30}" font-size="13.5" text-anchor="middle" fill="#333">{esc((T.get("period_labels") or ax.labels)[i])}</text>')
        if any(s["n"] for s in series):
            counts = " · ".join(format(s["n"][i], ",") for s in series if s["n"])
            o.append(f'<text x="{xs[i]:.1f}" y="{YB + 50}" font-size="12" text-anchor="middle" fill="#999">{counts}</text>')
    o.append(f'<line x1="{ax.right[-1]:.1f}" x2="{ax.right[-1]:.1f}" y1="{YB + 8}" y2="{YB + 14}" stroke="#111" stroke-width=".6"/>')
    if T.get("count_label"):
        o.append(f'<text x="{X1 + 16:.1f}" y="{YB + 50}" font-size="12" fill="#999">{esc(T["count_label"])}</text>')
    # with no events and no quote, close up the empty row kept for them above the plot
    if not T.get("events") and not quote and not L.get("keep_event_row"):
        top = yv(hi) - 12 - 16
        if any(s["v"] and max((v for v in s["v"] if v is not None), default=lo) >= hi - (hi - lo) * 0.05 for s in series):
            top -= 12                                    # a value label may sit above the top tick
        d = int(top - (header_bottom + 36))
        if d > 0:
            o.insert(hdr_end, f'<g transform="translate(0,{-d})">')
            o.append("</g>")
            o[0] = o[0].replace(f'viewBox="0 0 {W} {H}" width="{W}" height="{H}"', f'viewBox="0 0 {W} {H - d}" width="{W}" height="{H - d}"')
            o[1] = o[1].replace(f'height="{H}"', f'height="{H - d}"')
            H -= d
    for k, line in enumerate(T.get("footnote", [])):
        o.append(f'<text x="{X0}" y="{H - 34 + 18 * 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}: " + "; ".join(f"{s['label']} " + ", ".join(fmt(vfmt, v) if v is not None else "-" for v in s["v"]) for s in series))


if __name__ == "__main__":
    main()