import cv2
import numpy as np

def template_match_ocr(image, templates):
    """
    Basic template matching OCR implementation
    """
    results = []
    
    for char, template in templates.items():
        # Normalized cross-correlation
        result = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
        
        # Confidence threshold
        if max_val > 0.8:
            results.append({
                'character': char,
                'confidence': max_val,
                'position': max_loc
            })
    
    return sorted(results, key=lambda x: x['confidence'], reverse=True)

import torch
import torch.nn as nn

class OCR_CNN(nn.Module):
    def __init__(self, num_classes=62):  # 26 + 26 + 10
        super(OCR_CNN, self).__init__()
        
        # Feature extraction layers
        self.conv_layers = nn.Sequential(
            # First block: edge detection
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            
            # Second block: pattern detection  
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            
            # Third block: complex feature extraction
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            
            # Fourth block: high-level features
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d((4, 4))
        )
        
        # Classification layers
        self.classifier = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(256 * 4 * 4, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )
    
    def forward(self, x):
        features = self.conv_layers(x)
        features = features.view(features.size(0), -1)
        output = self.classifier(features)
        return output

from transformers import TrOCRProcessor, VisionEncoderDecoderModel

class TrOCRSystem:
    def __init__(self, model_name="microsoft/trocr-base-handwritten"):
        self.processor = TrOCRProcessor.from_pretrained(model_name)
        self.model = VisionEncoderDecoderModel.from_pretrained(model_name)
    
    def recognize_text(self, image):
        """
        End-to-end text recognition using TrOCR
        """
        # Preprocess image
        pixel_values = self.processor(image, return_tensors="pt").pixel_values
        
        # Generate text
        generated_ids = self.model.generate(pixel_values, max_length=50)
        
        # Decode output
        generated_text = self.processor.batch_decode(
            generated_ids, 
            skip_special_tokens=True
        )[0]
        
        return {
            'text': generated_text,
            'confidence': 0.95,  # Confidence calculation would use beam search probabilities
            'token_probabilities': self._get_token_probabilities(generated_ids)
        }

    def _get_token_probabilities(self, generated_ids):
        """
        Calculate per-token probability scores from generated token IDs.

        Args:
            generated_ids: Tensor of generated token IDs

        Returns:
            List of probability scores for each token
        """
        # Get logits from the model's last forward pass
        with torch.no_grad():
            # Apply softmax to convert logits to probabilities
            probs = torch.softmax(self.model.decoder.logits, dim=-1)

            # Extract probabilities for the generated tokens
            token_probs = torch.gather(
                probs,
                dim=-1,
                index=generated_ids.unsqueeze(-1)
            ).squeeze(-1)

        return token_probs.tolist()

import numpy as np

def calculate_sequence_confidence(char_confidences, method='geometric_mean'):
    """
    Calculate sequence confidence from character confidences
    """
    confidences = np.array(char_confidences)
    
    if method == 'arithmetic_mean':
        return np.mean(confidences)
    
    elif method == 'geometric_mean':
        # Prevents overconfidence in long sequences
        return np.power(np.prod(confidences), 1.0/len(confidences))
    
    elif method == 'harmonic_mean':
        # Emphasizes lowest confidence characters
        return len(confidences) / np.sum(1.0 / confidences)
    
    elif method == 'minimum':
        # Conservative approach - weakest link
        return np.min(confidences)
    
    elif method == 'weighted_geometric':
        # Weight by character position importance
        weights = np.ones_like(confidences)
        weights[0] = 1.5  # First character more important
        weights[-1] = 1.3  # Last character more important
        
        weighted_product = np.prod(np.power(confidences, weights))
        return np.power(weighted_product, 1.0/np.sum(weights))

import asyncio
from concurrent.futures import ThreadPoolExecutor
import time

class OptimizedOCRPipeline:
    def __init__(self, batch_size=16, max_workers=4):
        self.batch_size = batch_size
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        self.preprocessing_time = []
        self.inference_time = []
        self.postprocessing_time = []
    
    async def process_documents(self, document_paths):
        """Parallel document processing with performance monitoring"""
        
        # Batch processing for efficiency
        batches = [
            document_paths[i:i + self.batch_size] 
            for i in range(0, len(document_paths), self.batch_size)
        ]
        
        results = []
        for batch in batches:
            # Parallel preprocessing
            start_time = time.time()
            preprocessed = await asyncio.gather(*[
                asyncio.create_task(self._preprocess_image(path))
                for path in batch
            ])
            self.preprocessing_time.append(time.time() - start_time)

    async def _preprocess_image(self, path):
        """
        Preprocess a single image for OCR inference.

        Args:
            path: Path to image file

        Returns:
            Preprocessed image tensor
        """
        import cv2
        import numpy as np

        # Load image
        img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)

        # Resize to model input size (e.g., 384x384)
        img = cv2.resize(img, (384, 384), interpolation=cv2.INTER_CUBIC)

        # Normalize to [0, 1]
        img = img.astype(np.float32) / 255.0

        # Convert to tensor and add batch dimension
        img_tensor = torch.from_numpy(img).unsqueeze(0).unsqueeze(0)

        return img_tensor

            # Batch inference
            start_time = time.time()
            predictions = await self._batch_inference(preprocessed)
            self.inference_time.append(time.time() - start_time)

    async def _batch_inference(self, preprocessed):
        """
        Run batch inference on preprocessed images.

        Args:
            preprocessed: List of preprocessed image tensors

        Returns:
            List of model predictions
        """
        # Stack batch
        batch_tensor = torch.cat(preprocessed, dim=0)

        # Move to GPU if available
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        batch_tensor = batch_tensor.to(device)

        # Run inference
        with torch.no_grad():
            predictions = self.model(batch_tensor)

        return predictions
            
            # Parallel postprocessing
            start_time = time.time()
            processed = await asyncio.gather(*[
                asyncio.create_task(self._postprocess_result(pred))
                for pred in predictions
            ])
            self.postprocessing_time.append(time.time() - start_time)

    async def _postprocess_result(self, pred):
        """
        Postprocess a single prediction into formatted output.

        Args:
            pred: Raw model prediction (logits or token IDs)

        Returns:
            Formatted OCR result dictionary
        """
        # Decode prediction to text
        if isinstance(pred, torch.Tensor):
            # Convert logits to token IDs if needed
            if pred.dim() > 1:
                pred = torch.argmax(pred, dim=-1)

            # Decode tokens to text
            text = self.tokenizer.decode(pred, skip_special_tokens=True)
        else:
            text = str(pred)

        # Calculate confidence score
        confidence = self._calculate_confidence(pred)

        return {
            'text': text,
            'confidence': confidence,
            'length': len(text)
        }

    def _calculate_confidence(self, pred):
        """Calculate average confidence score for prediction."""
        # Simplified confidence calculation
        return 0.95  # In production, would use actual probability scores
            
            results.extend(processed)
        
        return results, self._get_performance_stats()
    
    def _get_performance_stats(self):
        return {
            'avg_preprocessing_time': np.mean(self.preprocessing_time),
            'avg_inference_time': np.mean(self.inference_time), 
            'avg_postprocessing_time': np.mean(self.postprocessing_time),
            'total_pipeline_time': sum([
                sum(self.preprocessing_time),
                sum(self.inference_time),
                sum(self.postprocessing_time)
            ])
        }