#!/usr/bin/env python3
"""Turn generated effect sheets (a 2x4 grid of frames on black) into the effect atlas of the
high-res style: every frame cut out (a margin off, the cell frames the model drew), the dark
background turned into transparency (alpha from the brightness, the colour divided by it, so
glows and smoke keep their soft edges), scaled to 64x64; one row of 8 frames per sheet.

Usage:
  make_fx.py out.png sheet1-raw.png sheet2-raw.png ... [--margin 0.08] [--floor 0.1]

--floor  brightness counted as background (the model's not quite black)
Writes out.png (8 x 64 wide, 64 per sheet high) and out-preview.png (on grey, 2x). Needs Pillow.
"""
import argparse

from PIL import Image


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("out")
    ap.add_argument("sheets", nargs="+")
    ap.add_argument("--margin", type=float, default=0.08)
    ap.add_argument("--floor", type=float, default=0.1)
    a = ap.parse_args()
    size = 64
    atlas = Image.new("RGBA", (8 * size, len(a.sheets) * size), (0, 0, 0, 0))
    for row, path in enumerate(a.sheets):
        img = Image.open(path).convert("RGB")
        w, h = img.size
        for k in range(8):
            r, c = divmod(k, 4)
            x0, x1 = c * w / 4, (c + 1) * w / 4
            y0, y1 = r * h / 2, (r + 1) * h / 2
            m = a.margin * (x1 - x0)
            cell = img.crop((int(x0 + m), int(y0 + m), int(x1 - m), int(y1 - m))).resize((size, size), Image.BOX)
            px = cell.load()
            out = Image.new("RGBA", (size, size))
            op = out.load()
            for y in range(size):
                for x in range(size):
                    rr, gg, bb = (v / 255 for v in px[x, y])
                    al = max(0.0, (max(rr, gg, bb) - a.floor) / (1 - a.floor))
                    if al < 0.03:
                        op[x, y] = (0, 0, 0, 0)
                        continue
                    k2 = max(rr, gg, bb)
                    op[x, y] = (int(255 * rr / k2), int(255 * gg / k2), int(255 * bb / k2), int(255 * min(1.0, al * 1.3)))
            atlas.paste(out, (k * size, row * size))
    atlas.save(a.out)
    prev = Image.new("RGBA", atlas.size, (70, 66, 74, 255))
    prev.alpha_composite(atlas)
    prev.resize((prev.width * 2, prev.height * 2), Image.NEAREST).save(a.out.replace(".png", "-preview.png"))
    print(a.out)


if __name__ == "__main__":
    main()