"""
Ground truth utilities for evaluation.

Handles loading, creating, and validating ground truth labels.
"""

import json
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass


@dataclass
class GroundTruthPair:
    """A query-gallery ground truth pair."""
    query_id: int
    gallery_ids: List[int]
    query_modality: str
    gallery_modality: str


class GroundTruth:
    """
    Ground truth labels for retrieval evaluation.
    
    Stores query-gallery pairs for same-modal and cross-modal evaluation.
    """
    
    def __init__(self):
        """Initialize empty ground truth."""
        self.pairs: List[GroundTruthPair] = []
        self._query_to_gallery: Dict[int, List[int]] = {}
    
    def add_pair(
        self,
        query_id: int,
        gallery_ids: List[int],
        query_modality: str,
        gallery_modality: str
    ) -> None:
        """
        Add a ground truth pair.
        
        Args:
            query_id: Query sample ID
            gallery_ids: List of matching gallery IDs
            query_modality: Modality of query
            gallery_modality: Modality of gallery
        """
        pair = GroundTruthPair(
            query_id=query_id,
            gallery_ids=gallery_ids,
            query_modality=query_modality,
            gallery_modality=gallery_modality
        )
        self.pairs.append(pair)
        self._query_to_gallery[query_id] = gallery_ids
    
    def get_gallery_ids(self, query_id: int) -> List[int]:
        """
        Get ground truth gallery IDs for a query.
        
        Args:
            query_id: Query sample ID
            
        Returns:
            List of matching gallery IDs
        """
        return self._query_to_gallery.get(query_id, [])
    
    @property
    def n_pairs(self) -> int:
        """Number of ground truth pairs."""
        return len(self.pairs)
    
    def save(self, path: str) -> None:
        """
        Save ground truth to JSON.
        
        Args:
            path: Output path
        """
        data = {
            "pairs": [
                {
                    "query_id": p.query_id,
                    "gallery_ids": p.gallery_ids,
                    "query_modality": p.query_modality,
                    "gallery_modality": p.gallery_modality,
                }
                for p in self.pairs
            ]
        }
        
        with open(path, "w") as f:
            json.dump(data, f, indent=2)
    
    @classmethod
    def load(cls, path: str) -> "GroundTruth":
        """
        Load ground truth from JSON.
        
        Args:
            path: Input path
            
        Returns:
            GroundTruth instance
        """
        with open(path, "r") as f:
            data = json.load(f)
        
        gt = cls()
        for pair_data in data["pairs"]:
            gt.add_pair(**pair_data)
        
        return gt
    
    def validate(self) -> bool:
        """
        Validate ground truth structure.
        
        Returns:
            True if valid
        """
        for pair in self.pairs:
            if not pair.gallery_ids:
                print(f"Warning: Query {pair.query_id} has no gallery matches")
                return False
            
            if pair.query_modality not in ["optical", "sar", "multispectral"]:
                print(f"Invalid query modality: {pair.query_modality}")
                return False
            
            if pair.gallery_modality not in ["optical", "sar", "multispectral"]:
                print(f"Invalid gallery modality: {pair.gallery_modality}")
                return False
        
        return True
    
    def get_same_modal_pairs(self) -> List[GroundTruthPair]:
        """Get pairs where query and gallery are same modality."""
        return [
            p for p in self.pairs
            if p.query_modality == p.gallery_modality
        ]
    
    def get_cross_modal_pairs(self) -> List[GroundTruthPair]:
        """Get pairs where query and gallery are different modalities."""
        return [
            p for p in self.pairs
            if p.query_modality != p.gallery_modality
        ]


def create_ground_truth_from_matches(
    matches: Dict[Tuple[str, str], List[Tuple[int, int]]]
) -> GroundTruth:
    """
    Create ground truth from modality matches.
    
    Args:
        matches: Dict mapping (query_modality, gallery_modality) to
                 list of (query_id, gallery_id) pairs
                 
    Returns:
        GroundTruth instance
    """
    gt = GroundTruth()
    
    for (query_mod, gallery_mod), pairs in matches.items():
        # Group by query_id
        query_to_galleries: Dict[int, List[int]] = {}
        for query_id, gallery_id in pairs:
            if query_id not in query_to_galleries:
                query_to_galleries[query_id] = []
            query_to_galleries[query_id].append(gallery_id)
        
        # Add to ground truth
        for query_id, gallery_ids in query_to_galleries.items():
            gt.add_pair(query_id, gallery_ids, query_mod, gallery_mod)
    
    return gt


# Self-check
if __name__ == "__main__":
    import tempfile
    
    print("Testing GroundTruth utilities...")
    
    # Create ground truth
    gt = GroundTruth()
    
    # Same-modal pairs
    gt.add_pair(0, [1, 2], "optical", "optical")
    gt.add_pair(1, [3, 4], "sar", "sar")
    
    # Cross-modal pairs
    gt.add_pair(2, [5, 6], "optical", "sar")
    gt.add_pair(3, [7, 8], "sar", "optical")
    
    print(f"Total pairs: {gt.n_pairs}")
    print(f"Same-modal: {len(gt.get_same_modal_pairs())}")
    print(f"Cross-modal: {len(gt.get_cross_modal_pairs())}")
    
    # Validate
    assert gt.validate(), "Validation failed"
    
    # Save/load roundtrip
    with tempfile.TemporaryDirectory() as tmpdir:
        save_path = Path(tmpdir) / "ground_truth.json"
        gt.save(save_path)
        
        loaded_gt = GroundTruth.load(save_path)
        
        assert loaded_gt.n_pairs == gt.n_pairs
        assert loaded_gt.validate()
    
    print("\nGroundTruth test passed!")