"""
Experiment: DINOv2-CLIP Hybrid vs Pure CLIP

Compares 4 approaches on the pre-computed EuroSAT gallery:
1. Pure CLIP (baseline)
2. CLIP + SAR Adapter
3. CLIP + DINOv2 patch features (hybrid)
4. Full hybrid (CLIP + SAR adapter + DINOv2)

Metrics: same-modal and cross-modal Recall@K, latency.
"""

import sys, time, json, traceback
import torch, numpy as np
from pathlib import Path
from PIL import Image
from dataclasses import dataclass, asdict
from typing import List, Optional

sys.path.insert(0, str(Path(__file__).parent))

DATA_DIR = Path("data")
PROCESSED_DIR = DATA_DIR / "processed"


@dataclass
class Result:
    model: str
    same_r1: float
    same_r5: float
    same_r10: float
    cross_r1: float
    cross_r5: float
    cross_r10: float
    latency_ms: float
    n_queries: int


def load_data():
    embeddings = torch.load(PROCESSED_DIR / "gallery_embeddings.pt", weights_only=True)
    with open(PROCESSED_DIR / "gallery_metadata.json") as f:
        metadata = json.load(f)
    return embeddings.numpy().astype(np.float32), metadata


def split(metadata):
    """Stratified 30/70 split: 30% of each (modality, class) pair goes to queries."""
    groups = {}
    for e in metadata:
        key = (e["modality"], e["class"])
        groups.setdefault(key, []).append(e)

    queries, gallery = [], []
    for key, entries in groups.items():
        n = max(1, int(len(entries) * 0.3))
        queries.extend(entries[:n])
        gallery.extend(entries[n:])
    return queries, gallery


def recall_at_k(retrieved, query_mod, query_class, metadata, k, mode="same"):
    hits = 0
    for idx in retrieved[:k]:
        m = metadata[idx]
        same_class = m["class"] == query_class
        same_mod = m["modality"] == query_mod
        if mode == "same" and same_class and same_mod:
            hits += 1
        elif mode == "cross" and same_class and not same_mod:
            hits += 1
    return hits


def evaluate(queries, all_emb, metadata, gallery_entries, extractor_fn, label):
    import faiss

    gal_idx = [e["index"] for e in gallery_entries]
    gal_emb = all_emb[gal_idx]
    dim = gal_emb.shape[1]
    index = faiss.IndexFlatIP(dim)
    index.add(gal_emb)

    sr1, sr5, sr10 = [], [], []
    cr1, cr5, cr10 = [], [], []
    latencies = []

    for q in queries:
        q_path = Path(q["gallery_path"])
        if not q_path.exists():
            continue
        img = Image.open(q_path).convert("RGB")

        start = time.perf_counter()
        try:
            emb = extractor_fn(img, q["modality"])
        except Exception:
            continue
        elapsed = (time.perf_counter() - start) * 1000
        latencies.append(elapsed)

        q_np = emb.reshape(1, -1).astype(np.float32)
        _, ids = index.search(q_np, 10)
        retrieved = [gal_idx[i] for i in ids[0] if 0 <= i < len(gal_idx)]

        sr1.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 1, "same"))
        sr5.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 5, "same"))
        sr10.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 10, "same"))
        cr1.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 1, "cross"))
        cr5.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 5, "cross"))
        cr10.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 10, "cross"))

    n = max(len(sr1), 1)
    return Result(
        model=label,
        same_r1=np.mean(sr1) / 1.0,
        same_r5=np.mean(sr5) / 5.0,
        same_r10=np.mean(sr10) / 10.0,
        cross_r1=np.mean(cr1) / 1.0,
        cross_r5=np.mean(cr5) / 5.0,
        cross_r10=np.mean(cr10) / 10.0,
        latency_ms=np.mean(latencies) if latencies else 0,
        n_queries=n,
    )


def main():
    print("=" * 72)
    print("  EXPERIMENT: DINOv2-CLIP Hybrid vs Pure CLIP")
    print("=" * 72)

    all_emb, metadata = load_data()
    queries, gallery = split(metadata)
    print(f"Gallery: {len(gallery)} | Queries: {len(queries)} | Dim: {all_emb.shape[1]}")

    from transformers import CLIPProcessor, CLIPModel
    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
    clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(device)
    clip_model.eval()
    print(f"CLIP loaded on {device}")

    @torch.no_grad()
    def clip_extract(img, modality):
        inputs = processor(images=img, return_tensors="pt").to(device)
        out = clip_model.vision_model(**inputs)
        pooled = out.last_hidden_state[:, 0, :]
        feat = clip_model.visual_projection(pooled).squeeze(0)
        return torch.nn.functional.normalize(feat, dim=-1).cpu().numpy()

    results = []

    print("\n[1/4] Pure CLIP ...")
    r = evaluate(queries, all_emb, metadata, gallery, clip_extract, "CLIP ViT-L/14")
    results.append(r)
    print(f"  Same R@5={r.same_r5:.4f} Cross R@5={r.cross_r5:.4f} Latency={r.latency_ms:.0f}ms")

    print("[2/4] CLIP + SAR Adapter ...")
    from src.features.sar_adapter import SARAdapter
    adapter = SARAdapter().eval()

    def clip_sar_extract(img, modality):
        if modality == "sar":
            arr = np.array(img).astype(np.float32) / 255.0
            t = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
            with torch.no_grad():
                adapted = adapter(t)
            img = Image.fromarray((adapted.squeeze(0).permute(1, 2, 0).numpy() * 255).clip(0, 255).astype(np.uint8))
        return clip_extract(img, modality)

    r = evaluate(queries, all_emb, metadata, gallery, clip_sar_extract, "CLIP + SAR Adapter")
    results.append(r)
    print(f"  Same R@5={r.same_r5:.4f} Cross R@5={r.cross_r5:.4f} Latency={r.latency_ms:.0f}ms")

    print("[3/4] CLIP + DINOv2 Hybrid ...")
    try:
        dinov2 = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14", pretrained=True)
        dinov2.to(device).eval()
        has_dino = True
        dino_embed_dim = dinov2.embed_dim  # 384 for vits14
        print(f"  DINOv2-ViT-S/14 loaded (embed_dim={dino_embed_dim})")
    except Exception as e:
        has_dino = False
        print(f"  DINOv2 load failed: {e}")

    if has_dino:
        from torchvision import transforms
        dino_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

        # Projection to match CLIP dim (768) if needed
        dino_proj = None
        if dino_embed_dim != 768:
            dino_proj = torch.nn.Linear(dino_embed_dim, 768, bias=False).to(device).eval()
            with torch.no_grad():
                torch.nn.init.eye_(dino_proj.weight)  # identity init — preserves features

        @torch.no_grad()
        def clip_dino_extract(img, modality):
            clip_feat = clip_extract(img, modality)
            t = dino_transform(img).unsqueeze(0).to(device)
            patch_feat = dinov2(t).squeeze(0)
            if dino_proj is not None:
                patch_feat = dino_proj(patch_feat)
            patch_feat = torch.nn.functional.normalize(patch_feat, dim=-1).cpu().numpy()
            hybrid = 0.7 * clip_feat + 0.3 * patch_feat
            return hybrid / (np.linalg.norm(hybrid) + 1e-8)

        r = evaluate(queries, all_emb, metadata, gallery, clip_dino_extract, "DINOv2-CLIP Hybrid")
        results.append(r)
        print(f"  Same R@5={r.same_r5:.4f} Cross R@5={r.cross_r5:.4f} Latency={r.latency_ms:.0f}ms")
    else:
        r = evaluate(queries, all_emb, metadata, gallery, clip_extract, "CLIP (DINOv2 unavailable)")
        results.append(r)

    print("[4/4] Full Hybrid (CLIP + SAR + DINOv2) ...")
    if has_dino:
        def full_extract(img, modality):
            if modality == "sar":
                arr = np.array(img).astype(np.float32) / 255.0
                t = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
                with torch.no_grad():
                    adapted = adapter(t)
                img = Image.fromarray((adapted.squeeze(0).permute(1, 2, 0).numpy() * 255).clip(0, 255).astype(np.uint8))
            return clip_dino_extract(img, modality)

        r = evaluate(queries, all_emb, metadata, gallery, full_extract, "Full Hybrid (CLIP+SAR+DINOv2)")
        results.append(r)
        print(f"  Same R@5={r.same_r5:.4f} Cross R@5={r.cross_r5:.4f} Latency={r.latency_ms:.0f}ms")
    else:
        r = evaluate(queries, all_emb, metadata, gallery, clip_sar_extract, "CLIP+SAR (DINOv2 unavailable)")
        results.append(r)

    print("\n" + "=" * 72)
    print("  RESULTS")
    print("=" * 72)
    hdr = f"{'Model':<35} {'S-R@1':>6} {'S-R@5':>6} {'S-R@10':>7} {'C-R@1':>6} {'C-R@5':>6} {'C-R@10':>7} {'ms':>6}"
    print(hdr)
    print("-" * 72)
    for r in results:
        print(f"{r.model:<35} {r.same_r1:>6.4f} {r.same_r5:>6.4f} {r.same_r10:>7.4f} {r.cross_r1:>6.4f} {r.cross_r5:>6.4f} {r.cross_r10:>7.4f} {r.latency_ms:>5.0f}")

    base_s5 = results[0].same_r5
    base_c5 = results[0].cross_r5
    print(f"\nDelta vs CLIP baseline (R@5):")
    for r in results[1:]:
        ds = r.same_r5 - base_s5
        dc = r.cross_r5 - base_c5
        print(f"  {r.model}: Same {'+' if ds >= 0 else ''}{ds:.4f}, Cross {'+' if dc >= 0 else ''}{dc:.4f}")

    out = PROCESSED_DIR / "experiment_results.json"
    with open(out, "w") as f:
        json.dump([asdict(r) for r in results], f, indent=2)
    print(f"\nSaved to {out}")


if __name__ == "__main__":
    try:
        main()
    except Exception:
        traceback.print_exc()