#!/usr/bin/env python3
"""
04_bubble_map.py — the global bubble map.

One bubble per testable country (power >= 0.8 to detect a medium departure
from uniformity).  Bubble AREA and REDNESS both grow with -log10(p) of the
chi-square test of the Moon's nakṣatra distribution against the
time-weighted uniform null.  No multiple-comparison correction is drawn;
the caption says how many bubbles chance alone would colour.

Two layers are written into one SVG — 27 and 28 asterisms — so the essay
page can toggle them; a standalone viewer sees the first. Every recorded
earthquake counts: the essay does not decluster.

Input : country_results.csv (from 03), countries-110m.json (outlines)
Output: bubble_map.svg
"""
import argparse
import json
import math

import numpy as np
import pandas as pd
from pyproj import Transformer

W, H = 1400, 720
BG = "#0b0e18"
LAND = "#1b2030"
LAND_EDGE = "#39415a"
GOLD = "#d4af37"
CREAM = "#f5f0e8"

_robin = Transformer.from_crs("EPSG:4326", "+proj=robin +lon_0=0", always_xy=True)
# Robinson extents (metres) for the world
_XMAX = 17005833.33
_YMAX = 8625154.47


def proj(lon, lat):
    x, y = _robin.transform(lon, lat)
    return (x / _XMAX * (W / 2 - 20) + W / 2,
            -y / _YMAX * (H / 2 - 20) + H / 2)


def land_paths(topo_path):
    import topojson
    from shapely.geometry import shape
    topo = json.load(open(topo_path))
    gj = json.loads(topojson.Topology(topo, object_name="countries").to_geojson())
    d = []
    for f in gj["features"]:
        g = shape(f["geometry"])
        polys = list(g.geoms) if g.geom_type == "MultiPolygon" else [g]
        for p in polys:
            for ring in [p.exterior] + list(p.interiors):
                pts = [proj(x, y) for x, y in ring.coords]
                d.append("M" + " L".join(f"{x:.1f} {y:.1f}" for x, y in pts) + "Z")
    return " ".join(d)


def colour(neglogp):
    """0 (p=1) -> muted blue-grey; 1.3 (p=.05) -> gold; >=3 -> deep red."""
    stops = [(0.0, (95, 110, 140)), (1.3, (212, 175, 55)),
             (2.0, (230, 120, 60)), (3.0, (214, 40, 40)), (6.0, (160, 10, 30))]
    v = max(0.0, min(neglogp, 6.0))
    for (a, ca), (b, cb) in zip(stops, stops[1:]):
        if a <= v <= b:
            t = (v - a) / (b - a)
            return "#%02x%02x%02x" % tuple(int(ca[i] + t * (cb[i] - ca[i])) for i in range(3))
    return "#a00a1e"


def radius(neglogp):
    return 3.5 + 7.5 * math.sqrt(min(neglogp, 8.0))


def layer(tab, tag, k, visible):
    t = tab[(tab.catalog == tag) & tab[f"testable{k}"]].copy()
    t["nlp"] = -np.log10(t[f"p{k}"].clip(lower=1e-12))
    t = t.sort_values("nlp")   # small first so big ones draw on top... then reverse
    hide = "" if visible else ' display="none"'
    out = [f'<g class="layer" data-layer="{tag}{k}"{hide}>']
    for _, r in t.iloc[::-1].iterrows():
        x, y = proj(r.lon, r.lat)
        rad = radius(r.nlp)
        ring = False
        p = r[f"p{k}"]
        ptxt = f"p-value {p:.3g}" if p >= 1e-4 else f"p-value {p:.1e}"
        title = (f"{r.country}\n{int(r.n)} earthquakes "
                 f"(every recorded event, "
                 f"{k} nakṣatras)\nuniformity: χ² = {r[f'chi2_{k}']:.1f}, {ptxt}"
                 f"\nstatistical power {r[f'power{k}']:.2f} (w = 0.3) — above the 0.8 gate")
        out.append(
            f'<circle cx="{x:.1f}" cy="{y:.1f}" r="{rad:.1f}" fill="{colour(r.nlp)}" '
            f'fill-opacity="0.78" stroke="{GOLD if ring else CREAM}" '
            f'stroke-width="{2.2 if ring else 0.6}" stroke-opacity="{1 if ring else 0.5}" '
            f'data-country="{r.country}" data-p="{p:.4g}"><title>{title}</title></circle>')
        if rad > 14:
            out.append(f'<text x="{x:.1f}" y="{y + 3.5:.1f}" text-anchor="middle" '
                       f'font-size="{max(9, min(13, rad * 0.75)):.0f}" fill="{CREAM}" '
                       f'font-family="Lato, system-ui, sans-serif" pointer-events="none">'
                       f'{short(r.country)}</text>')
    out.append("</g>")
    return "\n".join(out)


def short(name):
    return {"United States of America": "USA", "Papua New Guinea": "PNG",
            "Solomon Is.": "Solomon Is.", "Dominican Rep.": "Dom. Rep.",
            "New Zealand": "NZ", "Philippines": "Philippines"}.get(name, name)


def legend():
    xs = 60
    y = H - 48
    items = [(0.0, "p-value ≥ 0.5"), (1.3, "p-value 0.05"), (2.0, "p-value 0.01"),
             (3.0, "p-value 0.001"), (5.0, "p-value 10⁻⁵")]
    out = [f'<g font-family="Lato, system-ui, sans-serif" font-size="12" fill="{CREAM}" fill-opacity="0.85">']
    x = xs
    for v, lab in items:
        r = radius(v)
        out.append(f'<circle cx="{x + r:.0f}" cy="{y}" r="{r:.1f}" fill="{colour(v)}" fill-opacity="0.85" stroke="{CREAM}" stroke-width="0.6" stroke-opacity="0.5"/>')
        out.append(f'<text x="{x + 2 * r + 6:.0f}" y="{y + 4}">{lab}</text>')
        x += 2 * r + 6 + 7.2 * len(lab) + 10
    out.append(f'<text x="{x + 6}" y="{y + 4}" fill="{GOLD}">every bubble: statistical power ≥ 0.8 (n ≥ 258 events); countries below the gate are not drawn</text>')
    out.append("</g>")
    return "\n".join(out)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--results", default="country_results.csv")
    ap.add_argument("--outlines", default="countries-110m.json")
    ap.add_argument("--out", default="bubble_map.svg")
    a = ap.parse_args()
    tab = pd.read_csv(a.results)

    # graticule
    grat = []
    for lat in range(-60, 61, 30):
        grat.append("M" + " L".join(f"{x:.1f} {y:.1f}" for x, y in (proj(l, lat) for l in range(-180, 181, 5))))
    for lon in range(-180, 181, 30):
        grat.append("M" + " L".join(f"{x:.1f} {y:.1f}" for x, y in (proj(lon, l) for l in range(-90, 91, 5))))
    outline = "M" + " L".join(f"{x:.1f} {y:.1f}" for x, y in
                              [proj(-180, l) for l in range(-90, 91, 3)] +
                              [proj(180, l) for l in range(90, -91, -3)]) + "Z"

    svg = [f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {W} {H}" width="{W}" height="{H}" '
           f'role="img" aria-label="Bubble map: p-values of the nakṣatra-uniformity test, one bubble per country with statistical power of at least 0.8">',
           f'<rect width="{W}" height="{H}" fill="{BG}"/>',
           f'<path d="{outline}" fill="#0f1424" stroke="{LAND_EDGE}" stroke-width="0.8"/>',
           f'<path d="{" ".join(grat)}" fill="none" stroke="{LAND_EDGE}" stroke-width="0.4" stroke-opacity="0.5"/>',
           f'<path d="{land_paths(a.outlines)}" fill="{LAND}" stroke="{LAND_EDGE}" stroke-width="0.5" fill-rule="evenodd"/>']
    first = True
    for k in ("27", "28"):            # every recorded shaking counts; no declustered layer
        svg.append(layer(tab, "raw", k, first))
        first = False
    svg.append(legend())
    svg.append(f'<text x="{W - 24}" y="{H - 20}" text-anchor="end" font-size="11" fill="{CREAM}" fill-opacity="0.55" '
               f'font-family="Lato, system-ui, sans-serif">USGS ComCat 1980–2025, M ≥ 5.0 · Moon sidereal (Lahiri) · bubble = p-value of the uniformity test, drawn only where statistical power ≥ 0.8 · ayurastro.com</text>')
    svg.append("</svg>")
    open(a.out, "w", encoding="utf-8").write("\n".join(svg))
    print("wrote", a.out)


if __name__ == "__main__":
    main()