arabic_normalizer.py

"""
Main normalizer class - enhanced for massive dataset processing.
"""
import logging
import json
import pickle
from pathlib import Path
from typing import Optional
from tqdm import tqdm

from config import Config
from llm_processor import LLMProcessor
from text_utils import clean_text, is_valid_arabic_text, split_into_chunks, create_text_chunks


class ArabicNormalizer:
    """Main class for Arabic text normalization with massive dataset support."""
    
    def __init__(self, config: Optional[Config] = None):
        self.config = config or Config()
        self.llm = LLMProcessor(self.config)
        self.logger = self._setup_logging()
        
        # Checkpointing for huge datasets
        self.checkpoint_dir = Path("checkpoints")
        self.checkpoint_dir.mkdir(exist_ok=True)
    
    def setup_llm(self, api_key: Optional[str] = None):
        """Initialize the LLM processor."""
        self.llm.setup(api_key)
        self.logger.info("LLM setup completed for massive processing")
        self.logger.info(f"Configuration: {self.config.requests_per_minute} req/min, {self.config.api_delay}s delay")
    
    def normalize_text(self, text: str) -> str:
        """Normalize a single text string."""
        
        # Basic preprocessing
        text = clean_text(text)
        
        # Skip if too short
        if len(text) < self.config.min_paragraph_length:
            return ""
        
        # Process with LLM in chunks if needed
        chunks = split_into_chunks(text, self.config.chunk_size)
        normalized_chunks = []
        
        for i, chunk in enumerate(chunks):
            normalized_chunk = self.llm.normalize_text(chunk, i + 1, len(chunks))
            if normalized_chunk:
                normalized_chunks.append(normalized_chunk)
        
        result = '\n'.join(normalized_chunks)
        
        # Final validation
        if not is_valid_arabic_text(result, self.config.min_arabic_ratio):
            return ""
        
        # Length check
        if (self.config.max_paragraph_length and 
            len(result) > self.config.max_paragraph_length):
            result = result[:self.config.max_paragraph_length].rsplit(' ', 1)[0]
        
        return result
    
    def process_file_massive(self, input_path: str, output_path: str, resume: bool = True) -> bool:
        """Process a massive file with checkpointing and resume capability."""
        
        input_file = Path(input_path)
        output_file = Path(output_path)
        checkpoint_file = self.checkpoint_dir / f"{input_file.stem}_checkpoint.pkl"
        
        if not input_file.exists():
            self.logger.error(f"Input file not found: {input_path}")
            return False
        
        # Create output directory
        output_file.parent.mkdir(parents=True, exist_ok=True)
        
        # Load checkpoint if resuming
        processed_count = 0
        skipped_chunks = 0
        
        if resume and checkpoint_file.exists():
            try:
                with open(checkpoint_file, 'rb') as f:
                    checkpoint_data = pickle.load(f)
                processed_count = checkpoint_data.get('processed_count', 0)
                skipped_chunks = checkpoint_data.get('chunks_processed', 0)
                self.logger.info(f"Resuming from checkpoint: {processed_count} chunks processed")
            except Exception as e:
                self.logger.warning(f"Could not load checkpoint: {e}")
                processed_count = 0
                skipped_chunks = 0
        
        self.logger.info(f"Processing massive file: {input_path} → {output_path}")
        self.logger.info(f"Estimated size: {input_file.stat().st_size / (1024*1024*1024):.2f} GB")
        
        try:
            file_size = input_file.stat().st_size
            chunks_processed = 0
            
            # Open files with appropriate mode
            output_mode = 'a' if resume and processed_count > 0 else 'w'
            
            with (input_file.open('r', encoding='utf-8') as infile,
                  output_file.open(output_mode, encoding='utf-8') as outfile,
                  tqdm(total=file_size, unit='B', unit_scale=True, desc="Processing massive dataset") as pbar):
                
                for chunk in create_text_chunks(infile):
                    chunks_processed += 1
                    
                    # Skip already processed chunks
                    if chunks_processed <= skipped_chunks:
                        pbar.update(len(chunk.encode('utf-8')))
                        continue
                    
                    try:
                        normalized = self.normalize_text(chunk)
                        if normalized:
                            outfile.write(normalized + '\n')
                            outfile.flush()  # Ensure data is written
                            processed_count += 1
                        
                        pbar.update(len(chunk.encode('utf-8')))
                        
                        # Checkpoint progress with enhanced metrics
                        if chunks_processed % self.config.checkpoint_every == 0:
                            self._save_checkpoint(checkpoint_file, processed_count, chunks_processed)
                            stats = self.llm.get_stats()
                            self.logger.info(
                                f"🔄 CHECKPOINT #{chunks_processed//self.config.checkpoint_every} | "
                                f"Chunks: {processed_count}/{chunks_processed} | "
                                f"API calls: {stats['total_requests']} | "
                                f"Success: {stats['success_rate']:.1f}% | "
                                f"Session: {stats['session_duration_minutes']:.1f}m | "
                                f"Efficiency: {stats['processing_efficiency']:.1f}% | "
                                f"Rate: {stats['average_rate_per_minute']:.1f}/min"
                            )
                        
                    except Exception as e:
                        self.logger.warning(f"Error processing chunk {chunks_processed}: {e}")
                        continue
            
            # Final checkpoint and cleanup
            self._save_checkpoint(checkpoint_file, processed_count, chunks_processed)
            stats = self.llm.get_stats()
            
            # Show comprehensive session summary
            self.logger.info("🎉 MASSIVE PROCESSING COMPLETED!")
            self.logger.info(f"📊 FINAL STATS:")
            self.logger.info(f"   • Chunks processed: {processed_count}")
            self.logger.info(f"   • Total API calls: {stats['total_requests']}")
            self.logger.info(f"   • Success rate: {stats['success_rate']:.1f}%")
            self.logger.info(f"   • Session duration: {stats['session_duration_minutes']:.1f} minutes")
            self.logger.info(f"   • Processing efficiency: {stats['processing_efficiency']:.1f}%")
            self.logger.info(f"   • Average rate: {stats['average_rate_per_minute']:.1f} requests/minute")
            self.logger.info(f"   • Total wait time: {stats['total_wait_time_minutes']:.1f} minutes")
            self.logger.info(f"   • Output file: {output_file}")
            
            return True
            
        except Exception as e:
            self.logger.error(f"Massive file processing failed: {e}")
            return False
    
    def process_file(self, input_path: str, output_path: str) -> bool:
        """Process a normal-sized file (for backward compatibility)."""
        return self.process_file_massive(input_path, output_path, resume=False)
    
    def process_batch_massive(self, input_dir: str, output_dir: str, file_pattern: str = "*.txt") -> int:
        """Process multiple massive files in a directory."""
        
        input_path = Path(input_dir)
        output_path = Path(output_dir)
        
        if not input_path.exists():
            self.logger.error(f"Input directory not found: {input_dir}")
            return 0
        
        output_path.mkdir(parents=True, exist_ok=True)
        
        files = list(input_path.glob(file_pattern))
        processed_count = 0
        
        self.logger.info(f"Found {len(files)} massive files to process")
        
        for i, file in enumerate(files, 1):
            self.logger.info(f"[{i}/{len(files)}] Processing massive file: {file.name}")
            
            output_file = output_path / f"{file.stem}_normalized.txt"
            
            if self.process_file_massive(str(file), str(output_file), resume=True):
                processed_count += 1
                self.logger.info(f"✓ Completed: {file.name}")
            else:
                self.logger.error(f"✗ Failed: {file.name}")
        
        self.logger.info(f"Massive batch processing completed: {processed_count}/{len(files)} files")
        return processed_count
    
    def process_batch(self, input_dir: str, output_dir: str, file_pattern: str = "*.txt") -> int:
        """Process multiple files (delegates to massive processing)."""
        return self.process_batch_massive(input_dir, output_dir, file_pattern)
    
    def _save_checkpoint(self, checkpoint_file: Path, processed_count: int, chunks_processed: int):
        """Save processing checkpoint."""
        if self.config.enable_resume:
            checkpoint_data = {
                'processed_count': processed_count,
                'chunks_processed': chunks_processed,
                'timestamp': str(Path(__file__).stat().st_mtime)
            }
            try:
                with open(checkpoint_file, 'wb') as f:
                    pickle.dump(checkpoint_data, f)
            except Exception as e:
                self.logger.warning(f"Could not save checkpoint: {e}")
    
    def estimate_processing_time(self, file_path: str) -> dict:
        """Estimate processing time for a massive file."""
        file_size = Path(file_path).stat().st_size
        
        # Rough estimates based on configuration
        chars_per_chunk = self.config.chunk_size
        estimated_chunks = file_size / chars_per_chunk
        
        # Time per request (including delays and retries)
        time_per_request = self.config.api_delay * (1 + self.config.max_retries * 0.1)
        
        estimated_hours = (estimated_chunks * time_per_request) / 3600
        estimated_days = estimated_hours / 24
        
        return {
            'file_size_gb': file_size / (1024*1024*1024),
            'estimated_chunks': int(estimated_chunks),
            'estimated_hours': estimated_hours,
            'estimated_days': estimated_days,
            'requests_per_hour': 3600 / time_per_request
        }
    
    def _setup_logging(self):
        """Setup logging configuration."""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler('normalizer_massive.log', encoding='utf-8'),
                logging.StreamHandler()
            ]
        )
        return logging.getLogger(__name__)