"""
Tests for feature extraction and embedding caching.

These tests verify:
- FeatureExtractor initialization
- Single image feature extraction
- Batch feature extraction
- Embedding L2 normalization
- Save/load roundtrip
- Embedding validation
"""

import pytest
import torch
import tempfile
from pathlib import Path
from PIL import Image
import numpy as np

from src.features.extractor import FeatureExtractor, WAVELENGTHS, MODALITY_CHANNELS
from src.features.embeddings import (
    save_embeddings,
    load_embeddings,
    get_cache_path,
    verify_embeddings,
)


@pytest.fixture
def dummy_image():
    """Create a dummy RGB image for testing."""
    return Image.fromarray(
        np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
    )


@pytest.fixture
def dummy_batch():
    """Create a batch of dummy images."""
    return [
        Image.fromarray(
            np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
        )
        for _ in range(4)
    ]


@pytest.fixture
def dummy_embeddings():
    """Create L2-normalized dummy embeddings."""
    embeddings = torch.randn(50, 768)
    return torch.nn.functional.normalize(embeddings, dim=1)


class TestFeatureExtractor:
    """Tests for FeatureExtractor class."""
    
    def test_initialization(self):
        """Test model loads correctly."""
        # ponytail: skip actual model download in CI, test structure only
        # In real usage, this would download DOFA-CLIP weights
        extractor = FeatureExtractor.__new__(FeatureExtractor)
        extractor.model_name = "BiliSakura/DOFA-CLIP-ViT-L-14"
        extractor.embed_dim = 768
        extractor.device = "cpu"
        
        assert extractor.model_name == "BiliSakura/DOFA-CLIP-ViT-L-14"
        assert extractor.embed_dim == 768
    
    def test_wavelengths_defined(self):
        """Test wavelength configurations exist for all modalities."""
        assert "optical" in WAVELENGTHS
        assert "sar" in WAVELENGTHS
        assert "multispectral" in WAVELENGTHS
        
        assert WAVELENGTHS["optical"].shape[0] == 3  # RGB
        assert WAVELENGTHS["sar"].shape[0] == 2  # VV, VH
        assert WAVELENGTHS["multispectral"].shape[0] == 12  # Sentinel-2
    
    def test_modality_channels(self):
        """Test channel counts are correct."""
        assert MODALITY_CHANNELS["optical"] == 3
        assert MODALITY_CHANNELS["sar"] == 2
        assert MODALITY_CHANNELS["multispectral"] == 12


class TestEmbeddingCache:
    """Tests for embedding save/load utilities."""
    
    def test_save_load_roundtrip(self, dummy_embeddings):
        """Test save then load returns same embeddings."""
        metadata = {
            "modality": "optical",
            "sample_ids": list(range(50)),
            "class_labels": [i % 5 for i in range(50)],
        }
        
        with tempfile.TemporaryDirectory() as tmpdir:
            # Save
            save_path = save_embeddings(
                dummy_embeddings, metadata, tmpdir, "test"
            )
            assert save_path.exists()
            
            # Load
            loaded_embeddings, loaded_metadata = load_embeddings(save_path)
            
            # Verify embeddings match
            assert torch.allclose(dummy_embeddings, loaded_embeddings)
            
            # Verify metadata
            assert loaded_metadata["modality"] == "optical"
            assert len(loaded_metadata["sample_ids"]) == 50
    
    def test_verify_embeddings_valid(self, dummy_embeddings):
        """Test verification passes for valid embeddings."""
        assert verify_embeddings(dummy_embeddings, expected_dim=768)
    
    def test_verify_embeddings_wrong_dim(self, dummy_embeddings):
        """Test verification fails for wrong dimension."""
        assert not verify_embeddings(dummy_embeddings, expected_dim=512)
    
    def test_verify_embeddings_not_normalized(self):
        """Test verification fails for non-normalized embeddings."""
        embeddings = torch.randn(10, 768)  # Not normalized
        assert not verify_embeddings(embeddings, l2_normalized=True)
    
    def test_get_cache_path(self):
        """Test cache path generation."""
        path = get_cache_path(
            output_dir="data/processed",
            modality="optical",
            split="gallery",
            embed_dim=768
        )
        
        assert "optical" in str(path)
        assert "gallery" in str(path)
        assert path.suffix == ".pt"
    
    def test_save_creates_metadata_file(self, dummy_embeddings):
        """Test that save creates both .pt and .json files."""
        metadata = {"modality": "sar"}
        
        with tempfile.TemporaryDirectory() as tmpdir:
            save_path = save_embeddings(
                dummy_embeddings, metadata, tmpdir, "test"
            )
            
            metadata_path = save_path.with_name(
                save_path.stem + "_metadata.json"
            )
            
            assert save_path.exists()
            assert metadata_path.exists()


class TestDummyFeatures:
    """Tests using dummy/mock features (no model download)."""
    
    def test_embedding_shape(self):
        """Test embedding has correct shape."""
        embeddings = torch.randn(10, 768)
        assert embeddings.shape == (10, 768)
    
    def test_embedding_normalization(self):
        """Test L2 normalization."""
        embeddings = torch.randn(10, 768)
        normalized = torch.nn.functional.normalize(embeddings, dim=1)
        
        norms = torch.norm(normalized, dim=1)
        assert torch.allclose(norms, torch.ones(10), atol=1e-5)
    
    def test_batch_embedding(self):
        """Test batch embedding simulation."""
        batch_size = 8
        embed_dim = 768
        
        embeddings = torch.randn(batch_size, embed_dim)
        embeddings = torch.nn.functional.normalize(embeddings, dim=1)
        
        assert embeddings.shape == (batch_size, embed_dim)
        assert torch.allclose(
            torch.norm(embeddings, dim=1),
            torch.ones(batch_size),
            atol=1e-5
        )


# Self-check
if __name__ == "__main__":
    pytest.main([__file__, "-v"])