#!/usr/bin/env python3
"""Visualise PANORAMA's masks: one panel per phrase, plus a combined overlay.

    python visualize.py --image test.png --pred ./out               # after predict_panocaps.py
    python visualize.py --image ./my_images/ --pred ./out           # a whole folder

For each image it writes a folder <stem>_viz/ containing:
    00_<phrase>.png ...   the image with ONLY that phrase's mask(s) lit; everything else dimmed
    overlay.png           all masks tinted per phrase + a colour->phrase legend
    caption.txt           the caption verbatim, with the <p>...</p>[SEG] markers
    caption_clean.txt     the same caption with the markers stripped

Per-phrase panels dim the background instead of tinting the mask, so you judge the mask against the
object's real colours (a 50% tint turns a red mug green). The overlay is for coverage; the panels
are for correctness.

Pillow + numpy only — no GPU, no extra libraries.
"""
import argparse
import json
import os
import re

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

# Colour is per PHRASE, not per mask: one phrase can own several masks ("three chairs"), and they
# belong together. masks_per_phrase in the json gives the split.
# Okabe-Ito (colour-blind friendly), extended; cycles if a caption names more regions than colours.
# Names are the canonical Okabe-Ito ones for the first 8; the rest are our tints/extensions.
# Index == the phrase number printed in the per-phrase panels, so "phrase 03" is always yellow.
PALETTE_NAMES = [
    "orange",          # 0   (230, 159,   0)  Okabe-Ito
    "sky-blue",        # 1   ( 86, 180, 233)  Okabe-Ito
    "bluish-green",    # 2   (  0, 158, 115)  Okabe-Ito
    "yellow",          # 3   (240, 228,  66)  Okabe-Ito
    "blue",            # 4   (  0, 114, 178)  Okabe-Ito
    "vermillion",      # 5   (213,  94,   0)  Okabe-Ito
    "reddish-purple",  # 6   (204, 121, 167)  Okabe-Ito
    "grey",            # 7   (153, 153, 153)  Okabe-Ito (black slot, lightened)
    "salmon",          # 8   (255, 128, 128)  extension
    "mint",            # 9   (128, 255, 128)  extension
    "periwinkle",      # 10  (128, 128, 255)  extension
    "apricot",         # 11  (255, 200, 100)  extension
    "amethyst",        # 12  (170, 110, 200)  extension
    "teal",            # 13  (100, 200, 190)  extension
    "raspberry",       # 14  (200,  90, 140)  extension
]
PALETTE = [
    (230, 159, 0), (86, 180, 233), (0, 158, 115), (240, 228, 66), (0, 114, 178),
    (213, 94, 0), (204, 121, 167), (153, 153, 153), (255, 128, 128), (128, 255, 128),
    (128, 128, 255), (255, 200, 100), (170, 110, 200), (100, 200, 190), (200, 90, 140),
]
assert len(PALETTE_NAMES) == len(PALETTE), "palette names and colours out of sync"


def colour_name(i):
    """Name of the colour phrase `i` is drawn in (wraps like PALETTE does)."""
    return PALETTE_NAMES[i % len(PALETTE_NAMES)]
OVERLAY_ALPHA = 0.45   # combined view: strong enough to read coverage
DIM = 0.28             # per-phrase panels: how much of the image survives outside the mask


def phrase_of_mask(counts, n_masks):
    """mask index -> phrase index, from masks_per_phrase (falls back to 1:1)."""
    if not counts:
        return list(range(n_masks))
    out = []
    for pi, c in enumerate(counts):
        out.extend([pi] * int(c))
    if len(out) < n_masks:                    # trust the masks, not the counts
        out.extend([len(counts) - 1] * (n_masks - len(out)))
    return out[:n_masks]


def outline(mask):
    """1px boundary of a boolean mask."""
    e = np.zeros_like(mask)
    e[:-1, :] |= mask[:-1, :] ^ mask[1:, :]
    e[:, :-1] |= mask[:, :-1] ^ mask[:, 1:]
    return e


def _font(size=15):
    for p in ('DejaVuSans.ttf', '/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf'):
        try:
            return ImageFont.truetype(p, size)
        except Exception:
            pass
    return ImageFont.load_default()


def panel(img_arr, masks, colour, label):
    """One phrase: its mask(s) at full brightness, everything else dimmed."""
    m = np.zeros(img_arr.shape[:2], bool)
    for k in masks:
        m |= k
    out = (img_arr * DIM).astype(np.float32)         # dim everything
    out[m] = img_arr[m]                              # restore the mask region
    out[outline(m)] = colour                         # trace it
    im = Image.fromarray(out.astype(np.uint8))

    d = ImageDraw.Draw(im)
    f = _font(16)
    px = int(m.sum())
    txt = f'{label}   —   {px:,} px' + (f'  ({len(masks)} masks)' if len(masks) > 1 else '')
    x0, y0, x1, y1 = d.textbbox((0, 0), txt, font=f)
    d.rectangle([6, 6, 14 + (x1 - x0), 14 + (y1 - y0)], fill=(0, 0, 0))
    d.rectangle([6, 6, 12, 14 + (y1 - y0)], fill=colour)          # colour tab
    d.text((16, 10), txt, fill=(255, 255, 255), font=f)
    return im


def _wrap(draw, text, font, width):
    """Greedy word-wrap to `width` px."""
    lines, cur = [], ''
    for w in text.split():
        t = f'{cur} {w}'.strip()
        if draw.textlength(t, font=font) <= width or not cur:
            cur = t
        else:
            lines.append(cur)
            cur = w
    if cur:
        lines.append(cur)
    return lines


def overlay(img_arr, masks, phrases, m2p, caption=''):
    arr = img_arr.astype(np.float32)
    for i, m in enumerate(masks):
        if not m.any():
            continue
        c = np.array(PALETTE[m2p[i] % len(PALETTE)], np.float32)
        arr[m] = (1 - OVERLAY_ALPHA) * arr[m] + OVERLAY_ALPHA * c
    for i, m in enumerate(masks):
        if m.any():
            arr[outline(m)] = PALETTE[m2p[i] % len(PALETTE)]
    im = Image.fromarray(arr.astype(np.uint8))

    f, fc = _font(14), _font(15)
    probe = ImageDraw.Draw(im)
    # caption block: the [SEG] markers are noise on a picture that already shows the masks
    cap = re.sub(r'\s*\[SEG\]', '', caption).replace('<p>', '').replace('</p>', '')
    cap = re.sub(r'\s+', ' ', cap).strip()
    cap_lines = _wrap(probe, cap, fc, im.width - 24) if cap else []
    cap_h = (len(cap_lines) * 20 + 16) if cap_lines else 0

    lh = 20
    strip_h = cap_h + max(lh * len(phrases) + 12, 24)
    strip = Image.new('RGB', (im.width, strip_h), (255, 255, 255))
    sd = ImageDraw.Draw(strip)
    y = 8
    for ln in cap_lines:
        sd.text((12, y), ln, fill=(0, 0, 0), font=fc)
        y += 20
    if cap_lines:
        y += 8
        sd.line([12, y - 5, im.width - 12, y - 5], fill=(210, 210, 210))
    for pi, ph in enumerate(phrases):
        sd.rectangle([6, y + 3, 20, y + 15], fill=PALETTE[pi % len(PALETTE)], outline=(0, 0, 0))
        n = m2p.count(pi)
        sd.text((28, y), f'{pi:02d}  {ph.strip()}' + (f'   ({n} masks)' if n > 1 else ''),
                fill=(0, 0, 0), font=f)
        y += lh

    out = Image.new('RGB', (im.width, im.height + strip_h), (255, 255, 255))
    out.paste(im, (0, 0))
    out.paste(strip, (0, im.height))
    return out


def slug(s, n=40):
    s = re.sub(r'[^a-z0-9]+', '_', s.strip().lower()).strip('_')
    return (s[:n] or 'phrase')


def run(image_path, pred_dir):
    stem = os.path.splitext(os.path.basename(image_path))[0]
    jf, mf = os.path.join(pred_dir, f'{stem}.json'), os.path.join(pred_dir, f'{stem}_masks.npz')
    if not (os.path.isfile(jf) and os.path.isfile(mf)):
        miss = [os.path.basename(f) for f in (jf, mf) if not os.path.isfile(f)]
        print(f'  skip {stem}: missing {", ".join(miss)} in {pred_dir} — run predict_panocaps.py '
              f'with --out {pred_dir} first')
        return
    d = json.load(open(jf))
    masks = np.load(mf)['masks'].astype(bool)
    phrases = d['phrases']
    img = np.asarray(Image.open(image_path).convert('RGB'))
    if masks.size and masks.shape[1:] != img.shape[:2]:
        print(f'  !! {stem}: masks are {masks.shape[1:]} but the image is {img.shape[:2]}')
        return
    m2p = phrase_of_mask(d.get('masks_per_phrase'), len(masks))

    outdir = os.path.join(pred_dir, f'{stem}_viz')
    os.makedirs(outdir, exist_ok=True)

    # The caption, next to the panels, in both forms:
    #   caption.txt        — verbatim, with the <p>...</p>[SEG] markers. This is the alignment
    #                        record: the i-th <p>...</p> is the i-th phrase, hence panel 00, 01, ...
    #   caption_clean.txt  — the markers stripped, i.e. the prose on its own, for reading or for
    #                        feeding somewhere that does not want the markup.
    raw = d['caption']
    clean = re.sub(r'\s*\[SEG\]', '', raw).replace('<p>', '').replace('</p>', '')
    clean = re.sub(r'\s+', ' ', clean)
    # the model writes "<p> red mug </p> [SEG]." — the spaces live INSIDE the tags, so removing them
    # strands a space before the punctuation ("a red mug ."). Close it up.
    clean = re.sub(r'\s+([,.;:!?%)\]])', r'\1', clean)
    clean = re.sub(r'([(\[])\s+', r'\1', clean).strip()
    open(os.path.join(outdir, 'caption.txt'), 'w').write(raw.strip() + '\n')
    open(os.path.join(outdir, 'caption_clean.txt'), 'w').write(clean + '\n')

    for pi, ph in enumerate(phrases):
        mine = [masks[i] for i in range(len(masks)) if m2p[i] == pi]
        if not mine:
            continue
        panel(img, mine, PALETTE[pi % len(PALETTE)], f'{pi:02d}  {ph.strip()}').save(
            os.path.join(outdir, f'{pi:02d}_{slug(ph)}.png'))
    overlay(img, masks, phrases, m2p).save(os.path.join(outdir, 'overlay.png'))
    print(f'{stem}: {len(masks)} mask(s), {len(phrases)} phrase(s) -> {outdir}/  '
          f'({len(phrases)} panels + overlay.png + caption.txt/caption_clean.txt)')


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--image', required=True, help='the image file, or the directory you predicted on')
    ap.add_argument('--pred', default='./out',
                    help='the --out DIRECTORY predict_panocaps.py wrote (not a file)')
    a = ap.parse_args()

    if os.path.isfile(a.pred):                     # easy slip: passing <stem>_masks.npz
        a.pred = os.path.dirname(os.path.abspath(a.pred))
        print(f'--pred is a file; using its directory: {a.pred}')
    elif not os.path.isdir(a.pred):
        raise SystemExit(f'--pred: no such directory: {a.pred}')

    imgs = ([os.path.join(a.image, f) for f in sorted(os.listdir(a.image))
             if f.lower().endswith(('.jpg', '.jpeg', '.png', '.webp', '.bmp'))]
            if os.path.isdir(a.image) else [a.image])
    for p in imgs:
        run(p, a.pred)


if __name__ == '__main__':
    main()