from typing import Dict, Any, List
import numpy as np
from ..data.schema import DocumentRecord, OCRWord, CLASS_FIELD_SCHEMAS

class ProximityHeuristicsExtractor:
    """
    Baseline B: Spatial Bounding-Box Proximity Extractor.
    Associates values located directly to the right or below target key words.
    """
    def __init__(self):
        self.triggers = {
            "invoice_number": ["invoice", "inv", "number", "no", "#"],
            "invoice_date": ["date", "dated"],
            "total_amount": ["total", "due", "balance", "amount"],
            "company_name": ["store", "mart", "inc", "ltd", "cafe"],
            "date": ["date", "time"],
            "applicant_name": ["name", "applicant", "employee", "signature"],
            "address": ["address", "street", "city", "place"]
        }

    def _find_words_near(self, key_word: OCRWord, all_words: List[OCRWord], max_dx=300, max_dy=80) -> List[OCRWord]:
        kx0, ky0, kx1, ky1 = key_word.bbox
        candidates = []
        for w in all_words:
            if w == key_word:
                continue
            wx0, wy0, wx1, wy1 = w.bbox

            # 1. Directly to the right (same line)
            is_right = (wx0 >= kx0) and (wx0 - kx1 <= max_dx) and (abs(wy0 - ky0) <= 25)

            # 2. Directly below (within column)
            is_below = (wy0 >= ky1) and (wy0 - ky1 <= max_dy) and (abs(wx0 - kx0) <= 100)

            if is_right or is_below:
                dist = np.sqrt((wx0 - kx1)**2 + (wy0 - ky0)**2)
                candidates.append((dist, w))

        candidates.sort(key=lambda x: x[0])
        return [c[1] for c in candidates]

    def extract(self, record: DocumentRecord, doc_class: str = None) -> Dict[str, Any]:
        cls_name = doc_class or record.doc_class
        schema_fields = CLASS_FIELD_SCHEMAS.get(cls_name, [])
        extracted = {f: None for f in schema_fields}
        all_words = record.ocr_words

        for field in schema_fields:
            trig_list = self.triggers.get(field, [])
            matched_key_word = None
            for w in all_words:
                w_lower = w.text.lower()
                if any(t in w_lower for t in trig_list):
                    matched_key_word = w
                    break

            if matched_key_word:
                near_words = self._find_words_near(matched_key_word, all_words)
                if near_words:
                    # Collect tokens on same line
                    val_tokens = [near_words[0].text]
                    for nw in near_words[1:4]:
                        if abs(nw.bbox[1] - near_words[0].bbox[1]) <= 20:
                            val_tokens.append(nw.text)
                    extracted[field] = " ".join(val_tokens)

        return extracted