src/data_preparation.py

"""
Data preparation module for document text extraction.
Handles OCR, text cleaning, and dataset creation for NER training.
"""

import os
import json
import re
import pytesseract
from PIL import Image
import pandas as pd
import cv2
import numpy as np
from typing import List, Dict, Tuple, Optional
from pathlib import Path
import fitz  # PyMuPDF for PDF processing
from docx import Document
import easyocr


class DocumentProcessor:
    """Handles document processing, OCR, and text extraction."""
    
    def __init__(self, tesseract_path: Optional[str] = None):
        """Initialize document processor with OCR settings."""
        if tesseract_path:
            pytesseract.pytesseract.tesseract_cmd = tesseract_path
        
        # Initialize EasyOCR reader
        self.ocr_reader = easyocr.Reader(['en'])
        
        # Entity patterns for initial labeling
        self.entity_patterns = {
            'NAME': [
                r'\b[A-Z][a-z]+ [A-Z][a-z]+\b',  # First Last
                r'(?:Mr\.|Mrs\.|Ms\.|Dr\.)\s+[A-Z][a-z]+ [A-Z][a-z]+',  # Title + Name
            ],
            'DATE': [
                r'\b\d{1,2}[/\-]\d{1,2}[/\-]\d{2,4}\b',  # DD/MM/YYYY
                r'\b\d{4}[/\-]\d{1,2}[/\-]\d{1,2}\b',    # YYYY/MM/DD
                r'\b(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\s+\d{1,2},?\s+\d{2,4}\b'
            ],
            'INVOICE_NO': [
                r'(?:Invoice\s+(?:No|Number|#):\s*)?([A-Z]{2,4}[-]?\d{3,6})',
                r'(?:INV[-]?\d{3,6})',
            ],
            'AMOUNT': [
                r'\$\s*\d{1,3}(?:,\d{3})*(?:\.\d{2})?',  # $1,000.00
                r'\d{1,3}(?:,\d{3})*(?:\.\d{2})?\s*(?:USD|EUR|GBP)',  # 1000.00 USD
            ],
            'ADDRESS': [
                r'\d+\s+[A-Za-z\s]+(?:Street|St|Avenue|Ave|Road|Rd|Drive|Dr|Lane|Ln).*',
            ],
            'PHONE': [
                r'\+?\d{1,3}[-.\s]?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}',
                r'\(\d{3}\)\s*\d{3}-\d{4}',
            ],
            'EMAIL': [
                r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
            ]
        }
    
    def extract_text_from_pdf(self, pdf_path: str) -> str:
        """Extract text from PDF file."""
        try:
            doc = fitz.open(pdf_path)
            text = ""
            for page_num in range(len(doc)):
                page = doc.load_page(page_num)
                text += page.get_text()
            doc.close()
            return text
        except Exception as e:
            print(f"Error extracting text from PDF {pdf_path}: {e}")
            return ""
    
    def extract_text_from_docx(self, docx_path: str) -> str:
        """Extract text from DOCX file."""
        try:
            doc = Document(docx_path)
            text = ""
            for paragraph in doc.paragraphs:
                text += paragraph.text + "\n"
            return text
        except Exception as e:
            print(f"Error extracting text from DOCX {docx_path}: {e}")
            return ""
    
    def preprocess_image(self, image_path: str) -> np.ndarray:
        """Preprocess image for better OCR results."""
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"Could not load image: {image_path}")
        
        # Convert to grayscale
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # Apply Gaussian blur to reduce noise
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        
        # Apply adaptive threshold
        thresh = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
        )
        
        return thresh
    
    def extract_text_with_tesseract(self, image_path: str) -> str:
        """Extract text using Tesseract OCR."""
        try:
            preprocessed_img = self.preprocess_image(image_path)
            
            # Configure Tesseract
            custom_config = r'--oem 3 --psm 6'
            text = pytesseract.image_to_string(preprocessed_img, config=custom_config)
            
            return text
        except Exception as e:
            print(f"Error with Tesseract OCR on {image_path}: {e}")
            return ""
    
    def extract_text_with_easyocr(self, image_path: str) -> str:
        """Extract text using EasyOCR."""
        try:
            results = self.ocr_reader.readtext(image_path)
            text = " ".join([result[1] for result in results])
            return text
        except Exception as e:
            print(f"Error with EasyOCR on {image_path}: {e}")
            return ""
    
    def extract_text_from_image(self, image_path: str, use_easyocr: bool = True) -> str:
        """Extract text from image using OCR."""
        if use_easyocr:
            text = self.extract_text_with_easyocr(image_path)
            if not text.strip():  # Fallback to Tesseract
                text = self.extract_text_with_tesseract(image_path)
        else:
            text = self.extract_text_with_tesseract(image_path)
            if not text.strip():  # Fallback to EasyOCR
                text = self.extract_text_with_easyocr(image_path)
        
        return text
    
    def clean_text(self, text: str) -> str:
        """Clean and normalize extracted text."""
        # Remove extra whitespace
        text = re.sub(r'\s+', ' ', text)
        
        # Remove special characters but keep important punctuation
        text = re.sub(r'[^\w\s\.\,\:\;\-\$\(\)\[\]\/]', '', text)
        
        # Normalize whitespace around punctuation
        text = re.sub(r'\s*([,.;:])\s*', r'\1 ', text)
        
        return text.strip()
    
    def process_document(self, file_path: str) -> str:
        """Process any document type and extract text."""
        file_path = Path(file_path)
        file_ext = file_path.suffix.lower()
        
        if file_ext == '.pdf':
            text = self.extract_text_from_pdf(str(file_path))
        elif file_ext == '.docx':
            text = self.extract_text_from_docx(str(file_path))
        elif file_ext in ['.png', '.jpg', '.jpeg', '.tiff', '.bmp']:
            text = self.extract_text_from_image(str(file_path))
        else:
            raise ValueError(f"Unsupported file type: {file_ext}")
        
        return self.clean_text(text)


class NERDatasetCreator:
    """Creates NER training datasets from processed documents."""
    
    def __init__(self, document_processor: DocumentProcessor):
        self.document_processor = document_processor
        self.entity_labels = ['O', 'B-NAME', 'I-NAME', 'B-DATE', 'I-DATE', 
                             'B-INVOICE_NO', 'I-INVOICE_NO', 'B-AMOUNT', 'I-AMOUNT',
                             'B-ADDRESS', 'I-ADDRESS', 'B-PHONE', 'I-PHONE',
                             'B-EMAIL', 'I-EMAIL']
    
    def auto_label_text(self, text: str) -> List[Tuple[str, str]]:
        """Automatically label text using regex patterns."""
        words = text.split()
        labels = ['O'] * len(words)
        
        # Track word positions in original text
        word_positions = []
        start = 0
        for word in words:
            pos = text.find(word, start)
            word_positions.append((pos, pos + len(word)))
            start = pos + len(word)
        
        # Apply entity patterns
        for entity_type, patterns in self.document_processor.entity_patterns.items():
            for pattern in patterns:
                matches = list(re.finditer(pattern, text, re.IGNORECASE))
                for match in matches:
                    match_start, match_end = match.span()
                    
                    # Find which words overlap with this match
                    first_word_idx = None
                    last_word_idx = None
                    
                    for i, (word_start, word_end) in enumerate(word_positions):
                        if word_start >= match_start and word_end <= match_end:
                            if first_word_idx is None:
                                first_word_idx = i
                            last_word_idx = i
                        elif word_start < match_end and word_end > match_start:
                            # Partial overlap
                            if first_word_idx is None:
                                first_word_idx = i
                            last_word_idx = i
                    
                    # Apply BIO labeling
                    if first_word_idx is not None:
                        labels[first_word_idx] = f'B-{entity_type}'
                        for i in range(first_word_idx + 1, last_word_idx + 1):
                            labels[i] = f'I-{entity_type}'
        
        return list(zip(words, labels))
    
    def create_training_example(self, text: str) -> Dict:
        """Create a training example from text."""
        labeled_tokens = self.auto_label_text(text)
        
        tokens = [token for token, _ in labeled_tokens]
        labels = [label for _, label in labeled_tokens]
        
        return {
            'tokens': tokens,
            'labels': labels,
            'text': text
        }
    
    def create_sample_dataset(self) -> List[Dict]:
        """Create sample training data for demonstration."""
        sample_texts = [
            "Invoice sent to Robert White on 15/09/2025 Invoice No: INV-1024 Amount: $1,250",
            "Bill for Sarah Johnson dated March 10, 2025. Invoice Number: BL-2045. Total: $2,300.50",
            "Payment due from Michael Brown on 01/12/2025. Reference: PAY-3067. Sum: $890.00",
            "Receipt for Emma Wilson Invoice: REC-4089 Date: 2025-04-22 Amount: $1,750.25",
            "Dr. James Smith 123 Main Street Boston MA 02101 Phone: (555) 123-4567 Email: james@email.com",
            "Ms. Lisa Anderson 456 Oak Avenue New York NY 10001 Contact: +1-555-987-6543",
            "Invoice INV-5678 issued to David Lee on February 5, 2025 for $3,400.00",
            "Bill #BIL-9012 for Jennifer Garcia dated 2025-05-15. Total amount: $567.89"
        ]
        
        dataset = []
        for text in sample_texts:
            example = self.create_training_example(text)
            dataset.append(example)
        
        return dataset
    
    def process_documents_folder(self, folder_path: str) -> List[Dict]:
        """Process all documents in a folder and create training dataset."""
        folder_path = Path(folder_path)
        dataset = []
        
        if not folder_path.exists():
            print(f"Folder {folder_path} does not exist. Creating sample dataset instead.")
            return self.create_sample_dataset()
        
        supported_extensions = ['.pdf', '.docx', '.png', '.jpg', '.jpeg', '.tiff', '.bmp']
        
        for file_path in folder_path.rglob('*'):
            if file_path.suffix.lower() in supported_extensions:
                try:
                    print(f"Processing {file_path.name}...")
                    text = self.document_processor.process_document(str(file_path))
                    
                    if text.strip():  # Only process non-empty texts
                        example = self.create_training_example(text)
                        example['source_file'] = str(file_path)
                        dataset.append(example)
                        print(f"Processed {file_path.name}")
                    else:
                        print(f"No text extracted from {file_path.name}")
                        
                except Exception as e:
                    print(f"Error processing {file_path.name}: {e}")
        
        if not dataset:
            print("No documents processed. Creating sample dataset.")
            return self.create_sample_dataset()
        
        return dataset
    
    def save_dataset(self, dataset: List[Dict], output_path: str):
        """Save dataset to JSON file."""
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(dataset, f, indent=2, ensure_ascii=False)
        
        print(f"Dataset saved to {output_path}")
        print(f"Total examples: {len(dataset)}")
        
        # Print statistics
        all_labels = []
        for example in dataset:
            all_labels.extend(example['labels'])
        
        label_counts = {}
        for label in all_labels:
            label_counts[label] = label_counts.get(label, 0) + 1
        
        print("\nLabel distribution:")
        for label, count in sorted(label_counts.items()):
            print(f"  {label}: {count}")


def main():
    """Main function to demonstrate data preparation."""
    # Initialize components
    processor = DocumentProcessor()
    dataset_creator = NERDatasetCreator(processor)
    
    # Process documents (or create sample data)
    raw_data_path = "data/raw"
    dataset = dataset_creator.process_documents_folder(raw_data_path)
    
    # Save processed dataset
    output_path = "data/processed/ner_dataset.json"
    dataset_creator.save_dataset(dataset, output_path)
    
    print(f"\nData preparation completed!")
    print(f"Processed {len(dataset)} documents")


if __name__ == "__main__":
    main()