"""
optimize_gif.py - make a GIF small for Kamoji without changing a single pixel.

    pip install pillow
    # and gifsicle: brew install gifsicle / sudo apt install gifsicle / https://www.lcdf.org/gifsicle/

    python optimize_gif.py dance.gif                  # writes dance.kamoji.gif
    python optimize_gif.py *.gif --out optimized/     # many at once, into a folder
    python optimize_gif.py dance.gif --gifsicle /path/to/gifsicle

What it does, for each GIF:
  1. Fits it to Kamoji's 240 x 240 screen if it is another size (centre crop + resize of every
     frame). Kamoji draws a GIF at its own size from the top-left corner, it does not scale.
  2. Re-encodes it losslessly with gifsicle -O3: each frame keeps only the rectangle that
     changed, pixels that did not change become transparent, repeated frames are merged.
  3. Checks that the result shows exactly the same picture as the input at every moment of the
     loop, and keeps it only if it is also smaller.

Needs Kamoji firmware 0.1.56 or later to play correctly (older firmware shows artifacts with
GIFs optimised this way). Guide: https://docs.kamoji.bot/docs/modding/optimizing-gifs
"""
import argparse
import hashlib
import io
import os
import shutil
import subprocess
import sys
import tempfile

from PIL import Image, ImageSequence

SIDE = 240  # Kamoji's screen


def visible_frames(path):
    """[(start ms, end ms, picture hash)]: what is on screen when, equal neighbours merged.
    Transparent pixels count as black, as on Kamoji."""
    out, t = [], 0
    for frame in ImageSequence.Iterator(Image.open(path)):
        rgba = frame.convert("RGBA")
        shown = Image.new("RGB", rgba.size)
        shown.paste(rgba.convert("RGB"), mask=rgba.split()[3])
        key = hashlib.md5(shown.tobytes()).hexdigest()
        delay = frame.info.get("duration", 0)
        if out and out[-1][2] == key:
            out[-1] = (out[-1][0], t + delay, key)
        else:
            out.append((t, t + delay, key))
        t += delay
    return out


def fit_to_screen(path, dest):
    """Centre-crop and resize every frame to 240 x 240, all frames on ONE palette: resizing
    blends edge colours slightly differently on every frame, so with a palette per frame almost
    no pixel would stay identical and the frame deltas (step 2) would save nothing."""
    src = Image.open(path)
    frames, delays = [], []
    for frame in ImageSequence.Iterator(src):
        rgb = frame.convert("RGB")
        side = min(rgb.size)
        left, top = (rgb.width - side) // 2, (rgb.height - side) // 2
        frames.append(rgb.crop((left, top, left + side, top + side)).resize((SIDE, SIDE), Image.LANCZOS))
        delays.append(frame.info.get("duration", 100))
    # The palette: from frames spread over the whole animation, side by side.
    sample = frames[:: max(1, len(frames) // 8)][:8]
    mosaic = Image.new("RGB", (SIDE * len(sample), SIDE))
    for k, f in enumerate(sample):
        mosaic.paste(f, (SIDE * k, 0))
    palette = mosaic.quantize(256, dither=Image.Dither.NONE)
    paletted = [f.quantize(palette=palette, dither=Image.Dither.NONE) for f in frames]
    paletted[0].save(dest, "GIF", save_all=True, append_images=paletted[1:], duration=delays,
                     loop=src.info.get("loop", 0), optimize=False, disposal=1)


def optimize(path, out_path, gifsicle):
    """Returns (bytes before, bytes after, note)."""
    before = os.path.getsize(path)
    with tempfile.TemporaryDirectory() as tmp:
        source = path
        size = Image.open(path).size
        note = ""
        if size != (SIDE, SIDE):
            source = os.path.join(tmp, "fitted.gif")
            fit_to_screen(path, source)
            note = f"resized from {size[0]}x{size[1]} (pixels change), then "
        candidate = os.path.join(tmp, "optimized.gif")
        run = subprocess.run([gifsicle, "-O3", source, "-o", candidate], capture_output=True, text=True)
        if run.returncode != 0:
            return before, None, "gifsicle failed: " + run.stderr.strip()
        if visible_frames(candidate) != visible_frames(source):
            return before, None, "re-encoded picture differs (kept the input)"
        # Keep the smaller of the (fitted) source and the re-encoded file.
        best = candidate if os.path.getsize(candidate) < os.path.getsize(source) else source
        if best == path:
            return before, None, note + "already optimal (nothing written)"
        shutil.copyfile(best, out_path)
        return before, os.path.getsize(out_path), note + "lossless: identical picture"


def change(before, after):
    """"-22.5%" (smaller) or "+12.0%" (larger, e.g. after a resize)."""
    pct = 100 * (after - before) / before
    return f"{pct:+.1f}%"


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("gifs", nargs="+")
    parser.add_argument("--out", help="folder for the results (default: next to each GIF, as <name>.kamoji.gif)")
    parser.add_argument("--gifsicle", default=shutil.which("gifsicle"), help="path to the gifsicle binary")
    args = parser.parse_args()
    if not args.gifsicle:
        sys.exit("gifsicle not found: install it (brew install gifsicle / apt install gifsicle) or pass --gifsicle")
    if args.out:
        os.makedirs(args.out, exist_ok=True)

    total_before = total_after = 0
    for path in args.gifs:
        name = os.path.basename(path)
        out_path = (os.path.join(args.out, name) if args.out
                    else os.path.splitext(path)[0] + ".kamoji.gif")
        before, after, note = optimize(path, out_path, args.gifsicle)
        total_before += before
        total_after += after if after is not None else before
        if after is None:
            print(f"{name}: {before / 1024:.1f} KB, {note}")
        else:
            print(f"{name}: {before / 1024:.1f} KB -> {after / 1024:.1f} KB "
                  f"({change(before, after)}), {note} -> {out_path}")
    if len(args.gifs) > 1 and total_before:
        print(f"total: {total_before / 1024:.1f} KB -> {total_after / 1024:.1f} KB "
              f"({change(total_before, total_after)})")


if __name__ == "__main__":
    main()