import numpy as np
from PIL import Image, ImageDraw, ImageFont
from scipy.ndimage import gaussian_filter
 
OUT = sys.argv[1] if len(sys.argv) > 1 else "Essais_nucleaires_1945-2017_animation.gif"
 
RAW = {
    "United States": "1945:1 1946:2 1948:3 1951:16 1952:10 1953:11 1954:6 1955:18 1956:18 1957:32 1958:77 1961:10 "
                     "1962:96 1963:47 1964:45 1965:38 1966:48 1967:42 1968:56 1969:46 1970:39 1971:24 1972:27 "
                     "1973:24 1974:22 1975:22 1976:20 1977:20 1978:19 1979:15 1980:14 1981:16 1982:18 1983:18 "
                     "1984:18 1985:17 1986:14 1987:14 1988:15 1989:11 1990:8 1991:7 1992:6",
    "Russia": "1949:1 1951:2 1953:5 1954:10 1955:6 1956:9 1957:16 1958:34 1961:59 1962:79 1964:9 1965:14 1966:18 "
              "1967:17 1968:17 1969:19 1970:16 1971:23 1972:24 1973:17 1974:21 1975:19 1976:21 1977:24 1978:31 "
              "1979:31 1980:24 1981:21 1982:19 1983:25 1984:27 1985:10 1987:23 1988:16 1989:7 1990:1",
    "United Kingdom": "1952:1 1953:2 1956:6 1957:7 1958:5 1962:2 1964:2 1965:1 1974:1 1976:1 1978:2 1979:1 1980:3 "
                      "1981:1 1982:1 1983:1 1984:2 1985:1 1986:1 1987:1 1989:1 1990:1 1991:1",
    "France": "1960:3 1961:2 1962:1 1963:3 1964:3 1965:4 1966:7 1967:3 1968:5 1970:8 1971:5 1972:4 1973:6 1974:9 "
              "1975:2 1976:5 1977:9 1978:11 1979:10 1980:12 1981:12 1982:10 1983:9 1984:8 1985:8 1986:8 1987:8 "
              "1988:8 1989:9 1990:6 1991:6 1995:5 1996:1",
    "China": "1964:1 1965:1 1966:3 1967:2 1968:1 1969:2 1970:1 1971:1 1972:2 1973:1 1974:1 1975:1 1976:4 1977:1 "
             "1978:3 1979:1 1980:1 1982:1 1983:2 1984:2 1987:1 1988:1 1990:2 1992:2 1993:1 1994:2 1995:2 1996:2",
    "India": "1974:1 1998:2",
    "Pakistan": "1998:2",
    "North Korea": "2006:1 2009:1 2013:1 2016:2 2017:1",
}
COUNTRIES = [  # key, French label, colour
    ("United States", "États-Unis", (76, 141, 255)),
    ("Russia", "URSS", (255, 77, 77)),
    ("United Kingdom", "Royaume-Uni", (181, 124, 255)),
    ("France", "France", (255, 200, 69)),
    ("China", "Chine", (255, 138, 61)),
    ("India", "Inde", (43, 217, 230)),
    ("Pakistan", "Pakistan", (124, 224, 124)),
    ("North Korea", "Corée du Nord", (255, 111, 181)),
]
Y0, Y1 = 1945, 2017
DATA = {k: {int(a): int(b) for a, b in (p.split(":") for p in v.split())} for k, v in RAW.items()}
TOTAL = sum(sum(d.values()) for d in DATA.values())
 
EVENTS = {  # year: (text, emphasis)
    1945: ("Trinity : premier essai nucléaire", 1.4),
    1949: ("Premier essai soviétique", 1.0),
    1952: ("Premier essai britannique · Ivy Mike, première bombe H", 1.2),
    1959: ("Moratoire de 1959-1960", 1.0),
    1960: ("Premier essai français (Gerboise bleue)", 1.0),
    1961: ("Tsar Bomba : 50 mégatonnes", 1.2),
    1962: ("Record : 178 essais en un an", 1.4),
    1963: ("Traité d'interdiction partielle des essais", 1.3),
    1964: ("Premier essai chinois", 1.0),
    1974: ("Premier essai indien", 1.0),
    1990: ("Dernier essai soviétique", 1.0),
    1996: ("Traité d'interdiction complète des essais", 1.3),
    1998: ("Essais indiens et pakistanais", 1.0),
    2006: ("Premier essai nord-coréen", 1.0),
    2017: ("Dernier essai recensé (Corée du Nord)", 1.2),
}
 
# ----------------------------------------------------------------- layout
W, H, S = 420, 270, 4
PX0, PBASE = 34, 254                       # plot left, baseline
COLW, SQ = 5, 2                            # px per year, square pitch (2 per row)
GAP = 0.25                                 # thin gap so single tests stay visible
BG = (7, 9, 15)
GRID, AXIS, TXT, TXT2, TXT3 = (26, 32, 46), (58, 66, 88), (236, 240, 248), (160, 170, 190), (100, 110, 132)
F = {"bold": "/usr/share/fonts/opentype/inter/Inter-Bold.otf",
     "semi": "/usr/share/fonts/opentype/inter/Inter-SemiBold.otf",
     "med": "/usr/share/fonts/opentype/inter/Inter-Medium.otf",
     "reg": "/usr/share/fonts/opentype/inter/Inter-Regular.otf"}
 
 
def font(k, px):
    return ImageFont.truetype(F[k], round(px * S))
 
 
def year_squares(y):
    """List of (x, y_top, colour) squares for year y, US at the bottom."""
    out = []
    i = 0
    x_col = PX0 + (y - Y0) * COLW
    for key, _, col in COUNTRIES:
        for _ in range(DATA[key].get(y, 0)):
            row, c = divmod(i, 2)
            out.append((x_col + c * SQ, PBASE - (row + 1) * SQ, col))
            i += 1
    return out
 
 
SQUARES = {y: year_squares(y) for y in range(Y0, Y1 + 1)}
 
 
def column_top(y):
    n = sum(DATA[k].get(y, 0) for k in DATA)
    return PBASE - ((n + 1) // 2) * SQ
 
 
def render(year, phase, event_alpha=None, final=False, fade=1.0):
    """year: current year; phase 0 = falling, 1 = landed (glow), 2 = settled."""
    base = Image.new("RGB", (W * S, H * S), BG)
    d = ImageDraw.Draw(base)
    glow = Image.new("RGB", (W, H), (0, 0, 0))
    gd = ImageDraw.Draw(glow)
 
    # grid: 50 / 100 / 150 tests per year (1 px per test)
    for n in (50, 100, 150):
        yv = PBASE - n
        d.line([(PX0 * S, yv * S), ((PX0 + (Y1 - Y0 + 1) * COLW) * S, yv * S)], fill=GRID, width=S)
        d.text(((PX0 - 5) * S, yv * S), str(n), font=font("reg", 8), fill=TXT3, anchor="rm")
    d.line([(PX0 * S, PBASE * S), ((PX0 + (Y1 - Y0 + 1) * COLW) * S, PBASE * S)], fill=AXIS, width=S)
    for yr in range(1950, Y1 + 1, 10):
        x = PX0 + (yr - Y0) * COLW + 2
        d.line([(x * S, PBASE * S), (x * S, (PBASE + 3) * S)], fill=AXIS, width=S)
        d.text((x * S, (PBASE + 4) * S), str(yr), font=font("reg", 8.5), fill=TXT3, anchor="ma")
    d.text(((PX0 - 5) * S, (PBASE - 165) * S), "essais / an", font=font("reg", 8), fill=TXT3, anchor="ls")
 
    def fcol(c, a=1.0):
        return tuple(int(BG[i] + (c[i] - BG[i]) * a * fade) for i in range(3))
 
    last = Y1 if final else year
    for y in range(Y0, last + 1):
        for (x, yt, col) in SQUARES[y]:
            if y == year and not final and phase == 0:
                continue
            bright = 1.0
            if y == year and not final and phase == 1:
                c2 = tuple(min(255, int(v * 0.55 + 255 * 0.45)) for v in col)
                d.rectangle([x * S, yt * S, (x + SQ - GAP) * S - 1, (yt + SQ - GAP) * S - 1], fill=fcol(c2))
                gd.rectangle([x - 1, yt - 1, x + SQ, yt + SQ], fill=col)
            else:
                d.rectangle([x * S, yt * S, (x + SQ - GAP) * S - 1, (yt + SQ - GAP) * S - 1], fill=fcol(col, bright))
    if not final and phase == 0:                       # falling: squares above their slot, with streaks
        for (x, yt, col) in SQUARES[year]:
            off = 26
            d.rectangle([x * S, (yt - off) * S, (x + SQ) * S - 1, (yt - off + SQ) * S - 1], fill=col)
            for k in range(1, 10):
                a = (1 - k / 10) * 0.55
                yy = yt - off - k
                d.rectangle([x * S, yy * S, (x + SQ) * S - 1, (yy + 1) * S - 1], fill=fcol(col, a))
            gd.rectangle([x - 1, yt - off - 1, x + SQ, yt - off + SQ], fill=col)
 
    img = np.asarray(base.reduce(S)).astype(np.float32)
    g = np.asarray(glow).astype(np.float32)
    if g.any():
        g = gaussian_filter(g, (2.2, 2.2, 0)) * 1.1 + gaussian_filter(g, (6, 6, 0)) * 0.6
        img = 255 - (255 - img) * (1 - np.clip(g / 255, 0, 1))       # screen blend
    im = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
    t = ImageDraw.Draw(im)
 
    # header: title, year, legend with running totals
    t.text((12, 19), "Essais nucléaires dans le monde", font=ImageFont.truetype(F["semi"], 13), fill=TXT, anchor="ls")
    t.text((12, 32), "chaque carré représente un essai", font=ImageFont.truetype(F["reg"], 9), fill=TXT2, anchor="ls")
    t.text((W - 12, 38), str(Y1 if final else year), font=ImageFont.truetype(F["bold"], 30), fill=TXT, anchor="rs")
    upto = Y1 if final else year
    fl = ImageFont.truetype(F["med"], 8.5)
    fn = ImageFont.truetype(F["semi"], 8.5)
    for i, (key, lab, col) in enumerate(COUNTRIES):
        n = sum(v for yr, v in DATA[key].items() if yr <= upto)
        cx = 12 + (i % 4) * 101
        cy = 52 + (i // 4) * 13
        active = n > 0
        a = 1.0 if active else 0.35
        t.rectangle([cx, cy - 7, cx + 6, cy - 1], fill=tuple(int(BG[j] + (col[j] - BG[j]) * a) for j in range(3)))
        t.text((cx + 10, cy), lab, font=fl, fill=TXT2 if active else TXT3, anchor="ls")
        t.text((cx + 86, cy), f"{n:,}".replace(",", " "), font=fn, fill=TXT if active else TXT3, anchor="rs")
 
    # milestone annotation
    if event_alpha:
        ey, a = event_alpha
        text, _ = EVENTS[ey]
        fe = ImageFont.truetype(F["semi"], 9.5)
        wtxt = t.textlength(text, font=fe)
        xc = PX0 + (ey - Y0) * COLW + 2
        top = column_top(ey)
        by = 96 if ey not in (1962,) else 104
        bx = xc + 10 if xc + 10 + wtxt + 14 < W - 6 else xc - 10 - wtxt - 14
        if ey == 1962:
            bx = xc + 12
        col_line = tuple(int(BG[j] + (v - BG[j]) * a) for j, v in enumerate((210, 216, 230)))
        t.line([(xc, by + 8), (xc, max(top - 3, by + 8))], fill=tuple(int(BG[j] + (90 - BG[j]) * a) for j in range(3)),
               width=1)
        t.line([(xc, by + 8), (bx if bx > xc else bx + wtxt + 14, by + 8)],
               fill=tuple(int(BG[j] + (90 - BG[j]) * a) for j in range(3)), width=1)
        t.rounded_rectangle([bx, by, bx + wtxt + 14, by + 16], radius=4,
                            fill=tuple(int(BG[j] + (24 - BG[j]) * a) for j in range(3)),
                            outline=tuple(int(BG[j] + (70 - BG[j]) * a) for j in range(3)))
        t.text((bx + 7, by + 12), text, font=fe, fill=col_line, anchor="ls")
    if final:
        fe = ImageFont.truetype(F["semi"], 10.5)
        txt = f"{TOTAL:,} essais de 1945 à 2017".replace(",", " ")
        wtxt = t.textlength(txt, font=fe)
        bx, by = W - 18 - wtxt - 16, 100
        t.rounded_rectangle([bx, by, bx + wtxt + 16, by + 20], radius=5, fill=(24, 28, 40), outline=(70, 78, 100))
        t.text((bx + 8, by + 14), txt, font=fe, fill=TXT, anchor="ls")
    return np.asarray(im)
 
 
def timeline():
    """List of (frame, duration_ms)."""
    fr = []
    fin = render(Y1, 2, final=True)
    fr.append((fin, 3200))
    for i in range(1, 5):
        fr.append((render(Y1, 2, final=True, fade=1 - i / 5), 70))
    fr.append((render(Y0, 0), 450))
    shown = None                                   # (event year, remaining frames)
    for y in range(Y0, Y1 + 1):
        if y in EVENTS:
            shown = [y, 5 if y < 1966 else 7]
        two = y <= 1965 and sum(DATA[k].get(y, 0) for k in DATA) > 0
        phases = [0, 1] if two else [1]
        for ph in phases:
            ev = None
            if shown:
                ev = (shown[0], 1.0)
            fr.append((render(y, ph, ev), 70))
        if y in EVENTS:
            hold = int(650 * EVENTS[y][1])
            f, dms = fr[-1]
            fr[-1] = (f, dms + hold)
        if shown:
            shown[1] -= 1
            if shown[1] <= 0:
                shown = None
    fr.append((render(Y1, 2, (2017, 1.0)), 1400))
    return fr
 
 
# ----------------------------------------------------------------- palette + GIF
def nearest(cols, pal):
    dd = (cols ** 2).sum(1)[:, None] - 2 * cols @ pal.T + (pal ** 2).sum(1)[None]
    return dd.argmin(1)
 
 
def kmeans(X, w, k, iters=40, seed=5):
    rg = np.random.default_rng(seed)
    C = [X[rg.choice(len(X), p=w / w.sum())]]
    d2 = ((X - C[0]) ** 2).sum(1)
    for _ in range(1, min(k, len(X))):
        p = w * d2
        C.append(X[rg.choice(len(X), p=p / p.sum())])
        d2 = np.minimum(d2, ((X - C[-1]) ** 2).sum(1))
    C = np.array(C)
    for _ in range(iters):
        lab = nearest(X, C)
        sw = np.bincount(lab, weights=w, minlength=len(C))
        ok = sw > 0
        for ch in range(3):
            C[ok, ch] = np.bincount(lab, weights=w * X[:, ch], minlength=len(C))[ok] / sw[ok]
    return C
 
 
def ucols(arr):
    f_ = arr.reshape(-1, 3).astype(np.int64)
    key = (f_[:, 0] << 16) | (f_[:, 1] << 8) | f_[:, 2]
    uk, inv, cn = np.unique(key, return_inverse=True, return_counts=True)
    return np.stack([(uk >> 16) & 255, (uk >> 8) & 255, uk & 255], -1).astype(np.float64), inv, cn
 
 
def main():
    raw = timeline()
    frames, durations = [], []
    for f, dms in raw:
        if frames and np.array_equal(f, frames[-1]):
            durations[-1] += dms
        else:
            frames.append(f)
            durations.append(dms)
    cols, _, cn = ucols(np.concatenate([fr.reshape(-1, 3) for fr in frames]))
    if len(cols) > 60000:
        keep = np.random.default_rng(0).choice(len(cols), 60000, replace=False, p=np.sqrt(cn) / np.sqrt(cn).sum())
        cols_k, w_k = cols[keep], np.sqrt(cn[keep])
    else:
        cols_k, w_k = cols, np.sqrt(cn)
    fixed = np.array([BG, GRID, AXIS, TXT, TXT2, TXT3] + [c for _, _, c in COUNTRIES], np.float64)
    pal = np.clip(np.rint(np.vstack([fixed, kmeans(cols_k, w_k, 255 - len(fixed))])), 0, 255)
    pal_list = [int(v) for v in pal.flatten()] + [0, 0, 0] * (256 - len(pal))
    ims, prev, err = [], None, []
    for fr in frames:
        c, inv, _ = ucols(fr)
        ind = nearest(c, pal)[inv.ravel()].reshape(H, W).astype(np.uint8)
        q = pal[ind].astype(np.uint8)
        err.append(np.abs(q.astype(int) - fr.astype(int)).mean())
        out = ind if prev is None else np.where((q == prev).all(-1), 255, ind).astype(np.uint8)
        prev = q
        im = Image.fromarray(out, "P")
        im.putpalette(pal_list)
        ims.append(im)
    ims[0].save(OUT, save_all=True, append_images=ims[1:], duration=durations, loop=0,
                transparency=255, disposal=1, optimize=False)
    print(f"{OUT}: {os.path.getsize(OUT) / 1e6:.2f} MB, {W}x{H}, {len(frames)} frames, "
          f"{sum(durations) / 1000:.1f} s, {W * H * len(frames) / 1e6:.2f} MP, err {np.mean(err):.2f}, total tests {TOTAL}, "
          f"1962 = {sum(DATA[k].get(1962, 0) for k in DATA)}")
 
 
if __name__ == "__main__":
    main()