import numpy as np
from PIL import Image, ImageDraw, ImageFont
 
OUT = sys.argv[1] if len(sys.argv) > 1 else "Chaine_uranium_238_animation.gif"
W, H, S, DELAY = 300, 290, 4, 60
 
# (element name, symbol, Z, A, label suffix, half-life text, decay mode)
CHAIN = [
    ("Uranium", "U", 92, 238, "", "4,47 milliards d'années", "a"),
    ("Thorium", "Th", 90, 234, "", "24,1 jours", "b"),
    ("Protactinium", "Pa", 91, 234, "m", "1,17 minute", "b"),
    ("Uranium", "U", 92, 234, "", "245 500 ans", "a"),
    ("Thorium", "Th", 90, 230, "", "75 400 ans", "a"),
    ("Radium", "Ra", 88, 226, "", "1 600 ans", "a"),
    ("Radon", "Rn", 86, 222, "", "3,82 jours", "a"),
    ("Polonium", "Po", 84, 218, "", "3,10 minutes", "a"),
    ("Plomb", "Pb", 82, 214, "", "26,8 minutes", "b"),
    ("Bismuth", "Bi", 83, 214, "", "19,9 minutes", "b"),
    ("Polonium", "Po", 84, 214, "", "164 microsecondes", "a"),
    ("Plomb", "Pb", 82, 210, "", "22,2 ans", "b"),
    ("Bismuth", "Bi", 83, 210, "", "5,01 jours", "b"),
    ("Polonium", "Po", 84, 210, "", "138 jours", "a"),
    ("Plomb", "Pb", 82, 206, "", "stable", None),
]
LOG_T = [17.15, 6.32, 1.85, 12.89, 12.38, 10.70, 5.52, 2.27, 3.21, 3.08, -3.78, 8.84, 5.64, 7.08]  # log10(s)
SYMBOLS = {80: "Hg", 81: "Tl", 82: "Pb", 83: "Bi", 84: "Po", 85: "At", 86: "Rn", 87: "Fr",
           88: "Ra", 89: "Ac", 90: "Th", 91: "Pa", 92: "U", 93: "Np"}
STABLE = {(81, 203), (81, 205), (82, 204), (82, 206), (82, 207), (82, 208), (83, 209)}
 
# chart geometry: N on x, Z on y
N0, N1, Z0, Z1 = 122, 147, 81, 93
CELL = 10.0
GX0 = 34.0
GY1 = 270.0                                  # bottom of the Z0 row
GY0 = GY1 - (Z1 - Z0 + 1) * CELL
 
 
def cell_xy(Z, N):
    """Top-left corner of the cell (1x px)."""
    return GX0 + (N - N0) * CELL, GY1 - (Z - Z0 + 1) * CELL
 
 
BG = (255, 255, 255)
GRID, GRID_TXT, TXT, TXT2 = (234, 236, 240), (114, 119, 125), (32, 33, 34), (84, 89, 93)
ALPHA_C, BETA_C, STABLE_C = (247, 209, 70), (88, 160, 226), (32, 33, 34)
ARROW_A, ARROW_B = (196, 140, 0), (30, 100, 180)
F = {"semi": "/usr/share/fonts/opentype/inter/Inter-SemiBold.otf",
     "bold": "/usr/share/fonts/opentype/inter/Inter-Bold.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 mixc(c, a, base=BG):
    return tuple(int(base[i] + (c[i] - base[i]) * a) for i in range(3))
 
 
def P(x, y):
    return x * S, y * S
 
 
def N_of(i):
    return CHAIN[i][3] - CHAIN[i][2]
 
 
def centre(i):
    x, y = cell_xy(CHAIN[i][2], N_of(i))
    return x + CELL / 2, y + CELL / 2
 
 
def render(cur, prog, fade=1.0, final=False):
    """cur: index of the current nucleus; prog in [0,1]: transition to cur+1 (0 = resting)."""
    img = Image.new("RGB", (W * S, H * S), BG)
    d = ImageDraw.Draw(img)
    # grid and stable nuclides
    for Z in range(Z0, Z1 + 1):
        for N in range(N0, N1 + 1):
            x, y = cell_xy(Z, N)
            fill = STABLE_C if (Z, Z + N) in STABLE else BG
            d.rectangle([P(x, y), P(x + CELL, y + CELL)], fill=mixc(fill, fade) if fill != BG else BG,
                        outline=GRID, width=S // 2)
    # visited chain cells
    upto = len(CHAIN) - 1 if final else cur
    for i in range(upto + 1):
        mode = CHAIN[i][6]
        col = ALPHA_C if mode == "a" else (BETA_C if mode == "b" else STABLE_C)
        x, y = cell_xy(CHAIN[i][2], N_of(i))
        d.rectangle([P(x + 0.5, y + 0.5), P(x + CELL - 0.5, y + CELL - 0.5)], fill=mixc(col, fade))
    # arrows between visited nuclides
    def arrow(i, p):
        mode = CHAIN[i][6]
        col = ARROW_A if mode == "a" else ARROW_B
        (x0, y0), (x1, y1) = centre(i), centre(i + 1)
        xe, ye = x0 + (x1 - x0) * p, y0 + (y1 - y0) * p
        d.line([P(x0, y0), P(xe, ye)], fill=mixc(col, fade), width=round(1.5 * S))
        if p > 0.6:
            v = np.array([x1 - x0, y1 - y0])
            v /= np.linalg.norm(v)
            n = np.array([-v[1], v[0]])
            tip = np.array([xe, ye]) - v * 2.2
            pts = [tip + v * 3.2, tip - v * 0.6 + n * 2.6, tip - v * 0.6 - n * 2.6]
            d.polygon([P(*q) for q in pts], fill=mixc(col, fade))
    last_full = len(CHAIN) - 1 if final else cur
    for i in range(min(last_full, len(CHAIN) - 1)):
        arrow(i, 1.0)
    if not final and prog > 0 and cur < len(CHAIN) - 1:
        arrow(cur, min(prog / 0.7, 1.0))
    # current nucleus marker
    if not final and fade > 0.99:
        j = cur
        x, y = cell_xy(CHAIN[j][2], N_of(j))
        d.rectangle([P(x - 1.2, y - 1.2), P(x + CELL + 1.2, y + CELL + 1.2)], outline=TXT, width=round(1.6 * S))
    # emitted particle flying away
    if not final and 0 < prog < 1 and cur < len(CHAIN) - 1:
        mode = CHAIN[cur][6]
        cx, cy = centre(cur)
        dirv = np.array([-0.35, -1.0]) if mode == "a" else np.array([0.45, 1.0])
        dirv /= np.linalg.norm(dirv)
        dist = 6 + 34 * prog
        px_, py_ = cx + dirv[0] * dist, cy + dirv[1] * dist
        a = 1 - max(0.0, (prog - 0.6) / 0.4)
        col = ARROW_A if mode == "a" else ARROW_B
        r = 3.6 if mode == "a" else 2.4
        d.ellipse([P(px_ - r, py_ - r), P(px_ + r, py_ + r)], fill=mixc(col, a))
        lab = "α" if mode == "a" else "β⁻"
        d.text(P(px_ + 5, py_ - 1), lab, font=font("bold", 10), fill=mixc(col, a), anchor="lm")
    # axes labels
    for Z in range(Z0 + 1, Z1):
        _, y = cell_xy(Z, N0)
        d.text(P(GX0 - 4, y + CELL / 2), SYMBOLS[Z], font=font("med", 7.5), fill=GRID_TXT, anchor="rm")
    for N in range(125, N1 + 1, 5):
        x, _ = cell_xy(Z0, N)
        d.text(P(x + CELL / 2, GY1 + 3), str(N), font=font("reg", 7.5), fill=GRID_TXT, anchor="ma")
    d.text(P(GX0 + (N1 - N0 + 1) * CELL, GY1 + 12), "N (neutrons)", font=font("med", 7.5),
           fill=GRID_TXT, anchor="ra")
    d.text(P(8, GY0 - 4), "Z", font=font("med", 8), fill=GRID_TXT, anchor="ls")
 
    # info panel
    j = len(CHAIN) - 1 if final else cur
    name, sym, Z, A, suf, half, mode = CHAIN[j]
    a = fade
    d.text(P(16, 20), "Famille de l'uranium 238", font=font("semi", 10.5), fill=mixc(TXT2, a), anchor="ls")
    step = f"étape {j} / {len(CHAIN) - 1}" if j else "noyau de départ"
    d.text(P(W - 14, 20), step, font=font("reg", 9.5), fill=mixc(GRID_TXT, a), anchor="rs")
    # big nuclide notation: superscript mass number + symbol
    fs, fm = font("bold", 30), font("semi", 13)
    mass = f"{A}{suf}"
    x0 = 16.0
    d.text(P(x0, 44), mass, font=fm, fill=mixc(TXT, a), anchor="ls")
    xs = x0 + fm.getlength(mass) / S + 1
    col_sym = ALPHA_C if mode == "a" else (BETA_C if mode == "b" else STABLE_C)
    d.text(P(xs, 66), sym, font=fs, fill=mixc(TXT, a), anchor="ls")
    xr = max(xs + fs.getlength(sym) / S + 14, 92)
    d.text(P(xr, 44), f"{name} {A}{suf}", font=font("semi", 12), fill=mixc(TXT, a), anchor="ls")
    d.text(P(xr, 60), "demi-vie : " + half if mode else "noyau stable : fin de la chaîne",
           font=font("reg", 10), fill=mixc(TXT2, a), anchor="ls")
    if mode:
        nxt = CHAIN[j + 1]
        lab = "désintégration α" if mode == "a" else "désintégration β⁻"
        col = ARROW_A if mode == "a" else ARROW_B
        d.rounded_rectangle([P(xr - 1, 68), P(xr + 7 + font("semi", 9.5).getlength(lab) / S, 81)],
                            radius=3 * S, fill=mixc(col_sym, 0.35 * a))
        d.text(P(xr + 3, 78), lab, font=font("semi", 9.5), fill=mixc(col, a), anchor="ls")
        d.text(P(xr + 12 + font("semi", 9.5).getlength(lab) / S, 78), f"→ {nxt[1]} {nxt[3]}{nxt[4]}",
               font=font("reg", 9.5), fill=mixc(TXT2, a), anchor="ls")
    # legend
    lx, ly = 16, 97
    for col, lab in ((ALPHA_C, "émetteur α"), (BETA_C, "émetteur β⁻"), (STABLE_C, "stable")):
        d.rectangle([P(lx, ly - 7), P(lx + 7, ly)], fill=mixc(col, a))
        d.text(P(lx + 10, ly), lab, font=font("reg", 8.5), fill=mixc(TXT2, a), anchor="ls")
        lx += 12 + font("reg", 8.5).getlength(lab) / S + 12
    return np.asarray(img.reduce(S))
 
 
def timeline():
    st = [(0, 0.0, 1.0, True)] * 25
    st += [(0, 0.0, 1 - i / 5, True) for i in range(1, 5)]
    for i in range(len(CHAIN) - 1):
        # resting time grows (gently) with the half-life: billions of years vs microseconds
        rest = int(round(np.interp(LOG_T[i], [-4, 17], [1, 7])))
        st += [(i, 0.0, 1.0, False)] * rest
        for k in range(1, 7):
            st.append((i, k / 7, 1.0, False))
    st += [(len(CHAIN) - 1, 0.0, 1.0, False)] * 8
    st += [(0, 0.0, 1.0, True)] * 4
    return [render(*s) for s in st]
 
 
# ----------------------------------------------------------------- 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 fr in raw:
        if frames and np.array_equal(fr, frames[-1]):
            durations[-1] += DELAY
        else:
            frames.append(fr)
            durations.append(DELAY)
    cols, _, cn = ucols(np.concatenate([fr.reshape(-1, 3) for fr in frames]))
    fixed = np.array([BG, GRID, TXT, TXT2, ALPHA_C, BETA_C, STABLE_C, ARROW_A, ARROW_B], np.float64)
    pal = np.clip(np.rint(np.vstack([fixed, kmeans(cols, np.sqrt(cn), 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}")
 
 
if __name__ == "__main__":
    main()