#!/usr/bin/env python3
"""
Landvex Vision — OCR-pipeline för typskyltar
Extraherar tillverkare, modell, serienummer från foton.
"""
import json
from pathlib import Path
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass

import pytesseract
from PIL import Image
import cv2
import numpy as np

@dataclass
class OCRResultat:
    text: str
    konfidens: float
    bbox: Tuple[int, int, int, int]  # x, y, w, h
    falt: Dict[str, str]  # Extraherade fält

class LandvexOCR:
    """OCR-pipeline optimerad för infrastrukturtypskyltar."""
    
    def __init__(self, tesseract_cmd: Optional[str] = None):
        if tesseract_cmd:
            pytesseract.pytesseract.tesseract_cmd = tesseract_cmd
        
        # Fältmönster per objektklass
        self.falt_monster = {
            "LVX-ELN-0101": {  # Kabelskåp
                "tillverkare": r"(ABB|Schneider|Siemens|GE)\s*[\w-]*",
                "modell": r"(CDC|Prisma|Okken|IM\s*\w+)",
                "ar": r"20\d{2}",
            },
            "LVX-TRP-0102": {  # Vägbelysningsarmatur
                "tillverkare": r"(Philips|Thorn|Schreder|Acuity|Cree)",
                "modell": r"(SL-\w+|ER\w+|Vista\w+)",
                "effekt": r"(\d+)\s*W",
            },
            "LVX-VAT-0101": {  # Brunnsbetäckning
                "tillverkare": r"(ULMA|ACO|Wrede|GDK)",
                "klass": r"(A15|B125|C250|D400|E600|F900)",
                "material": r"(gjutjärn|stål|komposit|betong)",
            },
        }
    
    def forbehandla_bild(self, bild: np.ndarray) -> np.ndarray:
        """Förbättra bildkvalitet för OCR."""
        # Konvertera till gråskala
        if len(bild.shape) == 3:
            gray = cv2.cvtColor(bild, cv2.COLOR_RGB2GRAY)
        else:
            gray = bild
        
        # Brusreducering
        denoised = cv2.fastNlMeansDenoising(gray)
        
        # Kontrastförbättring (CLAHE)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(denoised)
        
        # Skärpa
        kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
        sharpened = cv2.filter2D(enhanced, -1, kernel)
        
        return sharpened
    
    def hitta_text_regioner(self, bild: np.ndarray) -> List[Tuple[int, int, int, int]]:
        """Hitta regioner som troligen innehåller text."""
        # MSER (Maximally Stable Extremal Regions)
        mser = cv2.MSER_create()
        regions, _ = mser.detectRegions(bild)
        
        # Filtrera små regioner
        h, w = bild.shape
        min_area = (h * w) * 0.001  # Minst 0.1% av bilden
        
        bboxes = []
        for region in regions:
            x, y, w, h = cv2.boundingRect(region)
            if w * h > min_area and w > h * 2:  # Text är oftast bredare än hög
                bboxes.append((x, y, w, h))
        
        return bboxes
    
    def kora_ocr(self, bild: Image.Image, lvx_id: Optional[str] = None) -> List[OCRResultat]:
        """Kör OCR på bild och extrahera fält."""
        # Konvertera till numpy
        img_array = np.array(bild)
        
        # Förbehandla
        processed = self.forbehandla_bild(img_array)
        
        # Hitta textregioner
        regioner = self.hitta_text_regioner(processed)
        
        resultat = []
        for x, y, w, h in regioner[:5]:  # Max 5 regioner
            # Beskär region
            roi = processed[y:y+h, x:x+w]
            
            # Kör Tesseract
            text = pytesseract.image_to_string(roi, lang='eng+swe')
            conf_data = pytesseract.image_to_data(roi, output_type=pytesseract.Output.DICT)
            
            # Beräkna medelkonfidens
            konfidenser = [c for c in conf_data['conf'] if c > 0]
            medel_konf = np.mean(konfidenser) if konfidenser else 0
            
            # Extrahera fält om vi vet objektklassen
            falt = {}
            if lvx_id and lvx_id in self.falt_monster:
                import re
                for falt_namn, monster in self.falt_monster[lvx_id].items():
                    match = re.search(monster, text, re.IGNORECASE)
                    if match:
                        falt[falt_namn] = match.group(1)
            
            resultat.append(OCRResultat(
                text=text.strip(),
                konfidens=medel_konf / 100.0,  # Normalisera till 0-1
                bbox=(x, y, w, h),
                falt=falt
            ))
        
        # Sortera efter konfidens
        resultat.sort(key=lambda x: x.konfidens, reverse=True)
        return resultat
    
    def extrahera_falt(self, text: str, lvx_id: str) -> Dict[str, str]:
        """Extrahera strukturerade fält från OCR-text."""
        import re
        
        falt = {}
        monster = self.falt_monster.get(lvx_id, {})
        
        for falt_namn, pattern in monster.items():
            matches = re.findall(pattern, text, re.IGNORECASE)
            if matches:
                falt[falt_namn] = matches[0]
        
        return falt

def main():
    """Demo: OCR på syntetisk bild."""
    print("🔤 Landvex OCR Pipeline")
    
    ocr = LandvexOCR()
    
    # Skapa syntetisk testbild med text
    from PIL import ImageDraw, ImageFont
    img = Image.new('RGB', (400, 200), color='white')
    draw = ImageDraw.Draw(img)
    
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 24)
    except:
        font = ImageFont.load_default()
    
    draw.text((20, 20), "ABB Kabeldon", fill='black', font=font)
    draw.text((20, 60), "CDC-LT 400", fill='black', font=font)
    draw.text((20, 100), "2023", fill='black', font=font)
    
    # Kör OCR
    resultat = ocr.kora_ocr(img, lvx_id="LVX-ELN-0101")
    
    print("\n🎯 OCR-resultat:")
    for i, r in enumerate(resultat[:3]):
        print(f"  Region {i+1}:")
        print(f"    Text: {r.text[:100]}")
        print(f"    Konfidens: {r.konfidens:.2f}")
        print(f"    Fält: {r.falt}")
    
    print("\n✅ OCR klar!")

if __name__ == "__main__":
    main()