"""
StateEmbedding - Couche 3 : Fusion Multi-Modale

Crée un "fingerprint" unique de l'écran en fusionnant :
- Image embedding (screenshot complet)
- Text embedding (texte détecté)
- Title embedding (titre de fenêtre)
- UI embedding (éléments UI)
"""

from dataclasses import dataclass, field
from typing import Dict, Optional, Any
from pathlib import Path
import numpy as np
import json


@dataclass
class EmbeddingComponent:
    """Composante d'un State Embedding"""
    weight: float
    vector_id: str
    source_text: Optional[str] = None
    
    def to_dict(self) -> Dict[str, Any]:
        result = {
            "weight": self.weight,
            "vector_id": self.vector_id
        }
        if self.source_text:
            result["source_text"] = self.source_text
        return result
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'EmbeddingComponent':
        return cls(
            weight=data["weight"],
            vector_id=data["vector_id"],
            source_text=data.get("source_text")
        )


@dataclass
class StateEmbedding:
    """
    State Embedding - Vecteur unique représentant un état d'écran
    
    Fusion multi-modale :
    - 50% Image (screenshot complet)
    - 30% Texte (texte détecté)
    - 10% Titre (fenêtre)
    - 10% UI (éléments détectés)
    """
    embedding_id: str
    vector_id: str  # Chemin vers fichier .npy
    dimensions: int
    fusion_method: str  # "weighted" ou "concat_projection"
    components: Dict[str, EmbeddingComponent] = field(default_factory=dict)
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    # Cache du vecteur en mémoire
    _vector_cache: Optional[np.ndarray] = field(default=None, repr=False, compare=False)
    
    def get_vector(self) -> np.ndarray:
        """Charger le vecteur depuis le fichier (avec cache)"""
        if self._vector_cache is None:
            vector_path = Path(self.vector_id)
            if vector_path.exists():
                self._vector_cache = np.load(vector_path)
            else:
                raise FileNotFoundError(f"Embedding vector not found: {self.vector_id}")
        return self._vector_cache
    
    def set_vector(self, vector: np.ndarray) -> None:
        """Définir le vecteur et le mettre en cache"""
        if vector.shape[0] != self.dimensions:
            raise ValueError(
                f"Vector dimensions mismatch: expected {self.dimensions}, "
                f"got {vector.shape[0]}"
            )
        self._vector_cache = vector
    
    def save_vector(self, vector: np.ndarray) -> None:
        """Sauvegarder le vecteur dans un fichier .npy"""
        vector_path = Path(self.vector_id)
        vector_path.parent.mkdir(parents=True, exist_ok=True)
        np.save(vector_path, vector)
        self._vector_cache = vector
    
    def compute_similarity(self, other: 'StateEmbedding') -> float:
        """
        Calculer similarité cosinus avec autre embedding
        
        Property 5: State Embedding Similarity Symmetry
        Property 6: State Embedding Similarity Bounds
        """
        vec1 = self.get_vector()
        vec2 = other.get_vector()
        
        # Similarité cosinus
        dot_product = np.dot(vec1, vec2)
        norm1 = np.linalg.norm(vec1)
        norm2 = np.linalg.norm(vec2)
        
        if norm1 == 0 or norm2 == 0:
            return 0.0
        
        similarity = dot_product / (norm1 * norm2)
        
        # Clamp entre -1 et 1 (pour éviter erreurs numériques)
        similarity = np.clip(similarity, -1.0, 1.0)
        
        return float(similarity)
    
    def is_normalized(self, tolerance: float = 1e-6) -> bool:
        """
        Vérifier si le vecteur est normalisé (L2 norm = 1.0)
        
        Property 4: State Embedding Normalization
        """
        vector = self.get_vector()
        norm = np.linalg.norm(vector)
        return abs(norm - 1.0) < tolerance
    
    def to_dict(self) -> Dict[str, Any]:
        """Sérialiser en JSON (sans le vecteur)"""
        return {
            "embedding_id": self.embedding_id,
            "vector_id": self.vector_id,
            "dimensions": self.dimensions,
            "fusion_method": self.fusion_method,
            "components": {
                name: comp.to_dict() 
                for name, comp in self.components.items()
            },
            "metadata": self.metadata
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'StateEmbedding':
        """Désérialiser depuis JSON"""
        components = {
            name: EmbeddingComponent.from_dict(comp_data)
            for name, comp_data in data.get("components", {}).items()
        }
        
        return cls(
            embedding_id=data["embedding_id"],
            vector_id=data["vector_id"],
            dimensions=data["dimensions"],
            fusion_method=data["fusion_method"],
            components=components,
            metadata=data.get("metadata", {})
        )
    
    def to_json(self) -> str:
        """Sérialiser en JSON string"""
        return json.dumps(self.to_dict(), indent=2)
    
    @classmethod
    def from_json(cls, json_str: str) -> 'StateEmbedding':
        """Désérialiser depuis JSON string"""
        data = json.loads(json_str)
        return cls.from_dict(data)
    
    def save_to_file(self, filepath: Path) -> None:
        """Sauvegarder métadonnées dans un fichier JSON"""
        filepath.parent.mkdir(parents=True, exist_ok=True)
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(self.to_dict(), f, indent=2)
    
    @classmethod
    def load_from_file(cls, filepath: Path) -> 'StateEmbedding':
        """Charger métadonnées depuis un fichier JSON"""
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        return cls.from_dict(data)


# Configuration par défaut des poids de fusion
DEFAULT_FUSION_WEIGHTS = {
    "image": 0.5,   # 50% - Screenshot complet
    "text": 0.3,    # 30% - Texte détecté
    "title": 0.1,   # 10% - Titre fenêtre
    "ui": 0.1       # 10% - Éléments UI
}

# Méthodes de fusion supportées
FUSION_METHODS = [
    "weighted",           # Fusion pondérée simple
    "concat_projection"   # Concaténation + projection
]