import json, os, subprocess, tempfile, importlib.util
from concurrent.futures import ProcessPoolExecutor
B = "/tmp/smol-bench"
def lib():
    spec = importlib.util.spec_from_file_location("s", f"{B}/ssim_lib.py"); m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m); return m
def entries(n):
    e = []
    add = lambda tool, fmt, setting, path: e.append((n, tool, fmt, setting, f"{B}/{path}"))
    for q in (50,60,70,75,80,85,90,95): add("Smol","jpeg",q,f"smol/j{q}/kodim{n}_compressed.jpg")
    for q in (60,75,82,90): add("Smol","webp",q,f"smol/w{q}/kodim{n}_compressed.webp")
    for q in (45,55,65,75,85): add("Smol","avif",q,f"smol/a{q}/kodim{n}_compressed.avif")
    for q in (30,40,50,60,65,70,75,80,85,90): add("sips","jpeg",q,f"sips/kodim{n}-q{q}.jpg")
    for q in (45,55,65,75,85): add("sips","avif",q,f"sips/kodim{n}-q{q}.avif")
    for q in (50,60,70,75,80,85,90,95):
        add("Caesium","jpeg",q,f"cae/j{q}/kodim{n}.jpg"); add("Caesium","webp",q,f"cae/w{q}/kodim{n}.webp")
    add("Caesium","png","lossless",f"cae/lossless/kodim{n}.png")
    for l in (1,2,3,4,5,6): add("Zipic","jpeg",f"level {l}",f"zipic/jpeg-l{l}/kodim{n}.jpeg")
    add("Zipic","webp","level 3",f"zipic/webp-l3/kodim{n}.webp")
    for q in (60,75,85): add("Squoosh","jpeg",q,f"squoosh/kodim{n}-mozJPEG-q{q}.jpg")
    add("Squoosh","webp",75,f"squoosh/kodim{n}-webP.webp"); add("Squoosh","avif","default",f"squoosh/kodim{n}-avif.avif"); add("Squoosh","png","oxipng default",f"squoosh/kodim{n}-oxiPNG.png")
    return e
_m = None; _refs = {}
def score(item):
    global _m
    n, tool, fmt, setting, path = item
    if not os.path.exists(path): return None
    if _m is None: _m = lib()
    ref = f"{B}/src/kodim{n}.png"
    if n not in _refs: _refs[n] = _m.load(ref)
    c = _m.load(path)
    tmp = tempfile.mktemp(suffix=".png"); subprocess.run(["magick", path, "-strip", tmp], check=True)
    s2 = float(subprocess.run(["ssimulacra2", ref, tmp], capture_output=True, text=True).stdout.strip()); os.unlink(tmp)
    return dict(image=f"kodim{n}", tool=tool, format=fmt, setting=setting, bytes=os.path.getsize(path), ssim=round(_m.ssim(_refs[n], c), 5), ssimulacra2=round(s2, 2))
if __name__ == "__main__":
    items = [e for i in range(1, 25) for e in entries(f"{i:02d}")]
    with ProcessPoolExecutor(8) as ex: out = [r for r in ex.map(score, items, chunksize=8) if r]
    json.dump(out, open(f"{B}/results24.json", "w"))
    from collections import Counter
    print(len(out), "measurements;", dict(Counter((r['tool'], r['format']) for r in out)))