#!/usr/bin/env python3
"""Download each claim image, measure colour, delete it. One JSON line per claim, so a
killed run resumes and disk never holds more than one image.

Colourfulness is Hasler & Suesstrunk (2003) 'Measuring colourfulness in natural images',
the metric that paper validated against human ratings -- not a saturation average I made up.
"""
import json, os, urllib.request, numpy as np
from PIL import Image
UA = "Mozilla/5.0 (compatible; agentatwork/1.0; +https://agentatwork.xyz)"
OUT = "colour.jsonl"
TMP = "/tmp/claimimg.bin"

done = {json.loads(l)["id"] for l in open(OUT)} if os.path.exists(OUT) else set()
recs = [json.loads(l) for l in open("meta.jsonl")]

def measure(path):
    im = Image.open(path)
    im = im.convert("RGB")
    w, h = im.size
    im.thumbnail((256, 256))
    a = np.asarray(im).astype(np.float64)
    R, G, B = a[..., 0], a[..., 1], a[..., 2]
    rg = R - G
    yb = 0.5 * (R + G) - B
    colourfulness = float(np.hypot(rg.std(), yb.std()) + 0.3 * np.hypot(rg.mean(), yb.mean()))
    mx, mn = a.max(2), a.min(2)
    sat = np.where(mx > 0, (mx - mn) / np.maximum(mx, 1e-9), 0.0)   # HSV S
    val = mx / 255.0                                               # HSV V
    return dict(w=w, h=h, colourfulness=round(colourfulness, 3),
                sat=round(float(sat.mean()), 4), val=round(float(val.mean()), 4),
                dark_frac=round(float((val < 0.25).mean()), 4),
                bright_frac=round(float((val > 0.75).mean()), 4),
                grey_frac=round(float((sat < 0.10).mean()), 4))

with open(OUT, "a") as fh:
    for r in recs:
        if r["id"] in done:
            continue
        rec = {"id": r["id"], "url": r["image"]}
        try:
            req = urllib.request.Request(r["image"], headers={"User-Agent": UA})
            with urllib.request.urlopen(req, timeout=90) as x, open(TMP, "wb") as f:
                rec["bytes"] = f.write(x.read())
            rec.update(measure(TMP))
        except Exception as e:
            rec["error"] = f"{type(e).__name__}: {e}"[:160]
        finally:
            if os.path.exists(TMP):
                os.remove(TMP)
        fh.write(json.dumps(rec) + "\n"); fh.flush()
        print(("." if "colourfulness" in rec else "E"), end="", flush=True)
print()