"""Diffusion-based synthetic photographs of runic inscriptions (latent diffusion + ControlNet-Canny).

Follows the recipe of the recognition thesis: the transliteration is mapped to runes, rendered
black on white on a 1024x1024 canvas with automatic font size and line wrapping, converted to a
Canny edge map (thresholds 50/150, one-pixel dilation) and used as the ControlNet condition with
a fixed positive/negative prompt pair, 28 steps and guidance 7.0. Adapted for SDXL: a 1216x832 canvas instead of
1024x1024 and conditioning scale 0.8 instead of 0.65, with extra negative terms, because SDXL fills empty
stone with invented signs at the thesis settings.
SDXL-base with the SDXL Canny ControlNet replaces SD3-Medium, whose weights are gated.
Resumable per shard; labels are appended after every image.
"""
import argparse
import json
import os
import random
import sys

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

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from synth import FONT_DIR, RUNE, lexicon, sample_phrase  # noqa: E402

POSITIVE = ("deeply carved runes on a dark grey granite runestone, weathered eroded stone surface, "
            "plain stone around the inscription, even diffuse daylight, archaeological macro photograph, sharp focus")
NEGATIVE = ("polished, smooth, printed font, fantasy, illustration, painting, cartoon, glowing, "
            "text overlay, watermark, signature, blurry, extra runes, additional inscriptions, symbols, ornaments, "
            "patterns, tiles, bricks, frame")


def render_condition(translit, rng, width=1216, height=832):
    runes = " ".join("".join(RUNE.get(c, "") for c in w) for w in translit.split())
    runes = runes.replace(" ", " ᛫ ")  # runic single punctuation as word divider
    font_path = os.path.join(FONT_DIR, "NotoSansRunic-Regular.ttf")
    size = width
    canvas = Image.new("L", (width, height), 255)
    draw = ImageDraw.Draw(canvas)
    fsize = 200
    while fsize > 24:
        font = ImageFont.truetype(font_path, fsize)
        words, lines, cur = runes.split(" "), [], ""
        for w in words:
            trial = (cur + " " + w).strip()
            if draw.textlength(trial, font=font) <= size * 0.86:
                cur = trial
            else:
                lines.append(cur)
                cur = w
        lines.append(cur)
        line_h = int(fsize * 1.35)
        if len(lines) * line_h <= height * 0.8 and all(draw.textlength(l, font=font) <= size * 0.86 for l in lines):
            break
        fsize -= 8
    y0 = (height - len(lines) * line_h) // 2
    for i, l in enumerate(lines):
        w = draw.textlength(l, font=font)
        x = (size - w) / 2 + rng.randint(-20, 20)
        draw.text((x, y0 + i * line_h), l, fill=0, font=font)
    edges = cv2.Canny(np.array(canvas), 50, 150)
    edges = cv2.dilate(edges, np.ones((3, 3), np.uint8), iterations=1)
    return Image.fromarray(edges).convert("RGB")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--corpus", required=True)
    ap.add_argument("--out_dir", required=True)
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--n", type=int, default=1000)
    ap.add_argument("--scale", type=float, default=0.8)
    args = ap.parse_args()

    import torch
    from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionXLControlNetPipeline

    hub = "/home/jovyan/.cache/huggingface/hub"  # job pods have a different $HOME
    snap = lambda repo: sorted(os.listdir(f"{hub}/models--{repo.replace('/', '--')}/snapshots"))[0]  # noqa: E731
    base = f"{hub}/models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/{snap('stabilityai/stable-diffusion-xl-base-1.0')}"
    cn = f"{hub}/models--diffusers--controlnet-canny-sdxl-1.0/snapshots/{snap('diffusers/controlnet-canny-sdxl-1.0')}"
    vae = f"{hub}/models--madebyollin--sdxl-vae-fp16-fix/snapshots/{snap('madebyollin/sdxl-vae-fp16-fix')}"
    controlnet = ControlNetModel.from_pretrained(cn, torch_dtype=torch.float16, variant="fp16")
    pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
        base, controlnet=controlnet, vae=AutoencoderKL.from_pretrained(vae, torch_dtype=torch.float16),
        torch_dtype=torch.float16, variant="fp16").to("cuda")
    pipe.set_progress_bar_config(disable=True)

    rng = random.Random(1000 + args.shard)
    lex = lexicon(args.corpus)
    words = sorted(lex)
    weights = [lex[w] ** 0.7 for w in words]
    os.makedirs(os.path.join(args.out_dir, "img"), exist_ok=True)
    label_path = os.path.join(args.out_dir, f"labels_{args.shard}.jsonl")
    done = set()
    if os.path.exists(label_path):
        done = {json.loads(l)["id"] for l in open(label_path)}
    phrases = [sample_phrase(rng, words, weights) for _ in range(args.n)]
    with open(label_path, "a") as f:
        for i, t in enumerate(phrases):
            uid = f"dif{args.shard}_{i:05d}"  # v2 parameters (1216x832 canvas, scale 0.8)
            if uid in done:
                continue
            cond = render_condition(t, rng)
            g = torch.Generator("cuda").manual_seed(args.shard * 100000 + i)
            img = pipe(POSITIVE, negative_prompt=NEGATIVE, image=cond, width=cond.width, height=cond.height, num_inference_steps=28, guidance_scale=7.0,
                       controlnet_conditioning_scale=args.scale, generator=g).images[0]
            img.thumbnail((768, 768))
            name = f"img/{uid}.jpg"
            img.save(os.path.join(args.out_dir, name), quality=90)
            f.write(json.dumps({"id": uid, "image": name, "target": t}, ensure_ascii=False) + "\n")
            f.flush()
            if i % 50 == 0:
                print("generated", i, flush=True)
    print("DONE", args.shard, flush=True)


if __name__ == "__main__":
    main()