import cv2
import numpy as np
from scipy import ndimage
from skimage import exposure

class FadedInkEnhancer:
    @staticmethod
    def adaptive_histogram_equalization(image_path, output_path, clip_limit=2.0):
        """
        Apply CLAHE (Contrast Limited Adaptive Histogram Equalization).

        CLAHE works on local regions, making it effective for documents
        with varying background degradation across the page.

        Args:
            image_path: Path to input image
            output_path: Path for enhanced output
            clip_limit: Contrast limiting threshold (1.0-4.0)
        """
        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)

        # Apply CLAHE with specific tile size for document structure
        clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8))
        enhanced = clahe.apply(img)

        cv2.imwrite(output_path, enhanced)
        return enhanced

    @staticmethod
    def homomorphic_filtering(image_path, output_path, gamma_h=2.0, gamma_l=0.5):
        """
        Homomorphic filtering separates illumination from reflectance.

        Particularly effective for documents with uneven lighting or
        background discoloration that varies across the page.

        Args:
            image_path: Path to input image
            output_path: Path for enhanced output
            gamma_h: High frequency gain (enhance detail)
            gamma_l: Low frequency gain (suppress background)
        """
        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float32)

        # Avoid log(0) by adding small constant
        img = np.log1p(img)

        # Fourier transform
        dft = cv2.dft(img, flags=cv2.DFT_COMPLEX_OUTPUT)
        dft_shift = np.fft.fftshift(dft)

        # Create homomorphic filter
        rows, cols = img.shape
        crow, ccol = rows // 2, cols // 2

        # Generate Gaussian high-pass filter
        y, x = np.ogrid[-crow:rows-crow, -ccol:cols-ccol]
        mask = np.exp(-(x*x + y*y) / (2.0 * (min(rows, cols) / 8)**2))

        # Apply gamma transformations
        H = (gamma_h - gamma_l) * (1 - mask) + gamma_l

        # Expand H to match DFT dimensions
        H = np.expand_dims(H, axis=2)
        H = np.repeat(H, 2, axis=2)

        # Apply filter
        filtered_dft = dft_shift * H

        # Inverse transform
        f_ishift = np.fft.ifftshift(filtered_dft)
        img_back = cv2.idft(f_ishift)
        img_back = cv2.magnitude(img_back[:, :, 0], img_back[:, :, 1])

        # Exponentiate to reverse log
        result = np.expm1(img_back)

        # Normalize to 0-255
        result = cv2.normalize(result, None, 0, 255, cv2.NORM_MINMAX)
        result = result.astype(np.uint8)

        cv2.imwrite(output_path, result)
        return result

    @staticmethod
    def unsharp_masking(image_path, output_path, sigma=1.0, strength=1.5):
        """
        Unsharp masking enhances edges and fine details.

        Effective for recovering faded text by emphasizing character edges.

        Args:
            image_path: Path to input image
            output_path: Path for enhanced output
            sigma: Gaussian blur sigma (controls detail scale)
            strength: Enhancement strength (1.0-3.0 recommended)
        """
        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float32)

        # Create blurred version
        blurred = cv2.GaussianBlur(img, (0, 0), sigma)

        # Calculate unsharp mask
        sharpened = cv2.addWeighted(img, 1.0 + strength, blurred, -strength, 0)

        # Clip to valid range
        sharpened = np.clip(sharpened, 0, 255).astype(np.uint8)

        cv2.imwrite(output_path, sharpened)
        return sharpened

    @staticmethod
    def rolling_ball_background_subtraction(
        image_path,
        output_path,
        radius=50
    ):
        """
        Rolling ball algorithm for background estimation and removal.

        Particularly effective for documents with non-uniform staining or
        discoloration that varies gradually across the page.

        Args:
            image_path: Path to input image
            output_path: Path for cleaned output
            radius: Rolling ball radius in pixels (larger = smoother background)
        """
        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)

        # Estimate background using morphological opening with circular kernel
        kernel = cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE,
            (2 * radius + 1, 2 * radius + 1)
        )

        background = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel)

        # Subtract background
        foreground = cv2.subtract(img, background)

        # Enhance contrast after subtraction
        foreground = cv2.normalize(
            foreground, None, 0, 255,
            cv2.NORM_MINMAX, cv2.CV_8U
        )

        cv2.imwrite(output_path, foreground)
        return foreground

class AdaptiveBinarization:
    @staticmethod
    def sauvola_binarization(image_path, output_path, window_size=25, k=0.2):
        """
        Sauvola's method for local adaptive thresholding.

        Particularly effective for historical documents with varying
        background intensity and faded ink.

        Args:
            image_path: Path to input image
            output_path: Path for binarized output
            window_size: Local window size (odd number, 15-51 typical)
            k: Sauvola parameter controlling sensitivity (0.2-0.5)
        """
        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float32)

        # Calculate local mean using integral images for efficiency
        mean = cv2.blur(img, (window_size, window_size))

        # Calculate local standard deviation
        mean_sq = cv2.blur(img * img, (window_size, window_size))
        variance = mean_sq - mean * mean
        std = np.sqrt(np.maximum(variance, 0))

        # Sauvola threshold formula
        R = 128  # Dynamic range of standard deviation
        threshold = mean * (1 + k * ((std / R) - 1))

        # Apply threshold
        binary = (img > threshold).astype(np.uint8) * 255

        cv2.imwrite(output_path, binary)
        return binary

    @staticmethod
    def wolf_jolion_binarization(
        image_path,
        output_path,
        window_size=25,
        k=0.3
    ):
        """
        Wolf-Jolion method improves on Sauvola for very low contrast.

        Adds adaptive noise estimation for robust performance on
        severely degraded documents.

        Args:
            image_path: Path to input image
            output_path: Path for binarized output
            window_size: Local window size
            k: Sensitivity parameter
        """
        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float32)

        # Local mean
        mean = cv2.blur(img, (window_size, window_size))

        # Local standard deviation
        mean_sq = cv2.blur(img * img, (window_size, window_size))
        variance = mean_sq - mean * mean
        std = np.sqrt(np.maximum(variance, 0))

        # Minimum standard deviation (noise level)
        min_std = np.min(std[std > 0]) if np.any(std > 0) else 1.0

        # Wolf-Jolion threshold
        threshold = mean - k * (mean - min_std) * (1 - std / np.max(std))

        # Apply threshold
        binary = (img > threshold).astype(np.uint8) * 255

        cv2.imwrite(output_path, binary)
        return binary

    @staticmethod
    def combined_method(image_path, output_path):
        """
        Combined preprocessing pipeline for maximum faded ink recovery.

        Applies multiple techniques in sequence for optimal results
        on severely degraded documents.

        Args:
            image_path: Path to input image
            output_path: Path for final binarized output
        """
        # Step 1: Rolling ball background subtraction
        enhancer = FadedInkEnhancer()
        step1 = enhancer.rolling_ball_background_subtraction(
            image_path,
            "temp_step1.png",
            radius=50
        )

        # Step 2: CLAHE contrast enhancement
        step2 = enhancer.adaptive_histogram_equalization(
            "temp_step1.png",
            "temp_step2.png",
            clip_limit=3.0
        )

        # Step 3: Unsharp masking for edge enhancement
        step3 = enhancer.unsharp_masking(
            "temp_step2.png",
            "temp_step3.png",
            sigma=1.0,
            strength=1.5
        )

        # Step 4: Sauvola binarization
        final = AdaptiveBinarization.sauvola_binarization(
            "temp_step3.png",
            output_path,
            window_size=25,
            k=0.2
        )

        # Clean up temporary files
        import os
        for temp in ["temp_step1.png", "temp_step2.png", "temp_step3.png"]:
            if os.path.exists(temp):
                os.remove(temp)

        return final

import torch
import torch.nn as nn

class DocumentEnhancementNet(nn.Module):
    def __init__(self):
        """
        U-Net architecture for document image enhancement.

        Encoder-decoder structure with skip connections enables
        both local detail recovery and global background understanding.
        """
        super(DocumentEnhancementNet, self).__init__()

        # Encoder (downsampling path)
        self.enc1 = self._conv_block(1, 64)
        self.enc2 = self._conv_block(64, 128)
        self.enc3 = self._conv_block(128, 256)
        self.enc4 = self._conv_block(256, 512)

        # Bottleneck
        self.bottleneck = self._conv_block(512, 1024)

        # Decoder (upsampling path)
        self.upconv4 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
        self.dec4 = self._conv_block(1024, 512)

        self.upconv3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.dec3 = self._conv_block(512, 256)

        self.upconv2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.dec2 = self._conv_block(256, 128)

        self.upconv1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.dec1 = self._conv_block(128, 64)

        # Final output layer
        self.out = nn.Conv2d(64, 1, 1)

        self.pool = nn.MaxPool2d(2)

    def _conv_block(self, in_channels, out_channels):
        """Double convolution block with batch normalization."""
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        """
        Forward pass with skip connections.

        Args:
            x: Input degraded document image (batch, 1, height, width)

        Returns:
            Enhanced document image (batch, 1, height, width)
        """
        # Encoder
        enc1 = self.enc1(x)
        enc2 = self.enc2(self.pool(enc1))
        enc3 = self.enc3(self.pool(enc2))
        enc4 = self.enc4(self.pool(enc3))

        # Bottleneck
        bottleneck = self.bottleneck(self.pool(enc4))

        # Decoder with skip connections
        dec4 = self.upconv4(bottleneck)
        dec4 = torch.cat([dec4, enc4], dim=1)
        dec4 = self.dec4(dec4)

        dec3 = self.upconv3(dec4)
        dec3 = torch.cat([dec3, enc3], dim=1)
        dec3 = self.dec3(dec3)

        dec2 = self.upconv2(dec3)
        dec2 = torch.cat([dec2, enc2], dim=1)
        dec2 = self.dec2(dec2)

        dec1 = self.upconv1(dec2)
        dec1 = torch.cat([dec1, enc1], dim=1)
        dec1 = self.dec1(dec1)

        # Output
        out = torch.sigmoid(self.out(dec1))

        return out


def train_enhancement_model(model, train_loader, val_loader, epochs=50):
    """
    Training loop for document enhancement network.

    Args:
        model: DocumentEnhancementNet instance
        train_loader: DataLoader with degraded/clean image pairs
        val_loader: Validation DataLoader
        epochs: Number of training epochs
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)

    criterion = nn.BCELoss()  # Binary Cross-Entropy for binary images
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', patience=5
    )

    best_val_loss = float('inf')

    for epoch in range(epochs):
        # Training
        model.train()
        train_loss = 0

        for degraded, clean in train_loader:
            degraded, clean = degraded.to(device), clean.to(device)

            optimizer.zero_grad()
            outputs = model(degraded)
            loss = criterion(outputs, clean)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()

        # Validation
        model.eval()
        val_loss = 0

        with torch.no_grad():
            for degraded, clean in val_loader:
                degraded, clean = degraded.to(device), clean.to(device)
                outputs = model(degraded)
                loss = criterion(outputs, clean)
                val_loss += loss.item()

        avg_train_loss = train_loss / len(train_loader)
        avg_val_loss = val_loss / len(val_loader)

        print(f"Epoch {epoch+1}/{epochs}")
        print(f"  Train Loss: {avg_train_loss:.4f}")
        print(f"  Val Loss: {avg_val_loss:.4f}")

        scheduler.step(avg_val_loss)

        # Save best model
        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            torch.save(model.state_dict(), 'best_enhancement_model.pt')
            print("  Saved best model")

    return model

import editdistance
from pathlib import Path

class PreprocessingEvaluator:
    def __init__(self, ocr_engine):
        """
        Evaluate preprocessing methods by OCR accuracy.

        Args:
            ocr_engine: Callable OCR function that takes image path and returns text
        """
        self.ocr_engine = ocr_engine

    def evaluate_method(
        self,
        method_func,
        test_images,
        ground_truths,
        method_params=None
    ):
        """
        Evaluate a preprocessing method on test set.

        Args:
            method_func: Preprocessing function
            test_images: List of test image paths
            ground_truths: Corresponding ground truth texts
            method_params: Dictionary of parameters for method_func

        Returns:
            Dictionary containing accuracy metrics
        """
        if method_params is None:
            method_params = {}

        predictions = []
        total_chars = 0
        total_errors = 0

        for img_path, gt_text in zip(test_images, ground_truths):
            # Apply preprocessing
            preprocessed_path = "temp_preprocessed.png"
            method_func(img_path, preprocessed_path, **method_params)

            # Run OCR
            predicted_text = self.ocr_engine(preprocessed_path)
            predictions.append(predicted_text)

            # Calculate [character errors](/articles/character-recognition-accuracy)
            errors = editdistance.eval(predicted_text, gt_text)
            total_errors += errors
            total_chars += len(gt_text)

        # Calculate metrics
        cer = (total_errors / total_chars) * 100 if total_chars > 0 else 0

        correct = sum(
            pred.strip() == gt.strip()
            for pred, gt in zip(predictions, ground_truths)
        )
        accuracy = (correct / len(predictions)) * 100

        return {
            'CER': cer,
            'Accuracy': accuracy,
            'Total_Samples': len(test_images),
            'Correct_Samples': correct
        }

    def compare_methods(self, methods_config, test_images, ground_truths):
        """
        Compare multiple preprocessing methods.

        Args:
            methods_config: List of dicts with 'name', 'func', 'params'
            test_images: Test image paths
            ground_truths: Ground truth texts

        Returns:
            Comparison results dictionary
        """
        results = {}

        for config in methods_config:
            print(f"Evaluating {config['name']}...")

            metrics = self.evaluate_method(
                config['func'],
                test_images,
                ground_truths,
                config.get('params', {})
            )

            results[config['name']] = metrics

            print(f"  CER: {metrics['CER']:.2f}%")
            print(f"  Accuracy: {metrics['Accuracy']:.2f}%")

        return results