#!/usr/bin/env python3
"""SAM 2.1 polygon annotation — GPU run.
Detection bbox'larını prompt olarak kullanıp RLE mask üretir.
Output: her kaynağın detection.coco.json dosyasına 'segmentation' alanı ekler.
"""
import json, os, argparse, time
from pathlib import Path

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

def main():
    import torch
    from PIL import Image
    import numpy as np
    
    ap = argparse.ArgumentParser()
    ap.add_argument('--source', default='all', help='all|hitit_local|compvis|yeni_veri|deepscribe')
    ap.add_argument('--limit', type=int, default=0)
    args = ap.parse_args()
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {device}", flush=True)
    
    # SAM 2.1 import (ultralytics YOLO'dan sam kullanabiliriz; veya segment_anything_2 paketi)
    try:
        from ultralytics import SAM
        sam = SAM("sam2.1_l.pt")  # auto-download
        model_name = "sam2.1_l (ultralytics)"
    except Exception as e:
        print(f"Ultralytics SAM yüklenmedi: {e}; alternatif deneniyor", flush=True)
        try:
            from segment_anything import SamPredictor, sam_model_registry
            print("segment_anything varsa — manual checkpoint path gerekir")
            return
        except ImportError:
            print("SAM kurulu değil. Install: pip install ultralytics", flush=True)
            return
    
    source_dirs = ['hitit_local', 'compvis', 'yeni_veri', 'deepscribe'] if args.source == 'all' else [args.source]
    
    try:
        from pycocotools import mask as mask_utils
    except ImportError:
        os.system("pip install --user --quiet pycocotools")
        from pycocotools import mask as mask_utils
    
    t0 = time.time()
    total_done = 0
    
    for src_name in source_dirs:
        src = SOURCES / src_name
        mp = src / "manifest_detection.jsonl"
        if not mp.exists(): continue
        
        records = []
        with open(mp) as f:
            for line in f:
                records.append(json.loads(line))
        if args.limit:
            records = records[:args.limit]
        
        print(f"\n{src_name}: {len(records)} record", flush=True)
        
        out_json = src / "detection.coco.json"
        coco = json.load(open(out_json)) if out_json.exists() else {
            "info": {"description": f"{src_name} detection with SAM 2.1 masks"},
            "images": [], "annotations": [], "categories": []
        }
        
        ann_id = 1 + (max([a['id'] for a in coco.get('annotations', [])], default=0))
        cat_map = {c['name']: c['id'] for c in coco.get('categories', [])}
        img_id = 1 + (max([i['id'] for i in coco.get('images', [])], default=0))
        existing_img_paths = {i.get('file_name') for i in coco.get('images', [])}
        
        for ri, r in enumerate(records):
            p = r.get('path')
            if not p or not os.path.exists(p): continue
            if p in existing_img_paths: continue
            
            # Bbox topla (YOLO label_path veya extra.bboxes)
            bboxes_xyxy = []
            labels = []
            w = r.get('width') or 0
            h = r.get('height') or 0
            if not (w and h):
                try:
                    with Image.open(p) as img:
                        w, h = img.size
                except: continue
            
            lp = r.get('label_path')
            if lp and os.path.exists(lp):
                try:
                    with open(lp) as f:
                        for ln in f:
                            parts = ln.strip().split()
                            if len(parts) == 5:
                                cls, cx, cy, bw, bh = parts
                                cx, cy, bw, bh = float(cx)*w, float(cy)*h, float(bw)*w, float(bh)*h
                                x1, y1, x2, y2 = cx-bw/2, cy-bh/2, cx+bw/2, cy+bh/2
                                bboxes_xyxy.append([x1, y1, x2, y2])
                                labels.append(f"class_{cls}")
                except: pass
            extra = r.get('extra') or {}
            if isinstance(extra, dict):
                for b in extra.get('bboxes', []):
                    if not isinstance(b, dict): continue
                    lbl = b.get('class_name') or b.get('mzl_label') or b.get('train_label') or 'unknown'
                    yolo = b.get('yolo_bbox')
                    xyxy = b.get('bbox')
                    if yolo and len(yolo) == 4:
                        cx, cy, bw, bh = yolo
                        cx, cy, bw, bh = cx*w, cy*h, bw*w, bh*h
                        bboxes_xyxy.append([cx-bw/2, cy-bh/2, cx+bw/2, cy+bh/2])
                        labels.append(str(lbl))
                    elif xyxy and len(xyxy) == 4:
                        bboxes_xyxy.append([float(x) for x in xyxy])
                        labels.append(str(lbl))
            
            if not bboxes_xyxy: continue
            
            # SAM predict
            try:
                results = sam(p, bboxes=bboxes_xyxy, verbose=False)
            except Exception as e:
                print(f"  skip {p}: {e}", flush=True)
                continue
            
            # Image entry
            coco['images'].append({
                "id": img_id,
                "file_name": p,
                "width": w, "height": h,
                "fold": r.get('fold'),
                "tablet_id": r.get('tablet_id'),
            })
            
            # Annotations
            if results and hasattr(results[0], 'masks') and results[0].masks is not None:
                masks = results[0].masks.data.cpu().numpy()  # (N, H, W) bool
                for mi, (mask, lbl, bb) in enumerate(zip(masks, labels, bboxes_xyxy)):
                    if lbl not in cat_map:
                        cat_map[lbl] = len(cat_map) + 1
                        coco['categories'].append({"id": cat_map[lbl], "name": lbl, "supercategory": "sign"})
                    # RLE encode
                    mask_u8 = (mask > 0).astype(np.uint8)
                    rle = mask_utils.encode(np.asfortranarray(mask_u8))
                    rle['counts'] = rle['counts'].decode('utf-8')
                    x1, y1, x2, y2 = bb
                    coco['annotations'].append({
                        "id": ann_id,
                        "image_id": img_id,
                        "category_id": cat_map[lbl],
                        "bbox": [x1, y1, x2-x1, y2-y1],
                        "area": float(mask_u8.sum()),
                        "iscrowd": 0,
                        "segmentation": rle,
                    })
                    ann_id += 1
            img_id += 1
            total_done += 1
            
            if (ri+1) % 100 == 0:
                elapsed = time.time() - t0
                rate = total_done / max(elapsed, 1)
                print(f"  [{src_name}] {ri+1}/{len(records)} ({rate:.1f} img/s)", flush=True)
        
        with open(out_json, 'w') as f:
            json.dump(coco, f, ensure_ascii=False)
        print(f"  {src_name}: yazıldı {out_json} — {len(coco['images'])} img, {len(coco['annotations'])} ann", flush=True)

if __name__ == '__main__':
    main()