#!/usr/bin/env python3
"""
Correction Analyzer - Analyze OCR correction patterns
Finds common mistakes and builds correction rules automatically
"""

import json
import os
from collections import Counter, defaultdict
from typing import List, Dict, Tuple
import re
from datetime import datetime

class CorrectionAnalyzer:
    """Analyze correction patterns from training data"""
    
    def __init__(self):
        self.corrections = []
        self.error_patterns = Counter()
        self.text_replacements = Counter()
        self.dimension_errors = Counter()
        self.thread_tolerance_errors = Counter()
        
    def load_correction_file(self, json_path: str):
        """Load a single correction JSON file"""
        try:
            with open(json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                self.corrections.append(data)
                return True
        except Exception as e:
            print(f"Error loading {json_path}: {e}")
            return False
    
    def load_all_corrections(self, directory: str = "."):
        """Load all correction JSON files from directory"""
        count = 0
        for filename in os.listdir(directory):
            if filename.startswith("ocr_corrections_") and filename.endswith(".json"):
                if self.load_correction_file(os.path.join(directory, filename)):
                    count += 1
        print(f"✅ Loaded {count} correction files")
        return count
    
    def analyze_corrections(self):
        """Analyze all corrections to find patterns"""
        print("\n🔍 Analyzing correction patterns...\n")
        
        for correction_data in self.corrections:
            original_zones = correction_data.get('original_zones', [])
            corrected_zones = correction_data.get('corrected_zones', [])
            
            for orig, corr in zip(original_zones, corrected_zones):
                orig_text = orig.get('text', '')
                corr_text = corr.get('text', '')
                correction_type = corr.get('correction_type', 'none')
                is_dimension = corr.get('is_dimension', False)
                
                # Only analyze actual corrections
                if correction_type == 'text_fixed' and orig_text != corr_text:
                    # General text replacements
                    self.text_replacements[(orig_text, corr_text)] += 1
                    
                    # Dimension-specific errors
                    if is_dimension:
                        self.dimension_errors[(orig_text, corr_text)] += 1
                    
                    # Thread tolerance patterns (6G, M6, 8H, etc.)
                    if self._is_thread_tolerance(corr_text):
                        self.thread_tolerance_errors[(orig_text, corr_text)] += 1
                    
                    # Character-level patterns
                    self._analyze_character_errors(orig_text, corr_text)
    
    def _is_thread_tolerance(self, text: str) -> bool:
        """Check if text is a thread tolerance (6G, M6, 8H, etc.)"""
        # Patterns: 6G, M6, 8H, M8x1.25, etc.
        return bool(re.match(r'^M?\d+[A-Z]?\d*\.?\d*$', text.strip()))
    
    def _analyze_character_errors(self, orig: str, corr: str):
        """Analyze character-level substitution errors"""
        if len(orig) == len(corr):
            for i, (o, c) in enumerate(zip(orig, corr)):
                if o != c:
                    context = f"{orig[max(0,i-1):i+2]}"
                    self.error_patterns[(o, c, context)] += 1
    
    def generate_report(self) -> str:
        """Generate analysis report"""
        report = []
        report.append("=" * 70)
        report.append("📊 OCR CORRECTION ANALYSIS REPORT")
        report.append("=" * 70)
        report.append(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        report.append(f"Total corrections analyzed: {len(self.corrections)}")
        report.append("")
        
        # Most common text replacements
        report.append("🔤 TOP 20 MOST COMMON TEXT CORRECTIONS:")
        report.append("-" * 70)
        for (orig, corr), count in self.text_replacements.most_common(20):
            report.append(f"  '{orig}' → '{corr}' ({count} times)")
        report.append("")
        
        # Thread tolerance errors
        if self.thread_tolerance_errors:
            report.append("🔩 THREAD TOLERANCE ERRORS:")
            report.append("-" * 70)
            for (orig, corr), count in self.thread_tolerance_errors.most_common(10):
                report.append(f"  '{orig}' → '{corr}' ({count} times)")
            report.append("")
        
        # Dimension errors
        if self.dimension_errors:
            report.append("📐 DIMENSION ERRORS:")
            report.append("-" * 70)
            for (orig, corr), count in self.dimension_errors.most_common(10):
                report.append(f"  '{orig}' → '{corr}' ({count} times)")
            report.append("")
        
        # Character-level patterns
        if self.error_patterns:
            report.append("🔠 TOP CHARACTER SUBSTITUTION ERRORS:")
            report.append("-" * 70)
            for (orig_char, corr_char, context), count in self.error_patterns.most_common(15):
                report.append(f"  '{orig_char}' → '{corr_char}' in context '{context}' ({count} times)")
            report.append("")
        
        report.append("=" * 70)
        
        return "\n".join(report)
    
    def generate_correction_rules(self) -> Dict:
        """Generate correction rules from patterns"""
        rules = {
            'exact_replacements': {},
            'regex_patterns': [],
            'character_fixes': {},
            'thread_tolerance_fixes': {},
            'dimension_fixes': {}
        }
        
        # Exact text replacements (high confidence)
        for (orig, corr), count in self.text_replacements.most_common():
            if count >= 2:  # At least 2 occurrences
                rules['exact_replacements'][orig] = corr
        
        # Thread tolerance patterns
        for (orig, corr), count in self.thread_tolerance_errors.most_common():
            if count >= 1:
                rules['thread_tolerance_fixes'][orig] = corr
        
        # Dimension-specific fixes
        for (orig, corr), count in self.dimension_errors.most_common():
            if count >= 2:
                rules['dimension_fixes'][orig] = corr
        
        # Generate regex patterns for common errors
        rules['regex_patterns'] = self._generate_regex_patterns()
        
        return rules
    
    def _generate_regex_patterns(self) -> List[Dict]:
        """Generate regex patterns from error analysis"""
        patterns = []
        
        # Pattern: 69 → 6G (thread tolerance)
        if any('9' in orig and 'G' in corr for (orig, corr), _ in self.thread_tolerance_errors.items()):
            patterns.append({
                'pattern': r'(\d)9\b',
                'replacement': r'\1G',
                'description': 'Fix thread tolerance: 69 → 6G'
            })
        
        # Pattern: Ø4B → Ø48
        if any('B' in orig and '8' in corr for (orig, corr), _ in self.dimension_errors.items()):
            patterns.append({
                'pattern': r'Ø(\d+)B',
                'replacement': r'Ø\g<1>8',
                'description': 'Fix OCR error: B → 8 in diameters'
            })
        
        # Pattern: 6g → 6G (case sensitivity)
        if any(orig.upper() == corr for (orig, corr), _ in self.thread_tolerance_errors.items()):
            patterns.append({
                'pattern': r'(\d+)([a-z])\b',
                'replacement': lambda m: m.group(1) + m.group(2).upper(),
                'description': 'Fix thread tolerance case: 6g → 6G'
            })
        
        return patterns
    
    def save_rules(self, output_path: str = "correction_rules.json"):
        """Save generated rules to JSON file"""
        rules = self.generate_correction_rules()
        
        # Convert lambda functions to strings for JSON
        for pattern in rules['regex_patterns']:
            if callable(pattern['replacement']):
                pattern['replacement'] = 'UPPER_CASE'  # Mark for special handling
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(rules, f, indent=2, ensure_ascii=False)
        
        print(f"✅ Correction rules saved to: {output_path}")
        print(f"   - {len(rules['exact_replacements'])} exact replacements")
        print(f"   - {len(rules['regex_patterns'])} regex patterns")
        print(f"   - {len(rules['thread_tolerance_fixes'])} thread tolerance fixes")
        print(f"   - {len(rules['dimension_fixes'])} dimension fixes")


def main():
    """Main function"""
    print("🔍 OCR Correction Pattern Analyzer")
    print("=" * 70)
    
    analyzer = CorrectionAnalyzer()
    
    # Load all correction files from current directory
    count = analyzer.load_all_corrections(".")
    
    if count == 0:
        print("\n⚠️  No correction files found!")
        print("   Place your ocr_corrections_*.json files in this directory")
        return
    
    # Analyze patterns
    analyzer.analyze_corrections()
    
    # Generate and print report
    report = analyzer.generate_report()
    print(report)
    
    # Save report to file
    with open("correction_analysis_report.txt", 'w', encoding='utf-8') as f:
        f.write(report)
    print(f"\n📄 Full report saved to: correction_analysis_report.txt")
    
    # Generate and save correction rules
    analyzer.save_rules("correction_rules.json")
    
    print("\n✅ Analysis complete!")
    print("\n💡 Next steps:")
    print("   1. Review correction_analysis_report.txt")
    print("   2. Check correction_rules.json")
    print("   3. Use correction_post_processor.py to apply rules to new OCR results")


if __name__ == "__main__":
    main()