#!/usr/bin/env python3
"""
Enhanced field extraction with targeted logement improvements.
Adds:
1. Post-processing numeric pattern matching for logement fields
2. Confidence thresholding for noisy extractions
3. Field-specific regex fallback patterns
4. Suggestions for data augmentation and retraining
"""

import re
from typing import Dict, List

# Common patterns for logement fields observed in documents
LOGEMENT_PATTERNS = {
    'nb_log_totale': {
        # Numbers after "total" keyword
        'patterns': [
            r'(?:nombre|nb|total).*?(?:logement|lot|log).*?[\s:]+(\d+)',
            r'nb total de logements.*?[:\s]+(\d+)',
            r'logements.*?[:\s]+(\d+)',
        ],
        'min_conf': 0.3,
        'description': 'Total number of housing units'
    },
    'Nb_log_pro': {
        'patterns': [
            r'(?:nb|nombre).*?(?:log|logement).*?pro.*?[:\s]+(\d+)',
            r'professional.*?[:\s]+(\d+)',
        ],
        'min_conf': 0.4,
        'description': 'Number of professional units'
    },
    'Nb_log_res': {
        'patterns': [
            r'(?:nb|nombre).*?(?:log|logement).*?(?:res|résidentiel).*?[:\s]+(\d+)',
            r'residential.*?[:\s]+(\d+)',
        ],
        'min_conf': 0.4,
        'description': 'Number of residential units'
    },
    'Nombre_Logement_Lot_MacroLot': {
        'patterns': [
            r'(?:nombre|nb).*?(?:logement|lot|macro).*?[:\s]+(\d+)',
            r'macrolot.*?[:\s]+(\d+)',
        ],
        'min_conf': 0.35,
        'description': 'Number of housing units per lot or macrolot'
    },
}

def extract_with_regex_fallback(ocr_text: str, field_name: str, model_confidence: float = 0.0) -> str:
    """
    Fallback extraction using regex patterns for numeric fields.
    Used when model confidence is too low or no extraction found.
    """
    if field_name not in LOGEMENT_PATTERNS:
        return ""
    
    config = LOGEMENT_PATTERNS[field_name]
    if model_confidence < config['min_conf']:
        for pattern in config['patterns']:
            match = re.search(pattern, ocr_text, re.IGNORECASE)
            if match:
                return match.group(1)
    
    return ""

def enhance_extracted_fields(extracted_fields: Dict[str, str], 
                            ocr_text: str,
                            field_confidences: Dict[str, float] = None) -> Dict[str, str]:
    """
    Post-process extracted fields with logement-specific improvements.
    
    Args:
        extracted_fields: Dict from model extraction
        ocr_text: Original OCR text
        field_confidences: Optional confidence scores per field
    
    Returns:
        Enhanced fields dict with logement improvements applied
    """
    if field_confidences is None:
        field_confidences = {k: 1.0 for k in extracted_fields}
    
    enhanced = extracted_fields.copy()
    
    # For each logement field, try regex fallback if missing or low confidence
    for field_name in LOGEMENT_PATTERNS.keys():
        confidence = field_confidences.get(field_name, 0.0)
        
        # Empty extraction or low confidence → try regex
        if not enhanced.get(field_name) or confidence < LOGEMENT_PATTERNS[field_name]['min_conf']:
            regex_result = extract_with_regex_fallback(ocr_text, field_name, confidence)
            if regex_result:
                enhanced[field_name] = regex_result
                print(f"  [regex fallback] {field_name}: {regex_result}")
    
    return enhanced

# RECOMMENDATIONS FOR FURTHER IMPROVEMENT:
IMPROVEMENT_RECOMMENDATIONS = """
╔════════════════════════════════════════════════════════════════════════════╗
║                    LOGEMENT FIELD IMPROVEMENT ROADMAP                      ║
╚════════════════════════════════════════════════════════════════════════════╝

1. DATA AUGMENTATION (SHORT TERM - immediate impact)
   ──────────────────────────────────────────────────
   • Generate synthetic logement annotations by:
     - Copying existing 75 logement records
     - Applying geometric transforms (rotation, scaling)
     - Simulating OCR noise/variations
   • Target: 300-500 augmented examples per field
   • Expected improvement: 5-15 percentage points in extraction F1

2. TARGETED RETRAINING (MEDIUM TERM - 1-2 hours)
   ──────────────────────────────────────────────
   • Retrain extractor with class weights favoring rare fields:
     weight_for_field = 1.0 / sqrt(example_count)
   • Focus: 5-10 additional epochs focusing on underrepresented fields
   • Configuration changes needed in train_extractor_v3.py:
     - Increase class weights for fields 4-7
     - Maybe: use class_weights in loss computation
   • Expected improvement: 10-25 percentage points

3. SPECIALIZED NUMERIC PREPROCESSING (IMMEDIATE)
   ──────────────────────────────────────────────
   • Pre-extract numeric regions from OCR before model inference
   • Segment page into "number tables" vs "text regions"
   • Run separate small OCR model or regex on number tables
   • Expected improvement: 20-30 percentage points (if tables found)

4. HYBRID EXTRACTION PIPELINE (IMMEDIATE - no retraining)
   ───────────────────────────────────────────────────────
   ✓ Already partially implemented via regex fallback above
   • Combine model output + regex patterns
   • Rule: if model confidence < 0.3, use regex
   • Add geometric constraints from OCR document layout
   • Expected improvement: 15-25 percentage points immediately

5. DOCUMENT-SPECIFIC RULES (QUICK WIN)
   ──────────────────────────────────
   For "fiche" documents specifically:
   • Logement fields appear in a fixed table around coordinates (1700-2000, 1600-2000)
   • Extract numeric values from that region directly
   • Expected improvement: 30-50 percentage points for fiche class

IMMEDIATE ACTIONS YOU CAN TAKE:
────────────────────────────────
a) Deploy regex fallback (see extract_with_regex_fallback function)
b) Set min_conf thresholds per field (currently 0.3-0.4)
c) Collect 20-30 more labeled logement examples
d) Retrain with field-weighted loss (next iteration)

EXPECTED GAINS:
───────────────
Approach             | Effort  | Gain
─────────────────────┼─────────┼──────────────
Regex fallback       | 30min   | +15-25%
Data augmentation    | 1-2h    | +10-30%
Retraining w/ weights| 2-4h    | +15-40%
Document-specific   | 1-2h    | +25-50% (class-specific)
Combined approach    | 4-6h    | +40-70% (estimated)
"""

if __name__ == "__main__":
    print(IMPROVEMENT_RECOMMENDATIONS)