"""
Extract 5 style reference images per writer from the full dataset.
Output: style_refs/{writer_id}/{filename}.png
"""
import csv
import os
import shutil

META = "/Users/admin/Downloads/UkrHandwritten_Words_CC_ULCleanNAF_v1/METAFILE_extended_trocr_local3_balanced_20260421_181218.tsv"
WORDS_DIR = "/Users/admin/Downloads/UkrHandwritten_Words_CC_ULCleanNAF_v1/words/words"
OUT_DIR = os.path.join(os.path.dirname(__file__), "style_refs")
N_REFS = 5

writer_images: dict[str, list[str]] = {}
with open(META, encoding="utf-8") as f:
    for row in csv.reader(f, delimiter="\t"):
        if len(row) < 2:
            continue
        fname = row[0]
        parts = fname.replace(".png", "").split("-")
        if len(parts) >= 3:
            wid = parts[2]
            writer_images.setdefault(wid, []).append(fname)

print(f"Writers found: {len(writer_images)}")

copied = skipped = 0
for wid, fnames in sorted(writer_images.items()):
    out_writer_dir = os.path.join(OUT_DIR, wid)
    os.makedirs(out_writer_dir, exist_ok=True)
    for fname in fnames[:N_REFS]:
        src = os.path.join(WORDS_DIR, fname)
        dst = os.path.join(out_writer_dir, fname)
        if os.path.exists(src):
            shutil.copy2(src, dst)
            copied += 1
        else:
            skipped += 1

print(f"Copied: {copied}  Skipped (missing): {skipped}")
print(f"Output: {OUT_DIR}")