import numpy as np
from PIL import Image, ImageDraw, ImageFont
 
OUT = sys.argv[1] if len(sys.argv) > 1 else "Palan_animation.gif"
W, H, S = 340, 330, 3
R = 12.0                                   # sheave radius (px)
XM1 = 96.0
XF1, XM2, XF2 = XM1 + 2 * R, XM1 + 4 * R, XM1 + 6 * R
XD, XG = XM1 - R, XF2 + R                  # standing end, hauling part
YF = 80.0                                  # fixed axle
YM0, LIFT = 232.0, 40.0                    # moving axle at rest, lift (px)
PX_PER_M = 80.0                            # 40 px = 0.5 m
YG0 = 112.0                                # end of the hauling part at rest
FLOOR = 298.0
 
INK, INK2, INK3 = (32, 33, 34), (84, 89, 93), (130, 136, 142)
ROPE = (150, 98, 48)
STEEL, STEEL_D = (176, 186, 198), (96, 108, 124)
SHEAVE, SHEAVE_D = (226, 230, 236), (86, 96, 110)
BLUE, ORANGE, RED = (51, 102, 204), (230, 120, 30), (200, 40, 40)
CRATE, CRATE_D = (214, 176, 120), (150, 112, 62)
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 P(x, y):
    return x * S, y * S
 
 
def ln(d, pts, col, w):
    d.line([P(*p) for p in pts], fill=col, width=round(w * S), joint="curve")
 
 
def disc(d, x, y, r, fill, outline=None, w=1.0):
    d.ellipse([P(x - r, y - r), P(x + r, y + r)], fill=fill, outline=outline,
              width=round(w * S) if outline else 0)
 
 
def arrow(d, x0, y0, x1, y1, col, w=1.8, head=6.0):
    ang = np.arctan2(y1 - y0, x1 - x0)
    bx, by = x1 - head * np.cos(ang), y1 - head * np.sin(ang)
    ln(d, [(x0, y0), (bx, by)], col, w)
    nx, ny = -np.sin(ang), np.cos(ang)
    d.polygon([P(x1, y1), P(bx + nx * head * 0.5, by + ny * head * 0.5), P(bx - nx * head * 0.5, by - ny * head * 0.5)],
              fill=col)
 
 
def sheave(d, x, y, phi, over):
    # rope groove arc (drawn later with the rope); sheave body with 3 spokes
    disc(d, x, y, R, SHEAVE, SHEAVE_D, 1.3)
    for k in range(3):
        a = phi + 2 * np.pi * k / 3
        ln(d, [(x + 2.2 * np.cos(a), y - 2.2 * np.sin(a)), (x + (R - 2.4) * np.cos(a), y - (R - 2.4) * np.sin(a))],
           SHEAVE_D, 1.4)
    disc(d, x, y, 2.4, STEEL_D)
 
 
def rope_arc(d, x, y, top):
    box = [P(x - R, y - R), P(x + R, y + R)]
    if top:
        d.arc(box, 180, 360, fill=ROPE, width=round(2.4 * S))
    else:
        d.arc(box, 0, 180, fill=ROPE, width=round(2.4 * S))
 
 
def fmt(v):
    return f"{v:.2f}".replace(".", ",") + " m"
 
 
def render(s):
    """s = lift in px (0 .. LIFT)."""
    img = Image.new("RGB", (W * S, H * S), (255, 255, 255))
    d = ImageDraw.Draw(img)
    ym = YM0 - s
    yg = YG0 + 4 * s
    # header
    d.text(P(12, 20), "Palan à 4 brins", font=font("bold", 14), fill=INK, anchor="ls")
    d.text(P(12, 33), "4 fois plus de corde à tirer, 4 fois moins d'effort", font=font("reg", 8.5), fill=INK2,
           anchor="ls")
    # beam and floor
    bx0, bx1 = XD - 24, XG + 14
    d.rectangle([P(bx0, 44), P(bx1, 50)], fill=(150, 156, 164))
    for x in np.arange(bx0 + 2, bx1, 6):
        ln(d, [(x, 44), (x + 5, 39)], INK3, 0.8)
    ln(d, [(20, FLOOR), (166, FLOOR)], INK2, 1.2)
    for x in np.arange(22, 166, 7):
        ln(d, [(x, FLOOR + 1), (x - 5, FLOOR + 6)], INK3, 0.8)
    # fixed block: hanger + plate
    ln(d, [((XF1 + XF2) / 2, 50), ((XF1 + XF2) / 2, YF - 10)], STEEL_D, 2.4)
    d.rounded_rectangle([P(XD - 5, YF - 5), P(XF2 + 5, YF + 5)], radius=round(3 * S), fill=STEEL, outline=STEEL_D,
                        width=S)
    # moving block plate + hook
    cxm = (XM1 + XM2) / 2
    d.rounded_rectangle([P(XM1 - 5, ym - 5), P(XM2 + 5, ym + 5)], radius=round(3 * S), fill=STEEL, outline=STEEL_D,
                        width=S)
    ln(d, [(cxm, ym + 5), (cxm, ym + 18)], STEEL_D, 2.4)
    d.arc([P(cxm - 5, ym + 15), P(cxm + 5, ym + 25)], 270 - 200, 270 + 120, fill=STEEL_D, width=round(2 * S))
    # load (crate) on slings
    ct, cb = ym + 30, ym + 66
    cw = 33
    ln(d, [(cxm, ym + 24), (cxm - cw + 5, ct)], INK2, 1.1)
    ln(d, [(cxm, ym + 24), (cxm + cw - 5, ct)], INK2, 1.1)
    d.rectangle([P(cxm - cw, ct), P(cxm + cw, cb)], fill=CRATE, outline=CRATE_D, width=round(1.2 * S))
    for yy in (ct + 12, ct + 24):
        ln(d, [(cxm - cw, yy), (cxm + cw, yy)], CRATE_D, 0.8)
    d.text(P(cxm - 7, ct + 18), "charge", font=font("semi", 8.5), fill=(90, 60, 25), anchor="mm")
    # sheaves (rotation from the kinematics)
    sheave(d, XM1, ym, s / R, False)
    sheave(d, XF1, YF, -2 * s / R, True)
    sheave(d, XM2, ym, 3 * s / R, False)
    sheave(d, XF2, YF, -4 * s / R, True)
    # rope: falls + arcs
    ln(d, [(XD, YF + 5), (XD, ym)], ROPE, 2.4)
    rope_arc(d, XM1, ym, False)
    ln(d, [(XM1 + R, ym), (XF1 - R, YF)], ROPE, 2.4)
    rope_arc(d, XF1, YF, True)
    ln(d, [(XF1 + R, YF), (XM2 - R, ym)], ROPE, 2.4)
    rope_arc(d, XM2, ym, False)
    ln(d, [(XM2 + R, ym), (XF2 - R, YF)], ROPE, 2.4)
    rope_arc(d, XF2, YF, True)
    ln(d, [(XG, YF), (XG, yg)], ROPE, 2.4)
    disc(d, XD, YF + 5, 2.2, STEEL_D)
    # handle
    d.rounded_rectangle([P(XG - 7, yg), P(XG + 7, yg + 5)], radius=round(2 * S), fill=INK2)
    # fall numbers
    ymid = (YF + ym) / 2 + 6
    for i, x in enumerate((XD, XF1 - R, XM2 - R, XF2 - R)):
        disc(d, x, ymid, 5.2, (255, 255, 255), ROPE, 1.1)
        d.text(P(x, ymid + 0.3), str(i + 1), font=font("bold", 7), fill=ROPE, anchor="mm")
    # forces
    arrow(d, XG, yg + 7, XG, yg + 24, RED, 1.8, 6)
    d.text(P(XG + 5, yg + 20), "F = P/4", font=font("bold", 9), fill=RED, anchor="ls")
    arrow(d, cxm + 22, ct + 5, cxm + 22, cb - 2, RED, 1.8, 6)
    d.text(P(cxm + 18, cb - 4), "P", font=font("bold", 9), fill=RED, anchor="rs")
    # labels with leader lines
    lf = font("med", 8)
    d.text(P(12, YF + 3), "moufle fixe", font=lf, fill=INK, anchor="ls")
    ln(d, [(58, YF), (XD - 7, YF)], INK3, 0.8)
    d.text(P(12, ym + 3), "moufle", font=lf, fill=INK, anchor="ls")
    d.text(P(12, ym + 13), "mobile", font=lf, fill=INK, anchor="ls")
    ln(d, [(44, ym), (XM1 - 7, ym)], INK3, 0.8)
    yd = YF + 26
    d.text(P(12, yd + 3), "dormant", font=lf, fill=INK, anchor="ls")
    ln(d, [(48, yd), (XD - 3, yd)], INK3, 0.8)
    d.text(P(XG + 6, YF + 30), "garant", font=lf, fill=INK, anchor="ls")
    d.text(P(12, ymid + 3), "brins", font=lf, fill=ROPE, anchor="ls")
    ln(d, [(37, ymid), (XD - 7, ymid)], INK3, 0.8)
 
    # right panel: gauges (same scale)
    gx0 = 232
    d.text(P(gx0, 112), "même échelle :", font=font("med", 8), fill=INK2, anchor="ls")
    base = FLOOR
    rope_m = 4 * s / PX_PER_M
    load_m = s / PX_PER_M
    for k, (x, val, col, l1, l2) in enumerate(((gx0 + 6, rope_m, BLUE, "corde", "tirée"),
                                               (gx0 + 62, load_m, ORANGE, "charge", "levée"))):
        hpx = val * PX_PER_M
        full = (4 * LIFT if k == 0 else LIFT)
        d.rectangle([P(x, base - full), P(x + 18, base)], outline=(222, 226, 231), width=S)
        if hpx > 0.05:
            d.rectangle([P(x, base - hpx), P(x + 18, base)], fill=col)
        d.text(P(x + 9, base - hpx - 4), fmt(val), font=font("bold", 9), fill=col, anchor="ms")
        d.text(P(x + 9, base + 11), l1, font=font("med", 8), fill=INK, anchor="ms")
        d.text(P(x + 9, base + 21), l2, font=font("med", 8), fill=INK, anchor="ms")
    ln(d, [(gx0, base), (gx0 + 88, base)], INK2, 1.2)
    return np.asarray(img.reduce(S))
 
 
def ease(x):
    return 0.5 - 0.5 * np.cos(np.pi * x)
 
 
def frames_list():
    out = [(render(LIFT), 2600)]
    nd, nu = 22, 50
    for i in range(1, nd + 1):                         # lowering back
        out.append((render(LIFT * (1 - ease(i / nd))), 60))
    out.append((render(0), 700))
    for i in range(1, nu + 1):                         # hauling
        out.append((render(LIFT * ease(i / nu)), 70))
    return out
 
 
# ----------------------------------------------------------------- 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=30, 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))):
        pr = w * d2
        C.append(X[rg.choice(len(X), p=pr / pr.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 = frames_list()
    frames, durations = [], []
    for fr, dms in raw:
        if frames and np.array_equal(fr, frames[-1]):
            durations[-1] += dms
        else:
            frames.append(fr)
            durations.append(dms)
    cols, _, cn = ucols(np.concatenate([fr.reshape(-1, 3) for fr in frames]))
    fixed = np.array([(255, 255, 255), INK, INK2, INK3, ROPE, STEEL, STEEL_D, SHEAVE, SHEAVE_D, BLUE, ORANGE, RED,
                      CRATE, CRATE_D], 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()