"""
SatCLIP-compatible image encoder using OpenAI CLIP ViT-L/14.

Replaces the custom SatCLIP ViT with OpenAI's CLIP (openai/clip-vit-large-patch14)
which produces actually discriminative embeddings for land-cover retrieval.

Interface preserved: .encode(tensor, normalize=True) -> (N, 768) tensor
"""

import torch
import torch.nn.functional as F
from transformers import CLIPModel, CLIPProcessor
from torchvision import transforms


class SatCLIPEncoder:
    """
    Image encoder for satellite image retrieval using OpenAI CLIP ViT-L/14.

    Handles multi-channel input by converting to 3-channel RGB internally.
    """

    def __init__(self, device: str = None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.embed_dim = 768

        print("Loading OpenAI CLIP ViT-L/14...")
        self.model = CLIPModel.from_pretrained(
            "openai/clip-vit-large-patch14").to(self.device)
        self.processor = CLIPProcessor.from_pretrained(
            "openai/clip-vit-large-patch14")
        self.model.eval()

        # For direct tensor input (bypass processor)
        self.normalize = transforms.Normalize(
            mean=[0.48145466, 0.4578275, 0.40821073],
            std=[0.26862954, 0.26130258, 0.27577711],
        )

    def _to_3ch(self, tensor: torch.Tensor) -> torch.Tensor:
        """Convert any channel-count tensor to 3 channels for CLIP."""
        n = tensor.shape[1]
        if n == 3:
            return tensor
        if n == 1:
            return tensor.repeat(1, 3, 1, 1)
        if n == 2:
            return tensor.repeat(1, 3, 1, 1)[:, :3]
        # n >= 3: take first 3 channels
        return tensor[:, :3]

    @torch.no_grad()
    def encode(self, image_tensor: torch.Tensor,
               normalize: bool = True) -> torch.Tensor:
        """
        Encode image tensor to embedding.

        Args:
            image_tensor: (N, C, 224, 224) tensor, values in [0, 1]
            normalize: L2-normalize output

        Returns:
            (N, 768) embedding tensor
        """
        # Convert to 3 channels
        x = self._to_3ch(image_tensor.to(self.device))

        # Apply CLIP normalization (ImageNet stats)
        x = self.normalize(x)

        # Pass through CLIP vision encoder
        vision_outputs = self.model.vision_model(x)
        features = vision_outputs.pooler_output
        # Apply the visual projection
        features = self.model.visual_projection(features)

        if normalize:
            features = F.normalize(features, dim=-1)
        return features


# -- Self-check --
if __name__ == "__main__":
    print("Testing SatCLIPEncoder (CLIP backend)...")
    encoder = SatCLIPEncoder()
    print(f"Embed dim: {encoder.embed_dim}")

    dummy = torch.randn(2, 3, 224, 224)
    emb = encoder.encode(dummy)
    print(f"Output shape: {emb.shape}")
    print(f"L2 norm: {emb.norm(dim=-1).tolist()}")

    # Test multi-channel handling
    dummy_1ch = torch.randn(2, 1, 224, 224)
    emb_1ch = encoder.encode(dummy_1ch)
    print(f"1ch -> 768: {emb_1ch.shape}, norm={emb_1ch.norm(dim=-1).tolist()}")

    dummy_13ch = torch.randn(2, 13, 224, 224)
    emb_13ch = encoder.encode(dummy_13ch)
    print(f"13ch -> 768: {emb_13ch.shape}, norm={emb_13ch.norm(dim=-1).tolist()}")

    print("OK")