#!/usr/bin/env python3
"""
Render the post's charts from out/analysis.json into out/charts/*.png
(plus the .html each is screenshotted from).
"""
import json
import shutil
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "out"
CHARTS = OUT / "charts"
CHROME = next(filter(None, [shutil.which("google-chrome"), shutil.which("chromium"),
                             "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"]))

WIDTH = 640
INK, SECONDARY, MUTED, GRID, EMPTY = "#0b0b0b", "#52514e", "#8a8984", "#e4e3df", "#dcdbd6"
REPUBLICAN, DEMOCRAT = "#e34948", "#2a78d6"
AMENDMENTS = {"q1": ("Abortion", "#6250d6"), "q2": ("Same-sex marriage", "#eb6834"), "q3": ("Felon voting rights", "#1baf7a")}


def one_in(rate):
    return "almost none" if rate < 0.01 else f"1 in {round(1 / rate)}"


def pct_and_one_in(rate):
    return f"{100 * rate:.0f}%" + ("" if rate < 0.01 else f" · 1 in {round(1 / rate)}")


def text(x, y, s, size=14, weight=400, color=INK, anchor="start"):
    return f'<text x="{x}" y="{y}" font-size="{size}" font-weight="{weight}" fill="{color}" text-anchor="{anchor}">{s}</text>'


def line(x1, y1, x2, y2, color, width=1, dash=""):
    extra = f' stroke-dasharray="{dash}"' if dash else ""
    return f'<line x1="{x1}" y1="{y1}" x2="{x2}" y2="{y2}" stroke="{color}" stroke-width="{width}" stroke-linecap="round"{extra}/>'


def dot(x, y, color, r=6):
    return f'<circle cx="{x}" cy="{y}" r="{r}" fill="{color}" stroke="#fff" stroke-width="2"/>'


def title(s, sub=None):
    parts = [text(0, 24, s, 22, 700)]
    if sub:
        parts.append(text(0, 50, sub, 16, color=SECONDARY))
    return parts


def footnote(height, s):
    return text(0, height - 6, s, 14, color=SECONDARY)


def expected_vs_actual(a):
    height = 320
    left, right = 230, 420
    expected = ["q3", "q2", "q1"]
    notes = {"q3": "a no-brainer", "q2": "it's 2026", "q1": "squeak by, or fail"}
    finished = sorted(AMENDMENTS, key=lambda q: -a["totals"]["all"][q])
    y = lambda rank: 110 + rank * 70
    parts = title("I got it backwards")
    parts += [text(left, 80, "What I expected", 16, 600, SECONDARY, "end"), text(right, 80, "How they finished", 16, 600, SECONDARY)]
    for q, (name, color) in AMENDMENTS.items():
        y1, y2 = y(expected.index(q)), y(finished.index(q))
        parts.append(line(left + 12, y1, right - 12, y2, color, 3))
        parts += [dot(left + 12, y1, color), dot(right - 12, y2, color)]
        parts += [text(left - 4, y1 + 1, name, 17, 600, anchor="end"), text(left - 4, y1 + 21, notes[q], 15, color=SECONDARY, anchor="end")]
        parts.append(text(right + 4, y2 + 6, name, 17, 600))
    parts.append(footnote(height, "All three passed. Order by statewide yes share."))
    return parts, height


def dot_grid(x, y, filled, color, spacing=11):
    return [f'<circle cx="{x + (k % 10) * spacing + 4}" cy="{y + (k // 10) * spacing + 4}" r="4" '
            f'fill="{color if k < filled else EMPTY}"/>' for k in range(100)]


def crossovers(a):
    est = a["estimates"]["all"]
    row, top = 130, 104
    height = top + 3 * row + 30
    columns = [("r_yes", "Republicans who voted yes", REPUBLICAN, 190), ("d_no", "Democrats who voted no", DEMOCRAT, 440)]
    parts = title("Who crossed party lines (estimated)", "Each grid is 100 voters")
    for _, label, _, x in columns:
        parts.append(text(x, 88, label, 15, 600, SECONDARY))
    for i, (q, (name, _)) in enumerate(AMENDMENTS.items()):
        y = top + i * row
        parts.append(text(0, y + 58, name, 17, 600))
        for key, _, color, x in columns:
            rate = max(0.0, est[q][key])
            parts += dot_grid(x, y, round(100 / round(1 / rate)) if rate >= 0.01 else 0, color)
            parts.append(text(x + 122, y + 62, one_in(rate), 19, 700))
    parts.append(footnote(height, "Republicans = Mizusawa voters, Democrats = Warner voters (US Senate race)."))
    return parts, height


def democrats_by_black_share(a):
    groups = a["by_black_share"]
    height, top, bottom, left, right = 400, 80, 300, 48, 380
    ymax = 30
    X = lambda i: left + 30 + i * (right - left - 60) / (len(groups) - 1)
    Y = lambda v: bottom - v / ymax * (bottom - top)
    parts = title("Democrats who voted no (estimated)", "By the Black share of adults in the precinct")
    for v in range(0, ymax + 1, 10):
        parts += [line(left, Y(v), right, Y(v), GRID), text(left - 8, Y(v) + 5, f"{v}%", 14, color=SECONDARY, anchor="end")]
    for i, g in enumerate(groups):
        parts.append(text(X(i), bottom + 24, g["label"], 14, color=SECONDARY, anchor="middle"))
    labels = []
    for k, (q, (name, color)) in enumerate(AMENDMENTS.items()):
        dodge = (k - 1) * 5
        pts = [(X(i) + dodge, max(0.0, g["estimates"][q]["d_no"])) for i, g in enumerate(groups)]
        parts.append(f'<path d="{" ".join(("M" if i == 0 else "L") + f"{x},{Y(100 * v)}" for i, (x, v) in enumerate(pts))}" '
                     f'fill="none" stroke="{color}" stroke-width="3" stroke-linejoin="round"/>')
        parts += [dot(x, Y(100 * v), color, 5) for x, v in pts]
        labels.append([Y(100 * pts[-1][1]) - 4, name, f"{one_in(pts[0][1])} → {one_in(pts[-1][1])}", color])
    labels.sort()
    for i in range(1, len(labels)):
        labels[i][0] = max(labels[i][0], labels[i - 1][0] + 40)
    for yy, name, change, color in labels:
        parts += [text(right + 16, yy, name, 15, 700, color), text(right + 16, yy + 18, change, 14, color=SECONDARY)]
    parts.append(footnote(height, "Precincts matched to the 2020 Census (about 3 in 4). Describes precincts, not voters."))
    return parts, height


def republicans_by_income(a):
    groups = a["by_income_split"]
    shades = ["#f2a3a2", REPUBLICAN]
    names = ["Under $80k", "$80k or more"]
    row, top, left = 84, 104, 200
    height = top + 3 * row + 30
    X = lambda v: left + v / 20 * (WIDTH - left - 130)
    parts = title("Republicans who voted yes (estimated)", "By the median household income of the county or city")
    for j, label in enumerate(names):
        x = 0 if j == 0 else 160
        parts += [f'<rect x="{x}" y="66" width="14" height="14" rx="3" fill="{shades[j]}"/>', text(x + 20, 79, label, 15, color=SECONDARY)]
    for i, (q, (name, _)) in enumerate(AMENDMENTS.items()):
        y = top + i * row
        parts.append(text(0, y + 32, name, 17, 600))
        for j, g in enumerate(groups):
            rate = max(0.0, g["estimates"][q]["r_yes"])
            by = y + j * 28
            parts.append(f'<rect x="{left}" y="{by}" width="{max(2, X(100 * rate) - left)}" height="22" rx="4" fill="{shades[j]}"/>')
            parts.append(text(X(100 * rate) + 8, by + 17, pct_and_one_in(rate), 15, 600))
    parts.append(footnote(height, "Median household income, Census ACS 2019–2023."))
    return parts, height


def render(name, parts, height):
    svg = (f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="-12 -8 {WIDTH + 24} {height + 16}" width="{WIDTH + 24}" '
           f'height="{height + 16}" font-family="system-ui,-apple-system,sans-serif">{"".join(parts)}</svg>')
    html = CHARTS / f"{name}.html"
    html.write_text(f'<!DOCTYPE html><html><body style="margin:0;background:#fff">{svg}</body></html>')
    png = CHARTS / f"{name}.png"
    subprocess.run([CHROME, "--headless=new", "--disable-gpu", "--hide-scrollbars", "--force-device-scale-factor=2",
                    f"--window-size={WIDTH + 24},{height + 16}", f"--screenshot={png}", html.as_uri()],
                   check=True, stderr=subprocess.DEVNULL)
    print(png)


def main():
    a = json.load(open(OUT / "analysis.json"))
    CHARTS.mkdir(parents=True, exist_ok=True)
    for name, build in [("expected-vs-actual", expected_vs_actual), ("crossovers", crossovers),
                        ("democrats-by-black-share", democrats_by_black_share),
                        ("republicans-by-income", republicans_by_income)]:
        render(name, *build(a))


if __name__ == "__main__":
    sys.exit(main())