#!/usr/bin/env python3
"""
Landvex Vision — Identify-modell
Foto in → rankade kandidater med konfidens
"""
import json
import hashlib
from pathlib import Path
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass

import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import numpy as np

@dataclass
class IdentifyKandidat:
    lvx_id: str
    posttyp: str
    namn_sv: str
    namn_en: str
    konfidens: float
    verifieringsniva: str
    kannetecken_match: List[str]
    embedding_distans: float

class LandvexIdentifyModel:
    """Vision-modell för infrastrukturidentifiering."""
    
    def __init__(self, model_path: Optional[Path] = None):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"  🖥️  Enhet: {self.device}")
        
        # Ladda förtränad ResNet som backbone
        self.backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
        self.backbone = nn.Sequential(*list(self.backbone.children())[:-1])  # Ta bort sista FC
        self.backbone = self.backbone.to(self.device)
        self.backbone.eval()
        
        # Transform för bilder
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        
        # Embedding-databas (lvx_id → embedding)
        self.embeddings: Dict[str, np.ndarray] = {}
        self.metadata: Dict[str, dict] = {}
        
        if model_path and model_path.exists():
            self.ladda(model_path)
    
    def bild_till_embedding(self, bild: Image.Image) -> np.ndarray:
        """Konvertera bild till embedding-vektor."""
        tensor = self.transform(bild).unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            embedding = self.backbone(tensor)
        
        return embedding.cpu().numpy().flatten()
    
    def lagg_till_klass(self, lvx_id: str, bilder: List[Image.Image], metadata: dict):
        """Lägg till en objektklass med träningsbilder."""
        embeddings = [self.bild_till_embedding(b) for b in bilder]
        medel_embedding = np.mean(embeddings, axis=0)
        
        self.embeddings[lvx_id] = medel_embedding / np.linalg.norm(medel_embedding)
        self.metadata[lvx_id] = metadata
    
    def identifiera(self, bild: Image.Image, top_k: int = 5) -> List[IdentifyKandidat]:
        """Identifiera objekt i bild. Returnera top-k kandidater."""
        if not self.embeddings:
            return []
        
        query_embedding = self.bild_till_embedding(bild)
        query_embedding = query_embedding / np.linalg.norm(query_embedding)
        
        # Beräkna kosinuslikhet
        resultat = []
        for lvx_id, emb in self.embeddings.items():
            distans = np.dot(query_embedding, emb)
            meta = self.metadata[lvx_id]
            
            # Konfidens = likhet * verifieringsnivå-faktor
            verif_faktor = {
                "obekraftad": 0.5,
                "kallbelagd": 0.75,
                "faltverifierad": 0.9,
                "tillverkarbekraftad": 0.95,
            }.get(meta.get("verifieringsniva", "obekraftad"), 0.5)
            
            konfidens = float(distans * verif_faktor)
            
            resultat.append(IdentifyKandidat(
                lvx_id=lvx_id,
                posttyp=meta.get("posttyp", "objektklass"),
                namn_sv=meta.get("namn_sv", ""),
                namn_en=meta.get("namn_en", ""),
                konfidens=konfidens,
                verifieringsniva=meta.get("verifieringsniva", "obekraftad"),
                kannetecken_match=meta.get("kannetecken", [])[:3],
                embedding_distans=float(distans),
            ))
        
        # Sortera efter konfidens
        resultat.sort(key=lambda x: x.konfidens, reverse=True)
        return resultat[:top_k]
    
    def spara(self, path: Path):
        """Spara modell och embeddings."""
        path.parent.mkdir(parents=True, exist_ok=True)
        data = {
            "embeddings": {k: v.tolist() for k, v in self.embeddings.items()},
            "metadata": self.metadata,
        }
        torch.save(data, path)
        print(f"  💾 Modell sparad: {path}")
    
    def ladda(self, path: Path):
        """Ladda modell och embeddings."""
        data = torch.load(path, map_location=self.device)
        self.embeddings = {k: np.array(v) for k, v in data["embeddings"].items()}
        self.metadata = data["metadata"]
        print(f"  📂 Modell laddad: {len(self.embeddings)} klasser")

def main():
    """Demo: träna på syntetiska data och identifiera."""
    from PIL import ImageDraw
    
    model = LandvexIdentifyModel()
    
    # Skapa syntetiska träningsbilder (i verkligheten: riktiga foton)
    def skapa_testbild(farg, storlek=(224, 224)):
        img = Image.new('RGB', storlek, farg)
        draw = ImageDraw.Draw(img)
        draw.rectangle([50, 50, 174, 174], outline="white", width=3)
        return img
    
    # Lägg till två klasser
    model.lagg_till_klass("LVX-TRP-0101", [
        skapa_testbild("gray"),
        skapa_testbild("lightgray"),
    ], {
        "posttyp": "objektklass",
        "namn_sv": "Belysningsstolpe",
        "namn_en": "Lighting column",
        "verifieringsniva": "kallbelagd",
        "kannetecken": ["Grå stolpe", "Ljustopp"],
    })
    
    model.lagg_till_klass("LVX-TRP-0102", [
        skapa_testbild("blue"),
        skapa_testbild("darkblue"),
    ], {
        "posttyp": "objektklass",
        "namn_sv": "Vägbelysningsarmatur",
        "namn_en": "Road lighting luminaire",
        "verifieringsniva": "kallbelagd",
        "kannetecken": ["Blå armatur", "LED-ljus"],
    })
    
    # Testa identifiering
    test_bild = skapa_testbild("gray")
    kandidater = model.identifiera(test_bild, top_k=2)
    
    print("\n🎯 Identifieringsresultat:")
    for k in kandidater:
        print(f"  {k.lvx_id}: {k.namn_sv} (konfidens: {k.konfidens:.3f})")
    
    # Spara modell
    model.spara(Path("/tmp/landvex-vision-model.pt"))

if __name__ == "__main__":
    main()