"""
Quick Start Demo Mode for IDP System
Works without trained models using rule-based extraction
Allows immediate testing while models train
"""

import re
from typing import Dict, List
import numpy as np


class DemoClassifier:
    """Simple rule-based classifier for demo"""
    
    def predict(self, text: str) -> Dict:
        text_lower = text.lower()
        
        # Simple heuristics
        if 'invoice' in text_lower:
            return {
                'predicted_class': 'INVOICE',
                'confidence': 0.85,
                'probabilities': {'INVOICE': 0.85, 'RECEIPT': 0.10, 'FORM': 0.03, 'OTHER': 0.02}
            }
        elif any(k in text_lower for k in ['bank statement', 'account statement', 'statement of account']):
             return {
                'predicted_class': 'BANK_STATEMENT',
                'confidence': 0.90,
                'probabilities': {'INVOICE': 0.05, 'RECEIPT': 0.05, 'FORM': 0.05, 'BANK_STATEMENT': 0.85}
            }
        elif 'receipt' in text_lower:
            return {
                'predicted_class': 'RECEIPT',
                'confidence': 0.80,
                'probabilities': {'INVOICE': 0.15, 'RECEIPT': 0.80, 'FORM': 0.03, 'OTHER': 0.02}
            }
        elif 'form' in text_lower or 'application' in text_lower:
            return {
                'predicted_class': 'FORM',
                'confidence': 0.75,
                'probabilities': {'INVOICE': 0.05, 'RECEIPT': 0.10, 'FORM': 0.75, 'OTHER': 0.10}
            }
        else:
            # Fallback for things with prices/totals that aren't explicitly receipts
            if 'total' in text_lower and '$' in text_lower:
                 return {
                    'predicted_class': 'RECEIPT',
                    'confidence': 0.60,
                    'probabilities': {'INVOICE': 0.20, 'RECEIPT': 0.60, 'FORM': 0.10, 'OTHER': 0.10}
                }
            
            return {
                'predicted_class': 'OTHER',
                'confidence': 0.70,
                'probabilities': {'INVOICE': 0.10, 'RECEIPT': 0.10, 'FORM': 0.10, 'OTHER': 0.70}
            }


class DemoNER:
    """Simple rule-based NER for demo"""
    
    def __init__(self):
        self.patterns = {
            'INVOICE_NUMBER': [
                r'INV[-\s]?\d+',
                r'Invoice\s*#?\s*:?\s*([A-Z0-9\-]+)',
                r'\b\d{6,}\b',
            ],
            'DATE': [
                r'\d{1,2}[/-]\d{1,2}[/-]\d{2,4}',
                r'\d{4}[/-]\d{1,2}[/-]\d{1,2}',
            ],
            'TOTAL_AMOUNT': [
                r'Total\s*:?\s*[₹$€£]\s*[\d,]+\.?\d*',
                r'[₹$€£]\s*[\d,]+\.?\d*',
            ],
            'TAX_AMOUNT': [
                r'Tax\s*:?\s*[₹$€£]\s*[\d,]+\.?\d*',
                r'GST\s*:?\s*[₹$€£]\s*[\d,]+\.?\d*',
            ],
        }
    
    def predict(self, text: str) -> Dict:
        entities = []
        
        for entity_type, patterns in self.patterns.items():
            for pattern in patterns:
                matches = re.finditer(pattern, text, re.IGNORECASE)
                for match in matches:
                    entities.append({
                        'entity': entity_type,
                        'text': match.group(0).strip(),
                        'confidence': 0.75,
                        'start': match.start(),
                        'end': match.end()
                    })
        
        return {'entities': entities, 'text': text}


def create_demo_models_directory():
    """Create models directory with demo flag"""
    import os
    
    os.makedirs('models/classifier', exist_ok=True)
    os.makedirs('models/ner', exist_ok=True)
    
    # Create demo marker
    with open('models/DEMO_MODE.txt', 'w') as f:
        f.write('Using demo mode - rule-based extraction\n')
        f.write('Train models with: python train_classifier.py && python train_ner.py\n')
    
    print("✓ Demo mode initialized")
    print("  Using rule-based extraction until models are trained")


if __name__ == "__main__":
    create_demo_models_directory()
    
    # Test demo
    demo_text = """
    INVOICE #INV-12345
    Date: 28/11/2025
    Total: $1,250.00
    Tax: $125.00
    """
    
    classifier = DemoClassifier()
    ner = DemoNER()
    
    print("\nDemo Classifier:")
    print(classifier.predict(demo_text))
    
    print("\nDemo NER:")
    print(ner.predict(demo_text))