"""
Benchmark cross-modal retrieval approaches.

Tests:
1. Single-index with modality filtering
2. Multi-index search
3. Hybrid search
4. Cross-modal alignment with projection heads
"""

import torch
import numpy as np
import time
from pathlib import Path
from typing import Dict, List, Tuple

# Add src to path
import sys
sys.path.insert(0, str(Path(__file__).parent))

from src.features.cross_modal import CrossModalAligner, CrossModalConfig
from src.retrieval.cross_modal_retrieval import CrossModalRetrieval


def load_gallery_data():
    """Load real gallery embeddings and metadata."""
    import json
    
    data_dir = Path("data/processed")
    
    # Load embeddings
    embeddings = torch.load(data_dir / "gallery_embeddings.pt", weights_only=True)
    
    # Load metadata
    with open(data_dir / "gallery_metadata.json") as f:
        metadata = json.load(f)
    
    return embeddings.numpy().astype(np.float32), metadata


def split_by_modality(embeddings: np.ndarray, metadata: List[dict]) -> Dict[str, np.ndarray]:
    """Split embeddings by modality."""
    modalities = {}
    for i, entry in enumerate(metadata):
        mod = entry["modality"]
        if mod not in modalities:
            modalities[mod] = []
        modalities[mod].append(embeddings[i])
    
    return {mod: np.array(embs) for mod, embs in modalities.items()}


def compute_recall_at_k(retrieved_indices: List[int], ground_truth_idx: int, k: int) -> float:
    """Compute Recall@K."""
    if ground_truth_idx in retrieved_indices[:k]:
        return 1.0
    return 0.0


def benchmark_single_index(
    embeddings: np.ndarray,
    metadata: List[dict],
    n_queries: int = 50,
    k: int = 5
) -> Dict:
    """Benchmark single-index approach."""
    print("\n=== Single-Index Benchmark ===")
    
    retrieval = CrossModalRetrieval(embed_dim=embeddings.shape[1])
    retrieval.build_single_index(embeddings, [m["modality"] for m in metadata], metadata)
    
    # Generate queries (use gallery images as queries)
    query_indices = np.random.choice(len(embeddings), n_queries, replace=False)
    
    results = {
        "same_modal": [],
        "cross_modal": [],
    }
    
    for idx in query_indices:
        query = embeddings[idx:idx+1]
        query_mod = metadata[idx]["modality"]
        query_class = metadata[idx]["class"]
        
        # Same-modal search
        same_result = retrieval.search(query, query_mod, target_modality=query_mod, k=k)
        same_recall = compute_recall_at_k(
            [metadata[i]["class"] for i in same_result.indices],
            query_class,
            k
        )
        results["same_modal"].append(same_recall)
        
        # Cross-modal search (find different modality with same class)
        cross_targets = [m for m in ["optical", "sar", "multispectral"] if m != query_mod]
        for target_mod in cross_targets:
            cross_result = retrieval.search(query, query_mod, target_mod, k=k)
            cross_recall = compute_recall_at_k(
                [metadata[i]["class"] for i in cross_result.indices],
                query_class,
                k
            )
            results["cross_modal"].append(cross_recall)
    
    return {
        "same_modal_recall@k": np.mean(results["same_modal"]),
        "cross_modal_recall@k": np.mean(results["cross_modal"]),
    }


def benchmark_multi_index(
    embeddings_by_mod: Dict[str, np.ndarray],
    metadata: List[dict],
    n_queries: int = 50,
    k: int = 5
) -> Dict:
    """Benchmark multi-index approach."""
    print("\n=== Multi-Index Benchmark ===")
    
    # Build metadata by modality
    metadata_by_mod = {}
    for entry in metadata:
        mod = entry["modality"]
        if mod not in metadata_by_mod:
            metadata_by_mod[mod] = []
        metadata_by_mod[mod].append(entry)
    
    retrieval = CrossModalRetrieval(embed_dim=768)
    retrieval.build_multi_index(embeddings_by_mod, metadata_by_mod)
    
    results = {
        "same_modal": [],
        "cross_modal": [],
    }
    
    # Generate queries
    all_embeddings = np.concatenate(list(embeddings_by_mod.values()))
    query_indices = np.random.choice(len(all_embeddings), n_queries, replace=False)
    
    for idx in query_indices:
        query = all_embeddings[idx:idx+1]
        
        # Determine query modality
        offset = 0
        query_mod = None
        for mod, embs in embeddings_by_mod.items():
            if idx < offset + len(embs):
                query_mod = mod
                break
            offset += len(embs)
        
        if query_mod is None:
            continue
        
        # Same-modal search
        same_result = retrieval.search(query, query_mod, target_modality=query_mod, k=k)
        same_recall = 1.0 if any(True for _ in same_result.indices) else 0.0
        results["same_modal"].append(same_recall)
        
        # Cross-modal search
        cross_targets = [m for m in embeddings_by_mod.keys() if m != query_mod]
        cross_result = retrieval.search(query, query_mod, k=k)
        cross_recall = 1.0 if any(True for _ in cross_result.indices) else 0.0
        results["cross_modal"].append(cross_recall)
    
    return {
        "same_modal_recall@k": np.mean(results["same_modal"]),
        "cross_modal_recall@k": np.mean(results["cross_modal"]),
    }


def benchmark_hybrid(
    embeddings_by_mod: Dict[str, np.ndarray],
    metadata: List[dict],
    n_queries: int = 50,
    k: int = 5
) -> Dict:
    """Benchmark hybrid search approach."""
    print("\n=== Hybrid Search Benchmark ===")
    
    metadata_by_mod = {}
    for entry in metadata:
        mod = entry["modality"]
        if mod not in metadata_by_mod:
            metadata_by_mod[mod] = []
        metadata_by_mod[mod].append(entry)
    
    retrieval = CrossModalRetrieval(embed_dim=768)
    retrieval.build_multi_index(embeddings_by_mod, metadata_by_mod)
    
    results = []
    
    all_embeddings = np.concatenate(list(embeddings_by_mod.values()))
    query_indices = np.random.choice(len(all_embeddings), n_queries, replace=False)
    
    for idx in query_indices:
        query = all_embeddings[idx:idx+1]
        
        offset = 0
        query_mod = None
        for mod, embs in embeddings_by_mod.items():
            if idx < offset + len(embs):
                query_mod = mod
                break
            offset += len(embs)
        
        if query_mod is None:
            continue
        
        result = retrieval.search_hybrid(query, query_mod, k=k)
        results.append(1.0 if len(result.indices) > 0 else 0.0)
    
    return {
        "hybrid_recall@k": np.mean(results),
    }


def benchmark_cross_modal_alignment(
    embeddings: np.ndarray,
    metadata: List[dict],
    n_queries: int = 50,
    k: int = 5
) -> Dict:
    """Benchmark cross-modal alignment with projection heads."""
    print("\n=== Cross-Modal Alignment Benchmark ===")
    
    config = CrossModalConfig(
        embed_dim=embeddings.shape[1],
        projection_dim=256,
        use_wavelength_encoding=True,
        use_domain_adaptation=True,
    )
    
    aligner = CrossModalAligner(config)
    
    # Project all embeddings
    projected = {}
    for mod in ["optical", "sar", "multispectral"]:
        mask = [m["modality"] == mod for m in metadata]
        mod_embeddings = embeddings[mask]
        projected[mod] = aligner.project(
            torch.tensor(mod_embeddings), mod
        ).detach().numpy()
    
    results = {
        "same_modal": [],
        "cross_modal": [],
    }
    
    query_indices = np.random.choice(len(embeddings), n_queries, replace=False)
    
    for idx in query_indices:
        query = embeddings[idx:idx+1]
        query_mod = metadata[idx]["modality"]
        query_class = metadata[idx]["class"]
        
        # Project query
        query_proj = aligner.project(
            torch.tensor(query), query_mod
        ).detach().numpy()
        
        # Same-modal search
        same_proj = projected[query_mod]
        similarities = query_proj @ same_proj.T
        topk_idx = np.argsort(similarities[0])[::-1][:k]
        
        same_recall = compute_recall_at_k(
            [metadata[i]["class"] for i in topk_idx],
            query_class,
            k
        )
        results["same_modal"].append(same_recall)
        
        # Cross-modal search
        for target_mod in ["optical", "sar", "multispectral"]:
            if target_mod == query_mod:
                continue
            
            target_proj = projected[target_mod]
            similarities = query_proj @ target_proj.T
            topk_idx = np.argsort(similarities[0])[::-1][:k]
            
            cross_recall = compute_recall_at_k(
                [metadata[i]["class"] for i in topk_idx],
                query_class,
                k
            )
            results["cross_modal"].append(cross_recall)
    
    return {
        "same_modal_recall@k": np.mean(results["same_modal"]),
        "cross_modal_recall@k": np.mean(results["cross_modal"]),
    }


def main():
    """Run all benchmarks."""
    print("=" * 60)
    print("Cross-Modal Retrieval Benchmark")
    print("=" * 60)
    
    # Load data
    print("\nLoading gallery data...")
    embeddings, metadata = load_gallery_data()
    print(f"Loaded {len(metadata)} embeddings of dimension {embeddings.shape[1]}")
    
    # Split by modality
    embeddings_by_mod = split_by_modality(embeddings, metadata)
    print(f"Modalities: {list(embeddings_by_mod.keys())}")
    for mod, embs in embeddings_by_mod.items():
        print(f"  {mod}: {len(embs)} images")
    
    # Run benchmarks
    n_queries = min(50, len(metadata))
    k = 5
    
    results = {}
    
    # 1. Single-index
    t0 = time.time()
    results["single"] = benchmark_single_index(embeddings, metadata, n_queries, k)
    results["single"]["time"] = time.time() - t0
    
    # 2. Multi-index
    t0 = time.time()
    results["multi"] = benchmark_multi_index(embeddings_by_mod, metadata, n_queries, k)
    results["multi"]["time"] = time.time() - t0
    
    # 3. Hybrid
    t0 = time.time()
    results["hybrid"] = benchmark_hybrid(embeddings_by_mod, metadata, n_queries, k)
    results["hybrid"]["time"] = time.time() - t0
    
    # 4. Cross-modal alignment
    t0 = time.time()
    results["alignment"] = benchmark_cross_modal_alignment(embeddings, metadata, n_queries, k)
    results["alignment"]["time"] = time.time() - t0
    
    # Print results
    print("\n" + "=" * 60)
    print("Results Summary")
    print("=" * 60)
    
    print(f"\n{'Method':<20} {'Same-Modal R@5':<18} {'Cross-Modal R@5':<18} {'Time (s)':<10}")
    print("-" * 66)
    
    for method, res in results.items():
        same = res.get("same_modal_recall@k", 0)
        cross = res.get("cross_modal_recall@k", 0) or res.get("hybrid_recall@k", 0)
        t = res.get("time", 0)
        print(f"{method:<20} {same:<18.4f} {cross:<18.4f} {t:<10.3f}")
    
    print("\n" + "=" * 60)
    print("Recommendation: Use the method with highest cross-modal recall")
    print("=" * 60)


if __name__ == "__main__":
    main()