# synthetic OCR samples: render text with PIL, return (PIL image, messages)
import random, string, os, glob, json, re
from PIL import Image, ImageDraw, ImageFont, ImageFilter
FONTS=[f for f in glob.glob("/usr/share/fonts/truetype/**/*.ttf",recursive=True) if re.search(r"DejaVu(Sans|Serif|SansMono)(-Bold)?\.ttf|Liberation(Sans|Serif|Mono)-(Regular|Bold)\.ttf|Ubuntu(Mono)?\[wght\]|NotoSansMono-Regular|UbuntuSans\[",f)]
assert FONTS, "no fonts"
PROMPTS=["What does the text in the image say?","Read the text in this image.","OCR this image.","Transcribe the text in the image.","What is written in this image?","Extract the text from the image.","Please read this image."]
_sent=None
def sentences():
    global _sent
    if _sent is None:
        _sent=[]
        for l in open("data/rendered/smoltalk_eval.jsonl"):
            t=json.loads(l)["text"]
            for m in re.finditer(r"<\|im_start\|>(?:user|assistant)\n(.*?)<\|im_end\|>",t,re.S):
                for s in re.split(r"(?<=[.!?])\s+|\n+",m.group(1)):
                    s=s.strip()
                    if 3<=len(s)<=90 and s.isprintable() and s.isascii() and not s.startswith(("<","#","```","-","*","|")): _sent.append(s)
            if len(_sent)>200000: break
    return _sent
def rand_word(rng,n): return "".join(rng.choice(string.ascii_letters+string.digits) for _ in range(n))
def rand_text(rng):
    k=rng.random()
    if k<0.45: return rng.choice(sentences())
    if k<0.60:  # multi-line paragraph
        return "\n".join(rng.choice(sentences()) for _ in range(rng.randint(2,4)))
    if k<0.70: return str(rng.randint(0,10**rng.randint(1,9)))
    if k<0.78: return f"{rng.randint(1,12):02d}/{rng.randint(1,28):02d}/{rng.randint(1990,2030)}"
    if k<0.85: return (rand_word(rng,rng.randint(4,12)) if rng.random()<0.5 else " ".join(rng.choice(sentences()).split()[:rng.randint(1,3)])).upper()
    if k<0.92: return f"{rand_word(rng,rng.randint(3,8)).lower()}@{rng.choice(['example.com','mail.org','test.net'])}"
    if k<0.94: return " ".join(rng.choice(sentences()).split()[:rng.randint(1,4)])
    if k<0.98: return rng.choice([f"{rand_word(rng,rng.randint(3,6)).lower()} {rng.randint(1000,2999)}",f"{rand_word(rng,rng.randint(2,5)).upper()}-{rng.randint(10,999)}",f"{rand_word(rng,rng.randint(3,6))}{rng.randint(1,99)}"])
    return f"${rng.randint(1,9999)}.{rng.randint(0,99):02d}"
def render(rng,text):
    fpath=rng.choice(FONTS); size=rng.randint(16,40)
    lines=text.split("\n")
    pad=rng.randint(4,24)
    dummy=ImageDraw.Draw(Image.new("RGB",(1,1)))
    while True:
        font=ImageFont.truetype(fpath,size)
        ws=[dummy.textbbox((0,0),l,font=font)[2] for l in lines]; hs=[dummy.textbbox((0,0),l,font=font)[3] for l in lines]
        if max(ws)+2*pad<=768 or size<=10: break
        size-=2
    lh=max(hs)+rng.randint(0,8); W=max(ws)+2*pad; H=lh*len(lines)+2*pad
    W=min(max(W,32),768); H=min(max(H,32),512)
    bg=tuple(rng.randint(200,255) for _ in range(3)) if rng.random()<0.8 else tuple(rng.randint(0,60) for _ in range(3))
    fg=tuple(rng.randint(0,70) for _ in range(3)) if sum(bg)>380 else tuple(rng.randint(190,255) for _ in range(3))
    img=Image.new("RGB",(W,H),bg); d=ImageDraw.Draw(img)
    for i,l in enumerate(lines): d.text((pad,pad+i*lh),l,font=font,fill=fg)
    if rng.random()<0.3: img=img.filter(ImageFilter.GaussianBlur(rng.uniform(0.3,1.0)))
    if rng.random()<0.2: img=img.rotate(rng.uniform(-3,3),expand=False,fillcolor=bg)
    return img
def make_sample(rng):
    text=rand_text(rng)
    img=render(rng,text)
    parts=[{"type":"image"},{"type":"text","text":rng.choice(PROMPTS)}]
    if rng.random()<0.5: parts.reverse()
    msgs=[{"role":"user","content":parts},{"role":"assistant","content":text}]
    return img,msgs
if __name__=="__main__":
    rng=random.Random(0); os.makedirs("ocr_samples",exist_ok=True)
    for i in range(8):
        img,m=make_sample(rng); img.save(f"ocr_samples/{i}.png"); print(img.size,repr(m[1]["content"]))
    print(len(FONTS),"fonts")

# ---- preprocessing that mirrors llama.cpp (mtmd) for Qwen3-VL: smart_resize target + align-corners bilinear
import numpy as np
from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
def lcpp_bilinear(img, W, H):
    src=np.asarray(img.convert("RGB")).astype(np.float32); h,w=src.shape[:2]
    xr=(w-1)/(W-1) if W>1 else 0.0; yr=(h-1)/(H-1) if H>1 else 0.0
    px=np.arange(W)*xr; py=np.arange(H)*yr
    x0=np.minimum(px.astype(np.int64),w-1); y0=np.minimum(py.astype(np.int64),h-1)
    x1=np.minimum(x0+1,w-1); y1=np.minimum(y0+1,h-1)
    xf=(px-x0)[None,:,None]; yf=(py-y0)[:,None,None]
    out=(src[y0[:,None],x0[None,:]]*(1-xf)*(1-yf)+src[y0[:,None],x1[None,:]]*xf*(1-yf)+src[y1[:,None],x0[None,:]]*(1-xf)*yf+src[y1[:,None],x1[None,:]]*xf*yf)
    return Image.fromarray(np.clip(np.round(out),0,255).astype(np.uint8))
def target_size(img, min_pixels=8*1024, max_pixels=4096*1024):
    h,w=smart_resize(img.height,img.width,factor=32,min_pixels=min_pixels,max_pixels=max_pixels); return w,h
def letterbox_lcpp(img, W, H):
    # llama.cpp (mtmd, PAD_CEIL): aspect-preserving bilinear resize, centered on a black canvas
    import math
    img=img.convert("RGB"); sc=min(W/img.width,H/img.height)
    nw=min(math.ceil(img.width*sc),W); nh=min(math.ceil(img.height*sc),H)
    canvas=Image.new("RGB",(W,H),(0,0,0)); canvas.paste(lcpp_bilinear(img,nw,nh),((W-nw)//2,(H-nh)//2)); return canvas
def preprocess(img, ip, rng=None, return_tensors="np"):
    # rng=None -> exact llama.cpp preprocessing; with rng -> 75% llama.cpp style, 25% HF-style stretch with random kernels
    W,H=target_size(img)
    if rng is not None and rng.random()<0.25:
        img=img.convert("RGB").resize((W,H),rng.choice([Image.BILINEAR,Image.BICUBIC,Image.NEAREST,Image.LANCZOS]))
    else:
        img=letterbox_lcpp(img,W,H)
    return ip(images=[img],do_resize=False,return_tensors=return_tensors)