Tufte-style charts

"""Shared pieces for the Tufte-style charts: palette, period axis, monotone smoothing, text, page."""
import json
import os
import re
import string
from datetime import date, timedelta

PAPER, INK = "#fffff8", "#111"
RED, BLUE = "#a33a2c", "#4f7392"          # grew, shrank
GREYS = ["#ddd8cd", "#e7e3da"]            # steady, alternating so neighbours separate
QUIET = "#a8a39a"                         # event labels
FOOT = "#777"                             # footnote
MONTHS = "Jan Feb Mar Apr May Jun Jul Aug Sep Oct Nov Dec".split()

FONT_CSS = """@font-face{font-family:"et-book";src:url("https://cdn.jsdelivr.net/npm/tufte-css@1.8.0/et-book/et-book-roman-line-figures/et-book-roman-line-figures.woff") format("woff");font-style:normal}
@font-face{font-family:"et-book";src:url("https://cdn.jsdelivr.net/npm/tufte-css@1.8.0/et-book/et-book-display-italic-old-style-figures/et-book-display-italic-old-style-figures.woff") format("woff");font-style:italic}
@font-face{font-family:"et-book";src:url("https://cdn.jsdelivr.net/npm/tufte-css@1.8.0/et-book/et-book-semi-bold-old-style-figures/et-book-semi-bold-old-style-figures.woff") format("woff");font-weight:600}"""


def esc(s):
    return str(s).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")


# ---- the spec: one JSON file of wording, with the data inline or in a file beside it

def _read_json(path, what):
    try:
        with open(path, encoding="utf-8") as f:
            return json.load(f)
    except FileNotFoundError:
        raise SystemExit(f"{what} not found: {path}")
    except ValueError as e:
        raise SystemExit(f"{what} is not valid JSON: {path}: {e}")


def load_spec(path):
    spec = _read_json(path, "spec")
    d = spec.get("data")
    if isinstance(d, str):
        spec["data"] = _read_json(os.path.join(os.path.dirname(os.path.abspath(path)), d),
                                  f"data file named in {os.path.basename(path)} (paths are relative to the spec)")
    elif d is None:
        raise SystemExit(f"{path}: the spec needs \"data\" (an object, or a JSON file name beside the spec)")
    return spec


class _Fmt(string.Formatter):
    """str.format, but a negative number is written with a true minus sign (−), and a missing
    value (None or NaN, e.g. from an empty bin) is written as "–"."""

    def format_field(self, value, spec):
        if value is None or (isinstance(value, float) and value != value):
            return "\u2013"
        out = super().format_field(value, spec)
        if isinstance(value, (int, float)) and not isinstance(value, bool):
            out = out.replace("-", "−")
        return out


_FMT = _Fmt()


def filler(values):
    def fill(s):
        try:
            return _FMT.vformat(s, (), values)
        except (KeyError, IndexError, AttributeError, TypeError) as e:
            names = ", ".join(sorted(k for k in values if k != "data"))
            raise SystemExit(f"cannot fill {{{e.args[0] if e.args else ''}}} in {s!r}: this chart's placeholders are {names}, "
                             "plus {data[...]} for values in the data file")
    return fill


def fmt(template, *args):
    """template.format(*args), with true minus signs."""
    return _FMT.format(template, *args)


def num(x, spec=""):
    return _FMT.format("{:" + spec + "}", x)


# ---- periods: "2026Q1", "2026-03", "2026", "2026-03-15", or any other label (evenly spaced)

def _kind(p):
    p = str(p)
    if len(p) == 6 and p[:4].isdigit() and p[4] == "Q" and p[5] in "1234":
        return "quarter"
    if len(p) == 7 and p[:4].isdigit() and p[4] == "-" and p[5:].isdigit():
        return "month"
    if len(p) == 4 and p.isdigit():
        return "year"
    if len(p) == 10 and p[:4].isdigit() and p[4] == "-" and p[7] == "-":
        return "day"
    return "label"


def p_start(p):
    k = _kind(p)
    if k == "quarter":
        return date(int(p[:4]), 3 * int(p[5]) - 2, 1)
    if k == "month":
        return date(int(p[:4]), int(p[5:7]), 1)
    if k == "year":
        return date(int(p), 1, 1)
    return date.fromisoformat(p)


def p_end(p):
    """The first day after the period."""
    k, s = _kind(p), p_start(p)
    if k == "quarter":
        return date(s.year + 1, 1, 1) if s.month == 10 else date(s.year, s.month + 3, 1)
    if k == "month":
        return date(s.year + 1, 1, 1) if s.month == 12 else date(s.year, s.month + 1, 1)
    if k == "year":
        return date(s.year + 1, 1, 1)
    return s + timedelta(days=1)


def p_label(p):
    k = _kind(p)
    if k == "quarter":
        return f"{p[4:]} {p[:4]}"
    if k == "month":
        return f"{MONTHS[int(p[5:7]) - 1]} {p[:4]}"
    return str(p)


class Axis:
    """Maps periods to x. Dated periods sit on a true time scale, so a missing period leaves a gap
    and dated events can be drawn; other labels are evenly spaced bins."""

    def __init__(self, periods, x0, x1, pad_days=None):
        self.periods = list(periods)
        self.dated = all(_kind(p) != "label" for p in self.periods)
        if self.dated:
            k = _kind(self.periods[0])
            if pad_days is None:  # a little room after the last bin; 13 days for quarters
                pad_days = {"quarter": 13, "month": 4, "year": 51, "day": 0}.get(k, 0)
            t0 = p_start(self.periods[0])
            t1 = p_end(self.periods[-1]) + timedelta(days=pad_days)
            span = (t1 - t0).days
            self._x = lambda d: x0 + (x1 - x0) * (d - t0).days / span
            self.left = [self._x(p_start(p)) for p in self.periods]
            self.right = [self._x(p_end(p)) for p in self.periods]
            self.mid = [self._x(p_start(p) + (p_end(p) - p_start(p)) / 2) for p in self.periods]
        else:
            n = len(self.periods)
            w = (x1 - x0) / (n + 0.14)
            self._x = None
            self.left = [x0 + i * w for i in range(n)]
            self.right = [x0 + (i + 1) * w for i in range(n)]
            self.mid = [(l + r) / 2 for l, r in zip(self.left, self.right)]
        self.labels = [p_label(p) for p in self.periods]

    def date_x(self, d):
        if not self._x:
            raise SystemExit("dated events and quotes need dated periods (2026Q1, 2026-03, 2026)")
        return self._x(d)


def _monotone(pts):
    """Fritsch-Carlson slopes for points sorted by x: the curve never overshoots a measured value."""
    xs, ys = [p[0] for p in pts], [p[1] for p in pts]
    n = len(pts)
    d = [(ys[i + 1] - ys[i]) / (xs[i + 1] - xs[i]) for i in range(n - 1)]
    m = [d[0]] + [0.0 if d[i - 1] * d[i] <= 0 else (d[i - 1] + d[i]) / 2 for i in range(1, n - 1)] + [d[-1]]
    for i in range(n - 1):
        if d[i] == 0:
            m[i] = m[i + 1] = 0.0
        else:
            a, b = m[i] / d[i], m[i + 1] / d[i]
            h = a * a + b * b
            if h > 9:
                t = 3 / h ** 0.5
                m[i], m[i + 1] = t * a * d[i], t * b * d[i]
    segs = []
    for i in range(n - 1):
        h = xs[i + 1] - xs[i]
        segs.append(((xs[i], ys[i]), (xs[i] + h / 3, ys[i] + m[i] * h / 3),
                     (xs[i + 1] - h / 3, ys[i + 1] - m[i + 1] * h / 3), (xs[i + 1], ys[i + 1])))
    return segs


def smooth_path(pts, reverse=False, start="M"):
    """SVG path through pts (sorted by x). reverse draws it right to left; start='L' continues a path."""
    segs = _monotone(pts)
    if reverse:
        segs = [(p1, c2, c1, p0) for p0, c1, c2, p1 in segs[::-1]]
    f = lambda p: f"{p[0]:.1f},{p[1]:.1f}"
    return f"{start}{f(segs[0][0])} " + " ".join(f"C{f(c1)} {f(c2)} {f(p1)}" for _, c1, c2, p1 in segs)


def events(svg, tx, items, row_y, y_bottom, special=None):
    """Faint dotted hairlines for dated events, labels packed into up to three rows above row_y."""
    special = special or {}
    last = [-1e9] * 3
    for d, name in items:
        x = tx(d)
        svg.append(f'<line x1="{x:.1f}" x2="{x:.1f}" y1="{row_y + 6}" y2="{y_bottom}" stroke="#000" stroke-opacity=".28" stroke-width=".7" stroke-dasharray="1.5 3"/>')
        half = 3.4 * len(name) + 4
        row = next((r for r in range(3) if x - half > last[r]), 2)
        last[row] = x + half
        anchor, dx = special.get(name, ("middle", 0))
        svg.append(f'<text x="{x + dx:.1f}" y="{row_y - row * 15}" font-size="12.5" text-anchor="{anchor}" fill="{QUIET}">{esc(name)}</text>')


def parse_events(spec):
    return [(date.fromisoformat(e["date"]), e["label"]) for e in spec.get("events", [])]


def text_width(s, size):
    """A rough width in px for ET Book / Palatino: digits and capitals are wider than the rest."""
    w = 0.0
    for ch in str(s):
        w += 0.55 if ch.isdigit() or ch.isupper() or ch in "%→—" else 0.25 if ch in " .,:;'’|il!" else 0.45
    return w * size


def spread(ys, gap, lo, hi):
    """Push label positions apart by at least gap, keeping them inside [lo, hi] where possible.
    Returns positions in the input order."""
    order = sorted(range(len(ys)), key=lambda i: ys[i])
    pos = [ys[i] for i in order]
    for i in range(1, len(pos)):
        pos[i] = max(pos[i], pos[i - 1] + gap)
    over = pos[-1] - hi if pos else 0
    if over > 0:
        pos = [p - over for p in pos]
        for i in range(len(pos) - 2, -1, -1):
            pos[i] = min(pos[i], pos[i + 1] - gap)
    out = [0.0] * len(ys)
    for i, p in zip(order, pos):
        out[i] = p
    return out


# ---- light and dark
# The scripts draw with the classic palette below. With the default transparent background,
# write() turns each of these colours into a CSS custom property whose light and dark values
# are set in a <style> inside the SVG, switched by prefers-color-scheme. The light values darken
# the faint greys so text meets WCAG AA (4.5:1) on #fffff8 and white; the dark values meet it on
# #111. With background "paper" nothing is rewritten: the page is byte-for-byte the classic one.
#            classic     name      light      dark
THEME = [("#fffff8", "gap", "#fffff8", "#111111"),   # band separators, halos behind labels
         ("#111", "ink", "#111", "#ecebe4"),
         ("#222", "ink2", "#222", "#e2e0d8"),
         ("#333", "ink3", "#333", "#d6d3ca"),
         ("#444", "ink4", "#444", "#cac7bd"),
         ("#555", "ink5", "#555", "#bab7ad"),
         ("#777", "foot", "#666666", "#a9a69d"),
         ("#888", "mute", "#6b6b6b", "#a3a097"),
         ("#999", "faint", "#737373", "#9c998f"),
         ("#a8a39a", "quiet", "#706b62", "#a39e94"),
         ("#8a8478", "grey", "#6e6859", "#a8a294"),
         ("#000", "hair", "#000", "#fff"),             # always drawn at 25-28% opacity
         ("#a33a2c", "red", "#a33a2c", "#e2735f"),
         ("#4f7392", "blue", "#4f7392", "#86acd0"),
         ("#ddd8cd", "steady1", "#ddd8cd", "#3a3833"),
         ("#e7e3da", "steady2", "#e7e3da", "#46433d")]
TOKEN = {c: n for c, n, _, _ in THEME}
NAMED_COLORS = {"ink": INK, "red": RED, "blue": BLUE, "grey": "#8a8478", "gray": "#8a8478"}
BACKGROUNDS = ("transparent", "paper")


def color(c):
    """A spec colour: a palette name (ink, red, blue, grey) or any CSS colour, kept as given."""
    return NAMED_COLORS.get(c, c)


def add_background_arg(ap):
    ap.add_argument("--paper", action="store_true",
                    help="classic fixed look: #fffff8 paper, one palette, no dark mode (same as spec \"background\": \"paper\")")


def background(T, args):
    b = "paper" if getattr(args, "paper", False) else T.get("background", "transparent")
    if b not in BACKGROUNDS:
        raise SystemExit(f"spec \"background\" must be one of {', '.join(BACKGROUNDS)}, not {b!r}")
    return b


def _vars(i):
    return ";".join(f"--tc-{t[1]}:{t[i]}" for t in THEME)


THEME_CSS = (".tc-chart{--tc-bg:transparent;" + _vars(2) + "}"
             "@media (prefers-color-scheme:dark){.tc-chart:not([data-theme=light]){" + _vars(3) + "}}"
             ".tc-chart[data-theme=dark]{" + _vars(3) + "}")
_ATTR = re.compile(r'\s(fill|stroke)="(#[0-9a-fA-F]{3,6})"')


def _theme_tag(m):
    tag = m.group(0)
    props = []

    def drop(a):
        c = a.group(2).lower()
        if c not in TOKEN:
            return a.group(0)                             # a colour outside the palette stays as given
        name = "bg" if c == PAPER and tag.startswith("<rect") and a.group(1) == "fill" else TOKEN[c]
        props.append(f"{a.group(1)}:var(--tc-{name})")
        return ""
    tag = _ATTR.sub(drop, tag)
    if not props:
        return tag
    end = len(tag) - (2 if tag.endswith("/>") else 1)
    return tag[:end] + f' style="{";".join(props)}"' + tag[end:]


def themed(svg):
    """The classic SVG with every palette colour as a light/dark custom property."""
    svg = re.sub(r"<[a-zA-Z][^>]*>", _theme_tag, svg)
    svg = svg.replace('<svg id="fig"', '<svg id="fig" class="tc-chart"', 1)
    i = svg.index(">", svg.index("<svg")) + 1
    return svg[:i] + f"<style>{THEME_CSS}</style>" + svg[i:]


def page(svg_parts, width, bg="paper"):
    if bg == "paper":
        body = f"body{{margin:0;background:{PAPER}}}"
        svg = "".join(svg_parts)
    else:  # the page follows the reader's scheme; the chart itself stays transparent
        body = "body{margin:0;background:#fffff8}@media (prefers-color-scheme:dark){body{background:#111}}"
        svg = themed("".join(svg_parts))
    return ("<!doctype html><html><head><meta charset=\"utf-8\"><style>\n" + FONT_CSS +
            f"\n{body} svg{{width:{width}px;display:block;font-family:et-book,Palatino,Georgia,serif}}\n"
            "</style></head><body>" + svg + "</body></html>")


def write(out, svg_out, parts, width, svg_font=True, bg="paper"):
    open(out, "w").write(page(parts, width, bg))
    if svg_out:
        s = "".join(parts)
        if bg != "paper":
            s = themed(s)
        if svg_font:
            s = s.replace('<svg id="fig"', '<svg id="fig" font-family="Palatino, Georgia, serif"', 1)
        head = s[:s.index(">")]
        if bg != "paper" and " width=" not in head:       # an <img> needs the size to lay it out
            w, h = re.search(r'viewBox="0 0 ([\d.]+) ([\d.]+)"', head).groups()
            s = s.replace("<svg ", f'<svg width="{w}" height="{h}" ', 1)
        open(svg_out, "w").write(s)