OCR/backend/engine/recognition/template_engine.py

import hashlib
import json
from typing import List, Dict, Any, Tuple
from sqlalchemy.orm import Session
from models.template_models import Template, TemplateDocument, DocumentLayout, TemplateRecognitionHistory
from schemas.template_schemas import TemplateRecognitionResult

class TemplateRecognitionEngine:
    def __init__(self, db: Session):
        self.db = db
        
    def generate_fingerprint(self, layouts: List[DocumentLayout]) -> str:
        """
        Generate a structural fingerprint based on the document's layout.
        We use relative positions of HEADER, VENDOR, TABLE_HEADER blocks.
        """
        key_blocks = []
        for layout in layouts:
            if layout.block_type in ["HEADER", "VENDOR", "TABLE_HEADER", "TOTAL"]:
                key_blocks.append({
                    "type": layout.block_type,
                    "text": layout.text_value[:50] if layout.text_value else "", # first 50 chars
                    "rx": round(float(layout.x_coordinate) / 100) if layout.x_coordinate else 0, # relative bucket x
                    "ry": round(float(layout.y_coordinate) / 100) if layout.y_coordinate else 0  # relative bucket y
                })
        
        # Sort top to bottom, left to right
        key_blocks.sort(key=lambda b: (b['ry'], b['rx']))
        
        fingerprint_data = json.dumps(key_blocks)
        return hashlib.sha256(fingerprint_data.encode('utf-8')).hexdigest()

    def recognize_template(self, document_id: int) -> TemplateRecognitionResult:
        """
        Matches a document against existing templates using a 4-level weighted score:
        Level 1: Vendor Match (40%)
        Level 2: Header Similarity (20%)
        Level 3: Layout/Fingerprint Similarity (20%)
        Level 4: Coordinate Similarity (20%)
        Threshold: 85%
        """
        layouts = self.db.query(DocumentLayout).filter(DocumentLayout.fk_document_id == document_id).all()
        if not layouts:
            return TemplateRecognitionResult(templateMatched=False)
            
        doc_fingerprint = self.generate_fingerprint(layouts)
        
        # Extract features for scoring
        doc_vendor = self._extract_vendor_name(layouts)
        doc_headers = self._extract_headers(layouts)
        
        templates = self.db.query(Template).filter(Template.active_flag == True).all()
        best_match = None
        highest_score = 0
        
        for template in templates:
            score = 0.0
            
            # Level 1: Fingerprint Exact Match (counts for Layout + Vendor + Header if exact)
            if template.template_fingerprint == doc_fingerprint:
                score += 100.0
            else:
                # Need to load template sample mapping to compare heuristics
                # In a real system, we'd compare against the mapped fields' coordinates
                score += self._calculate_heuristic_score(template, doc_vendor, doc_headers, layouts)
                
            if score > highest_score:
                highest_score = score
                best_match = template
                
        # Store history
        if best_match:
            history = TemplateRecognitionHistory(
                fk_template_id=best_match.pk_template_id,
                fk_document_id=document_id,
                recognition_score=highest_score,
                matched_flag=(highest_score >= 85.0)
            )
            self.db.add(history)
            self.db.commit()
            
        if best_match and highest_score >= 85.0:
            extracted_fields = self._auto_map_fields(best_match, layouts)
            return TemplateRecognitionResult(
                templateMatched=True,
                templateId=best_match.pk_template_id,
                confidence=highest_score,
                extractedFields=extracted_fields
            )
            
        return TemplateRecognitionResult(templateMatched=False)

    def _extract_vendor_name(self, layouts: List[DocumentLayout]) -> str:
        for layout in layouts:
            if layout.block_type == "VENDOR":
                return layout.text_value.lower()
        return ""
        
    def _extract_headers(self, layouts: List[DocumentLayout]) -> List[str]:
        headers = []
        for layout in layouts:
            if layout.block_type == "TABLE_HEADER" and layout.text_value:
                headers.append(layout.text_value.lower())
        return headers

    def _calculate_heuristic_score(self, template: Template, doc_vendor: str, doc_headers: List[str], layouts: List[DocumentLayout]) -> float:
        score = 0.0
        
        # This requires the template to have some stored metadata or we check its fields
        # E.g., if template has a field "Vendor Name" mapped to a specific text
        
        # 1. Vendor Match (40%)
        # For this mockup, we check if the template name matches the vendor
        if template.template_name.lower() in doc_vendor or doc_vendor in template.template_name.lower():
            score += 40.0
            
        # 2. Header Similarity (20%)
        # Check if template fields exist that match doc_headers
        field_labels = [f.field_label.lower() for f in template.fields]
        header_matches = sum(1 for h in doc_headers if any(h in fl or fl in h for fl in field_labels))
        if len(doc_headers) > 0:
             score += (header_matches / len(doc_headers)) * 20.0
             
        # 3. Layout / Coordinate Similarity
        # We can check the template's previous mappings coordinates against current document
        # If coordinates are within a tolerance, we add score
        mapping_count = 0
        match_count = 0
        for mapping in template.mappings:
            mapping_count += 1
            # Find a layout block in current document that is near mapping coordinates
            for layout in layouts:
                if layout.page_no == mapping.page_no:
                    # check distance
                    if mapping.x_coordinate and mapping.y_coordinate and layout.x_coordinate and layout.y_coordinate:
                        dx = abs(float(mapping.x_coordinate) - float(layout.x_coordinate))
                        dy = abs(float(mapping.y_coordinate) - float(layout.y_coordinate))
                        if dx < 50 and dy < 20: # arbitrary tolerance
                            match_count += 1
                            break
        
        if mapping_count > 0:
            coord_score = (match_count / mapping_count) * 40.0 # Layout 20% + Coord 20%
            score += coord_score
            
        return score

    def _auto_map_fields(self, template: Template, layouts: List[DocumentLayout]) -> List[dict]:
        extracted = []
        for mapping in template.mappings:
            # Find closest block in new document
            best_block = None
            min_dist = float('inf')
            
            for layout in layouts:
                if layout.page_no == mapping.page_no:
                    if mapping.x_coordinate and mapping.y_coordinate and layout.x_coordinate and layout.y_coordinate:
                        dx = abs(float(mapping.x_coordinate) - float(layout.x_coordinate))
                        dy = abs(float(mapping.y_coordinate) - float(layout.y_coordinate))
                        dist = dx**2 + dy**2
                        if dist < min_dist and dist < 5000: # tolerance squared
                            min_dist = dist
                            best_block = layout
                            
            if best_block:
                extracted.append({
                    "field_id": mapping.fk_template_field_id,
                    "field_label": mapping.field.field_label if mapping.field else "",
                    "value": best_block.text_value,
                    "confidence": 95.0, # arbitrary high confidence for coordinate match
                    "layout_id": best_block.pk_document_data_id
                })
                
        return extracted