#!/usr/bin/env python3
"""
Extract DepthAnythingV2 depth maps for cuneiform crops.
Saves as uint8 grayscale PNGs alongside each crop under depth/ subfolder.

HF model: depth-anything/Depth-Anything-V2-Small-hf (fastest).
"""
import argparse
import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image


def load_depth_model(device='cuda'):
    from transformers import AutoImageProcessor, AutoModelForDepthEstimation
    proc = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
    model = AutoModelForDepthEstimation.from_pretrained(
        "depth-anything/Depth-Anything-V2-Small-hf", torch_dtype=torch.float16
    ).to(device).eval()
    return proc, model


def depth_of(img, proc, model, device='cuda'):
    inputs = proc(images=img, return_tensors="pt").to(device)
    inputs["pixel_values"] = inputs["pixel_values"].to(torch.float16)
    with torch.no_grad():
        out = model(**inputs)
    pred = out.predicted_depth  # [1, H, W]
    pred = torch.nn.functional.interpolate(
        pred.unsqueeze(1), size=img.size[::-1], mode="bicubic", align_corners=False
    )[0, 0]
    arr = pred.float().cpu().numpy()
    # Normalize to 0-255
    arr = arr - arr.min()
    if arr.max() > 0:
        arr = arr / arr.max() * 255
    return arr.astype(np.uint8)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--manifest', required=True)
    ap.add_argument('--output-dir', default='/arf/scratch/stakan/hitit-proje/hitit_ocr/data/classification/depth')
    ap.add_argument('--batch', type=int, default=1)
    ap.add_argument('--limit', type=int, default=0)
    args = ap.parse_args()

    proc, model = load_depth_model()
    out_dir = Path(args.output_dir); out_dir.mkdir(parents=True, exist_ok=True)

    seen = set()
    n = 0
    with open(args.manifest) as f:
        for line in f:
            r = json.loads(line)
            p = r.get('path')
            if not p or p in seen:
                continue
            seen.add(p)
            # Mirror relative path if under classification/all, else flatten with hash prefix
            try:
                rel = Path(p).relative_to('/arf/scratch/stakan/hitit-proje/hitit_ocr/data/classification/all')
                out_p = out_dir / rel
            except ValueError:
                # external source: flatten with source tag
                parts = Path(p).parts
                key = '_'.join(parts[-3:])
                out_p = out_dir / '_external' / key
            out_p = out_p.with_suffix('.png')
            if out_p.exists():
                n += 1
                continue
            out_p.parent.mkdir(parents=True, exist_ok=True)
            try:
                img = Image.open(p).convert('RGB')
                d = depth_of(img, proc, model)
                Image.fromarray(d).save(out_p)
            except Exception as e:
                print(f"fail {p}: {e}")
                continue
            n += 1
            if n % 500 == 0:
                print(f"  {n} depth maps")
            if args.limit and n >= args.limit:
                break
    print(f"DONE: {n} depth maps → {out_dir}")


if __name__ == '__main__':
    main()