#!/usr/bin/env python3
"""
Landvex Vision — Vektordatabas
Lagrar och söker bild- och textembeddings.
"""
import json
import numpy as np
from pathlib import Path
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass

import faiss  # Facebook AI Similarity Search

@dataclass
class VektorsokResultat:
    lvx_id: str
    distans: float
    metadata: dict

class LandvexVektordatabas:
    """Vektordatabas för bild- och textembeddings."""
    
    def __init__(self, dimension: int = 2048, index_type: str = "FlatIP"):
        self.dimension = dimension
        self.index_type = index_type
        
        # FAISS index för snabb sökning
        if index_type == "FlatIP":
            self.index = faiss.IndexFlatIP(dimension)  # Inner product (cosine om normaliserad)
        elif index_type == "IVF":
            nlist = 100  # Antal kluster
            quantizer = faiss.IndexFlatIP(dimension)
            self.index = faiss.IndexIVFFlat(quantizer, dimension, nlist)
        else:
            raise ValueError(f"Okänd index-typ: {index_type}")
        
        # Mappning index_id → lvx_id
        self.id_map: List[str] = []
        self.metadata: Dict[str, dict] = {}
    
    def lagg_till(self, lvx_id: str, embedding: np.ndarray, metadata: dict):
        """Lägg till embedding i databasen."""
        # Normalisera för cosine similarity
        embedding = embedding / np.linalg.norm(embedding)
        embedding = embedding.reshape(1, -1).astype('float32')
        
        self.index.add(embedding)
        self.id_map.append(lvx_id)
        self.metadata[lvx_id] = metadata
    
    def sok(self, query: np.ndarray, top_k: int = 5) -> List[VektorsokResultat]:
        """Sök närmaste grannar."""
        query = query / np.linalg.norm(query)
        query = query.reshape(1, -1).astype('float32')
        
        distanser, indices = self.index.search(query, top_k)
        
        resultat = []
        for dist, idx in zip(distanser[0], indices[0]):
            if idx < 0 or idx >= len(self.id_map):
                continue
            
            lvx_id = self.id_map[idx]
            resultat.append(VektorsokResultat(
                lvx_id=lvx_id,
                distans=float(dist),
                metadata=self.metadata.get(lvx_id, {})
            ))
        
        return resultat
    
    def spara(self, path: Path):
        """Spara index och metadata."""
        path.mkdir(parents=True, exist_ok=True)
        
        # Spara FAISS-index
        faiss.write_index(self.index, str(path / "index.faiss"))
        
        # Spara metadata
        with open(path / "metadata.json", 'w') as f:
            json.dump({
                "id_map": self.id_map,
                "metadata": self.metadata,
                "dimension": self.dimension,
                "index_type": self.index_type,
            }, f, indent=2)
        
        print(f"  💾 Vektordatabas sparad: {path}")
    
    def ladda(self, path: Path):
        """Ladda index och metadata."""
        self.index = faiss.read_index(str(path / "index.faiss"))
        
        with open(path / "metadata.json", 'r') as f:
            data = json.load(f)
        
        self.id_map = data["id_map"]
        self.metadata = data["metadata"]
        self.dimension = data["dimension"]
        self.index_type = data["index_type"]
        
        print(f"  📂 Vektordatabas laddad: {len(self.id_map)} vektorer")

def main():
    """Demo: vektordatabas."""
    print("🔍 Landvex Vektordatabas")
    
    db = LandvexVektordatabas(dimension=128)  # Låg dimension för demo
    
    # Lägg till exempel-embeddings
    np.random.seed(42)
    for i in range(100):
        emb = np.random.randn(128)
        db.lagg_till(
            lvx_id=f"LVX-TRP-{1000+i:04d}",
            embedding=emb,
            metadata={
                "namn_sv": f"Testobjekt {i}",
                "verifieringsniva": "kallbelagd",
            }
        )
    
    # Sök
    query = np.random.randn(128)
    resultat = db.sok(query, top_k=5)
    
    print("\n🎯 Sökresultat:")
    for r in resultat:
        print(f"  {r.lvx_id}: distans={r.distans:.3f}, {r.metadata['namn_sv']}")
    
    # Spara och ladda
    db.spara(Path("/tmp/landvex-vectordb"))
    
    db2 = LandvexVektordatabas()
    db2.ladda(Path("/tmp/landvex-vectordb"))
    
    print(f"\n✅ Databas: {len(db2.id_map)} vektorer")

if __name__ == "__main__":
    main()