"""fix_blob_targets.py — rebuild exp011's blob targets with the CORRECT
fused_json schema (diagnosed 2026-07-18: polygon is a FLAT [x1,y1,x2,y2,...]
list, and coordinates live in the declared coord_space with image_size — not
the assumed NORM_0_1000 pair-list). Row order replicates _iter_fused exactly
(same files, same gates); patches blob/val_blob into the existing cache.
LOUD assert: >50% of rows must produce non-empty blobs. CPU-only.
"""
from __future__ import annotations

import json
import os
import sys

sys.path[:0] = ["pod2", "."]

import numpy as np
import torch

from dexp011_fused_multiband import _iter_fused, N_TRAIN, N_VAL, DATA_DIR


def rasterize_v2(fused_json_str: str, size: int = 64) -> np.ndarray:
    from PIL import Image, ImageDraw
    try:
        fj = json.loads(fused_json_str)
    except Exception:
        return np.zeros((size, size), dtype=np.float32)
    cs = str(fj.get("coord_space", "")).lower()
    im_size = fj.get("image_size")
    if isinstance(im_size, (list, tuple)) and len(im_size) >= 2:
        W, H = float(im_size[0]) or None, float(im_size[1]) or None
    elif isinstance(im_size, dict):
        W = float(im_size.get("width") or im_size.get("w") or 0) or None
        H = float(im_size.get("height") or im_size.get("h") or 0) or None
    else:
        W = H = None
    img = Image.new("L", (size, size), 0)
    draw = ImageDraw.Draw(img)
    for ent in (fj.get("entities") or []):
        poly = (ent.get("mask") or {}).get("polygon") or []
        if len(poly) >= 6 and all(isinstance(v, (int, float))
                                  for v in poly[:6]):
            xs, ys = poly[0::2], poly[1::2]
            if "1000" in cs or (max(xs + ys) <= 1000 and W is None):
                sx = sy = 1.0 / 1000.0
            elif W and H:
                sx, sy = 1.0 / W, 1.0 / H
            else:
                m = max(max(xs), max(ys), 1.0)
                sx = sy = 1.0 / m
            pts = [(min(max(x * sx, 0), 1) * (size - 1),
                    min(max(y * sy, 0), 1) * (size - 1))
                   for x, y in zip(xs, ys)]
            if len(pts) >= 3:
                draw.polygon(pts, fill=255)
    return np.asarray(img, dtype=np.float32) / 255.0


def main():
    blobs, n_empty, cs_seen = [], 0, {}
    for r in _iter_fused(N_TRAIN + N_VAL):
        s = str(r.get("fused_json", ""))
        b = rasterize_v2(s)
        if b.max() == 0:
            n_empty += 1
        try:
            cs = json.loads(s).get("coord_space", "?")
            cs_seen[str(cs)] = cs_seen.get(str(cs), 0) + 1
        except Exception:
            pass
        blobs.append(torch.from_numpy(b))
    blobs = torch.stack(blobs)
    frac = 1 - n_empty / len(blobs)
    print(f"[fix_blob] non-empty {frac:.1%} (empty {n_empty}/{len(blobs)}); "
          f"coord_spaces {cs_seen}; mean occupancy "
          f"{blobs.mean().item():.3f}", flush=True)
    assert frac > 0.5, f"still mostly empty ({frac:.1%}) — schema wrong again"
    f = os.path.join(DATA_DIR, "cache.pt")
    cache = torch.load(f, map_location="cpu", weights_only=True)
    cache["blob"] = blobs[:N_TRAIN]
    cache["val_blob"] = blobs[N_TRAIN:]
    cache["blob_empty_count"] = n_empty
    torch.save(cache, f)
    print(f"[fix_blob] cache patched: {f}", flush=True)


if __name__ == "__main__":
    main()