"""
Evaluate trained face-parsing checkpoints on the OTHER face databases in
/Users/ari/FaceSegmentation (no ground truth -> qualitative overlay grids +
distribution stats), and optionally on the held-out CelebAMask-HQ test split
(ground truth -> mIoU via the ultralytics semantic validator).

Per database it samples N images deterministically, center-crops to square
(matches the iOS app's preprocessing), runs the raw logits head, and writes:
  <out>/<tag>/<db>_grid.jpg      colorized overlay contact sheet
  <out>/<tag>/report.md          per-db stats (class coverage, confidence)

Usage:
  eval_other_dbs.py --weights runs/celeba_large/weights/best.pt --tag large
  eval_other_dbs.py --weights ... --tag large --test-split   # adds GT mIoU
"""
import argparse, glob, json, os, random

import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont

from palette import CLASS_NAMES, PALETTE
from ultralytics import YOLO

DBS = {
    "celebrity_faces": "/Users/ari/FaceSegmentation/Celebrity Faces Dataset",
    "human_faces_real": "/Users/ari/FaceSegmentation/Human Faces Dataset 2/Real Images",
    "human_faces_ai": "/Users/ari/FaceSegmentation/Human Faces Dataset 2/AI-Generated Images",
    "img_align_celeba": "/Users/ari/FaceSegmentation/img_align_celeba",
}
DATA_YAML = "/Users/ari/FaceSegmentation/dataset_celebamaskhq_semantic/data.yaml"
PAL = np.array(PALETTE, dtype=np.uint8)


def sample_images(root, n, seed=0):
    exts = ("*.jpg", "*.jpeg", "*.png", "*.JPG", "*.webp")
    files = []
    for e in exts:
        files += glob.glob(os.path.join(root, "**", e), recursive=True)
    files = sorted(files)
    random.Random(seed).shuffle(files)
    return files[:n]


def center_square(im):
    w, h = im.size
    s = min(w, h)
    return im.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))


@torch.no_grad()
def infer(core, im, R, device):
    x = torch.from_numpy(np.asarray(im, dtype=np.float32) / 255.0).permute(2, 0, 1)[None].to(device)
    z = core(x)
    z = z[0] if isinstance(z, (list, tuple)) else z          # (1,19,g,g)
    z = torch.nn.functional.interpolate(z.float(), size=(R, R), mode="bilinear", align_corners=False)
    prob = z.softmax(1)[0]                                    # (19,R,R)
    conf, cls = prob.max(0)
    return cls.byte().cpu().numpy(), conf.cpu().numpy()


def overlay(im, cls, alpha=0.55):
    rgb = np.asarray(im, dtype=np.float32)
    col = PAL[cls].astype(np.float32)
    mask = (cls > 0)[..., None].astype(np.float32) * alpha
    return Image.fromarray((rgb * (1 - mask) + col * mask).astype(np.uint8))


def contact_sheet(tiles, cols, pad=4, label=None):
    n = len(tiles)
    rows = (n + cols - 1) // cols
    w, h = tiles[0].size
    sheet = Image.new("RGB", (cols * (w + pad) + pad, rows * (h + pad) + pad + (28 if label else 0)), (18, 18, 18))
    for i, t in enumerate(tiles):
        r, c = divmod(i, cols)
        sheet.paste(t, (pad + c * (w + pad), (28 if label else 0) + pad + r * (h + pad)))
    if label:
        ImageDraw.Draw(sheet).text((8, 6), label, fill=(240, 240, 240))
    return sheet


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", required=True)
    ap.add_argument("--tag", required=True)
    ap.add_argument("--n", type=int, default=24)
    ap.add_argument("--imgsz", type=int, default=512)
    ap.add_argument("--tile", type=int, default=256)
    ap.add_argument("--device", default="mps")
    ap.add_argument("--out", default="/Users/ari/FaceSegmentation/eval_other_dbs")
    ap.add_argument("--test-split", action="store_true",
                    help="also run GT mIoU on the CelebAMask-HQ test split")
    args = ap.parse_args()

    outdir = os.path.join(args.out, args.tag)
    os.makedirs(outdir, exist_ok=True)
    y = YOLO(args.weights)
    core = y.model.eval().to(args.device)
    R = args.imgsz

    report = [f"# Cross-database evaluation — {args.tag}", f"weights: `{args.weights}`", ""]
    stats_all = {}
    for db, root in DBS.items():
        files = sample_images(root, args.n)
        if not files:
            report.append(f"## {db}\n(no images found at {root})\n")
            continue
        tiles, cov, confs = [], np.zeros(len(CLASS_NAMES)), []
        for p in files:
            try:
                im = center_square(Image.open(p).convert("RGB")).resize((R, R), Image.BILINEAR)
            except Exception:
                continue
            cls, conf = infer(core, im, R, args.device)
            cov += np.bincount(cls.ravel(), minlength=len(CLASS_NAMES)) / cls.size
            confs.append(float(conf.mean()))
            tiles.append(overlay(im, cls).resize((args.tile, args.tile), Image.BILINEAR))
        cov /= max(len(tiles), 1)
        sheet = contact_sheet(tiles, cols=6, label=f"{args.tag} · {db} · n={len(tiles)}")
        sheet_path = os.path.join(outdir, f"{db}_grid.jpg")
        sheet.save(sheet_path, quality=90)
        top = sorted(enumerate(cov), key=lambda t: -t[1])[:8]
        stats_all[db] = {"n": len(tiles), "mean_confidence": float(np.mean(confs)) if confs else 0.0,
                         "class_coverage": {CLASS_NAMES[i]: round(float(v), 4) for i, v in top}}
        report += [f"## {db}", f"- n={len(tiles)}  mean pixel confidence: {np.mean(confs):.3f}",
                   "- top class coverage: " + ", ".join(f"{CLASS_NAMES[i]} {v:.1%}" for i, v in top),
                   f"- grid: `{sheet_path}`", ""]

    if args.test_split:
        report.append("## CelebAMask-HQ test split (ground truth)")
        try:
            m = y.val(data=DATA_YAML, split="test", imgsz=R, device=args.device, plots=False)
            rd = {k: float(v) for k, v in m.results_dict.items()}
            stats_all["celebamask_test"] = rd
            report += ["```", json.dumps(rd, indent=2), "```", ""]
        except Exception as e:
            report.append(f"validator failed: {e!r}")

    with open(os.path.join(outdir, "stats.json"), "w") as f:
        json.dump(stats_all, f, indent=2)
    with open(os.path.join(outdir, "report.md"), "w") as f:
        f.write("\n".join(report))
    print("wrote", outdir)


if __name__ == "__main__":
    main()