import os
import re
import numpy as np
import cv2
from typing import List, Dict, Any
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.feature_extraction.text import TfidfVectorizer

CLASS_INDICATIVE_KEYWORDS = [
    "invoice", "receipt", "application", "applicant", "bill", "total", "form",
    "fatura", "tax", "cashier", "subtotal", "change", "signature"
]

class LayoutFeatureExtractor(BaseEstimator, TransformerMixin):
    """
    Extracts layout and statistical features from OCR tokens and document records:
    - token count
    - digit ratio
    - keyword presence indicators
    - aspect ratio
    - line count estimate
    """
    def __init__(self):
        self.invoice_kw = ["invoice", "tax", "po", "bill", "due", "balance", "vendor"]
        self.receipt_kw = ["receipt", "change", "cash", "cashier", "subtotal", "card", "store"]
        self.form_kw = ["applicant", "signature", "employee", "title", "department", "date", "form"]

    def fit(self, X, y=None):
        return self

    def transform(self, records: List[Any]) -> np.ndarray:
        feats = []
        for r in records:
            words = r.ocr_words if hasattr(r, "ocr_words") else r.get("ocr_words", [])
            text_tokens = [w.text if hasattr(w, "text") else w.get("text", "") for w in words]
            full_text = " ".join(text_tokens).lower()

            num_tokens = len(text_tokens)
            total_chars = max(len(full_text), 1)
            num_digits = sum(c.isdigit() for c in full_text)
            digit_ratio = num_digits / total_chars

            has_inv = float(any(kw in full_text for kw in self.invoice_kw))
            has_rec = float(any(kw in full_text for kw in self.receipt_kw))
            has_frm = float(any(kw in full_text for kw in self.form_kw))

            # Bounding box layout estimation
            y_coords = [w.bbox[1] if hasattr(w, "bbox") else w.get("bbox", [0, 0, 0, 0])[1] for w in words]
            unique_lines = len(set([y // 20 for y in y_coords])) if y_coords else 0

            # Aspect ratio from image if available
            aspect_ratio = 1.4  # standard portrait default
            img_path = r.image_path if hasattr(r, "image_path") else r.get("image_path", "")
            if img_path and os.path.exists(img_path):
                try:
                    from PIL import Image
                    with Image.open(img_path) as im:
                        w_img, h_img = im.size
                        aspect_ratio = float(w_img / max(h_img, 1))
                except Exception:
                    pass

            feats.append([
                float(num_tokens),
                float(digit_ratio),
                has_inv,
                has_rec,
                has_frm,
                float(unique_lines),
                float(aspect_ratio)
            ])
        return np.array(feats, dtype=np.float32)

def mask_class_keywords(text: str) -> str:
    """Masks obvious class-indicative keywords for robustness ablation."""
    pattern = re.compile(r'\b(' + '|'.join(CLASS_INDICATIVE_KEYWORDS) + r')\b', re.IGNORECASE)
    return pattern.sub("[MASKED]", text)

def apply_image_augmentations(img: np.ndarray, aug_type: str = "all") -> np.ndarray:
    """
    Applies image distortions for robustness testing:
    - blur (Gaussian blur)
    - rotation (+/- 5 degrees)
    - compression (JPEG encoding at low quality)
    - noise (Gaussian random noise)
    """
    out = img.copy()
    h, w = out.shape[:2]

    if aug_type in ["rotation", "all"]:
        angle = float(np.random.choice([-5.0, 5.0]))
        M = cv2.getRotationMatrix2D((w // 2, h // 2), angle, 1.0)
        out = cv2.warpAffine(out, M, (w, h), borderMode=cv2.BORDER_REPLICATE)

    if aug_type in ["blur", "all"]:
        out = cv2.GaussianBlur(out, (5, 5), 0)

    if aug_type in ["compression", "all"]:
        encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 30]
        _, enc = cv2.imencode('.jpg', out, encode_param)
        out = cv2.imdecode(enc, 1)

    if aug_type in ["noise", "all"]:
        noise = np.random.normal(0, 15, out.shape).astype(np.float32)
        out = np.clip(out.astype(np.float32) + noise, 0, 255).astype(np.uint8)

    return out