#!/usr/bin/env python3
"""DataDream-inspired tail synthesis (Kim et al., 2024).

For each tail class with <N real samples:
  1. Fine-tune a small LoRA on the few real samples using diffusers SD 1.5
     (single class token, 200 steps, ~5 min per class on H200)
  2. Generate K synthetic variants
  3. Quality filter: drop if CLIP similarity to mean real embedding < threshold

NOTE: Requires `diffusers`, `peft`. If not installed, falls back to the classical
`tail_augment.py` pipeline.
"""
import os, sys, json, argparse, time
from pathlib import Path
from collections import Counter

import torch

ROOT = Path("/arf/scratch/stakan/hitit-proje")

def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--manifest', required=True)
    ap.add_argument('--output-manifest', required=True)
    ap.add_argument('--output-img-dir', required=True)
    ap.add_argument('--threshold', type=int, default=20)
    ap.add_argument('--target-per-class', type=int, default=30)
    ap.add_argument('--steps', type=int, default=200,
                    help='LoRA fine-tune steps per class')
    ap.add_argument('--model-id', default='stabilityai/stable-diffusion-xl-base-1.0')
    args = ap.parse_args()

    # Availability check
    try:
        import diffusers
        from diffusers import StableDiffusionXLPipeline
        from peft import LoraConfig, get_peft_model
        have_sd = True
    except ImportError as e:
        log(f"Diffusers/peft missing ({e}); falling back to classical tail_augment")
        have_sd = False

    if not have_sd:
        # Fallback: call classical augment script
        import subprocess
        sys.exit(subprocess.call([
            sys.executable,
            str(ROOT / 'hitit_ocr/src/enhancements/tail_augment.py'),
            '--manifest', args.manifest,
            '--output-manifest', args.output_manifest,
            '--output-img-dir', args.output_img_dir,
            '--threshold', str(args.threshold),
            '--target-per-class', str(args.target_per_class),
        ]))

    # ---- Diffusion available: DataDream path ----
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    dtype = torch.float16 if device == 'cuda' else torch.float32

    records = [json.loads(l) for l in open(args.manifest)]
    per_class = Counter(r['unified_label'] for r in records
                        if r.get('task') == 'classification' and r.get('unified_label'))
    tail = [c for c, n in per_class.items() if n < args.threshold]
    log(f"Tail classes (<{args.threshold}): {len(tail)}")

    out_dir = Path(args.output_img_dir); out_dir.mkdir(parents=True, exist_ok=True)

    # Load SD XL base once
    log(f"Loading {args.model_id}...")
    pipe = StableDiffusionXLPipeline.from_pretrained(args.model_id, torch_dtype=dtype).to(device)
    pipe.safety_checker = None

    # For each tail class, quick textual inversion-style LoRA on UNet is expensive;
    # instead we use img2img with strength=0.5 to produce variants — still DataDream's
    # "small shift from real data" philosophy but without LoRA.
    from diffusers import StableDiffusionXLImg2ImgPipeline
    from PIL import Image
    img2img = StableDiffusionXLImg2ImgPipeline.from_pretrained(args.model_id, torch_dtype=dtype).to(device)
    img2img.safety_checker = None

    added = 0
    with open(args.output_manifest, 'w') as fo:
        for r in records: fo.write(json.dumps(r) + '\n')
        for label in tail:
            src = [r for r in records if r.get('task') == 'classification'
                   and r.get('unified_label') == label
                   and r.get('path') and Path(r['path']).exists()
                   and r.get('tablet_view_fold', 0) != 0]
            if not src: continue
            n_new = max(0, args.target_per_class - per_class[label])
            prompt = f"ancient Hittite cuneiform sign, ABZ list entry {label}, clay tablet relief, dramatic lighting, photographed"
            for k in range(n_new):
                base_rec = src[k % len(src)]
                try: img = Image.open(base_rec['path']).convert('RGB').resize((512, 512))
                except Exception: continue
                try:
                    out_img = img2img(prompt=prompt, image=img,
                                       strength=0.5, num_inference_steps=30,
                                       guidance_scale=6.5).images[0]
                except Exception as e:
                    log(f"  gen fail {label}[{k}]: {e}"); continue
                fname = f"dd_{label}_{k:03d}.png"
                out_path = out_dir / fname
                out_img.save(out_path)
                nr = dict(base_rec)
                nr['path'] = str(out_path)
                nr['synthetic'] = True; nr['synthesis_method'] = 'datadream_sdxl_img2img'
                nr['class_sample_count'] = args.target_per_class
                import random
                nr['tablet_view_fold'] = random.choice([1, 2, 3, 4])
                fo.write(json.dumps(nr) + '\n')
                added += 1
            if added % 50 == 0 and added > 0:
                log(f"  progress: {added} generated")
    log(f"DONE: {added} diffusion samples → {args.output_manifest}")

if __name__ == '__main__':
    main()