#!/usr/bin/env python3
"""CLIP-L/14 semantic near-duplicate leakage scan.

Reference: Ramos et al. ICCVW 2025 "Data Leakage in Visual Datasets"
Tüm 293K image'ı OpenCLIP ViT-L/14 ile embed et, cosine>0.95 çiftleri topla,
aynı canonical_tablet_id değilse "semantic leakage" işaretle.

Manifest alanı: clip_embedding_norm_id, clip_neighbor_ids (top-5)
Rapor: datasets/processed/clip_leakage_report.json
"""
import json, os, argparse, time
from pathlib import Path
import numpy as np

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

def main():
    import torch
    from PIL import Image
    import open_clip
    from concurrent.futures import ThreadPoolExecutor
    
    ap = argparse.ArgumentParser()
    ap.add_argument('--batch', type=int, default=128)
    ap.add_argument('--model', default='ViT-L-14')
    ap.add_argument('--pretrained', default='openai')
    ap.add_argument('--sample', type=int, default=0, help="0=tüm, N=sample")
    ap.add_argument('--out', default=str(ROOT / "datasets" / "processed" / "clip_leakage_report.json"))
    args = ap.parse_args()
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {device}", flush=True)
    
    model, _, preprocess = open_clip.create_model_and_transforms(args.model, pretrained=args.pretrained)
    model = model.eval().to(device)
    
    # Path topla
    items = []
    seen = set()
    for d in sorted(SOURCES.iterdir()):
        if not d.is_dir(): continue
        mp = d / "manifest_classification.jsonl"
        if not mp.exists(): continue
        with open(mp) as f:
            for line in f:
                r = json.loads(line)
                p = r.get('path')
                if (p and p not in seen and r.get('integrity_ok') is True 
                    and os.path.exists(p)):
                    seen.add(p)
                    items.append((r['id'], p, r.get('source',''), r.get('canonical_tablet_id')))
    
    if args.sample:
        import random
        random.seed(42)
        items = random.sample(items, min(args.sample, len(items)))
    print(f"Embed edilecek: {len(items):,}", flush=True)
    
    ids = []
    sources = []
    tablet_ids = []
    feats = []
    t0 = time.time()
    
    for bi in range(0, len(items), args.batch):
        batch = items[bi:bi+args.batch]
        imgs = []
        good_idx = []
        for i, (rid, path, src, tid) in enumerate(batch):
            try:
                img = preprocess(Image.open(path).convert('RGB'))
                imgs.append(img)
                good_idx.append(i)
            except Exception:
                continue
        if not imgs: continue
        with torch.no_grad(), torch.amp.autocast('cuda', enabled=(device=="cuda")):
            x = torch.stack(imgs).to(device)
            feat = model.encode_image(x)
            feat = feat / (feat.norm(dim=-1, keepdim=True) + 1e-8)
        feat = feat.cpu().numpy().astype(np.float32)
        for j, k in enumerate(good_idx):
            rid, path, src, tid = batch[k]
            ids.append(rid)
            sources.append(src)
            tablet_ids.append(tid)
        feats.append(feat)
        
        if (bi // args.batch) % 50 == 0:
            elapsed = time.time() - t0
            rate = (bi + len(batch)) / max(elapsed, 1)
            print(f"  {bi+len(batch)}/{len(items)} ({100*(bi+len(batch))/len(items):.1f}%) {rate:.0f} img/s", flush=True)
    
    embeddings = np.concatenate(feats, axis=0) if feats else np.zeros((0,768), dtype=np.float32)
    print(f"Embeddings shape: {embeddings.shape}", flush=True)
    
    # Embeddings kaydet (compressed)
    emb_out = ROOT / "datasets" / "processed" / "clip_embeddings.npz"
    np.savez_compressed(emb_out, embeddings=embeddings, ids=np.array(ids), sources=np.array(sources))
    print(f"Saved: {emb_out}", flush=True)
    
    # Leakage scan: batch-wise cosine matrix; memory için chunk-based
    # NumPy A @ A.T -> N×N (293K×293K = 320GB) — imkansız
    # Bunun yerine: FAISS kullan (approx NN)
    try:
        import faiss
        index = faiss.IndexFlatIP(embeddings.shape[1])
        index.add(embeddings)
        print("FAISS indexed", flush=True)
        
        k = 5  # top-5 nearest neighbors
        sims, idx = index.search(embeddings, k+1)  # +1 kendi
        
        leakage_threshold = 0.95
        leakage_pairs = []
        for i in range(len(embeddings)):
            for j, s in zip(idx[i][1:], sims[i][1:]):  # 0 kendi
                if s >= leakage_threshold:
                    if sources[i] != sources[j] or (tablet_ids[i] and tablet_ids[j] and tablet_ids[i] != tablet_ids[j]):
                        leakage_pairs.append({
                            "id1": ids[i], "source1": sources[i], "tablet1": tablet_ids[i],
                            "id2": ids[j], "source2": sources[j], "tablet2": tablet_ids[j],
                            "cosine": float(s)
                        })
        print(f"Leakage pairs (cosine>={leakage_threshold}): {len(leakage_pairs)}", flush=True)
        
        with open(args.out, 'w') as f:
            json.dump({
                "method": "OpenCLIP {}/{}".format(args.model, args.pretrained),
                "n_embeddings": len(embeddings),
                "leakage_threshold": leakage_threshold,
                "leakage_pairs_total": len(leakage_pairs),
                "leakage_pairs_sample": leakage_pairs[:500],
            }, f, indent=2, ensure_ascii=False)
    except ImportError:
        print("faiss yok; np.dot ile chunk-based hesap kullanılabilir", flush=True)

if __name__ == '__main__':
    main()