app/services/rule_based_ner.py

"""
Rule-based Vietnamese NER fallback.

Used by ``NERService`` when the trained ``Model_NER_Admin_v5`` directory is not
present in the repository. It reproduces the entity schema the rest of the
pipeline expects -- TOURISM, OLD_WARD, NEW_WARD, PROVINCE, DISTRICT, CITY --
using administrative dictionaries (``Wards.csv`` / ``Provinces.csv``), the POI
list (``dulich.xlsx``) and a set of head-noun + proper-noun regular patterns.

It returns the same record shape produced by a HuggingFace ``ner`` pipeline with
``aggregation_strategy="simple"``::

    {"entity_group": str, "score": float, "word": str, "start": int, "end": int}

This is intentionally lightweight and deterministic so the full ``/search`` flow
runs end-to-end without any GPU or downloaded model. Replace it with the trained
model (see ``experiments/train_ner.py``) for production-grade accuracy.
"""

from __future__ import annotations

import csv
import logging
import re
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

logger = logging.getLogger(__name__)

# Repo root = parents[2] of app/services/rule_based_ner.py
_REPO_ROOT = Path(__file__).resolve().parents[2]
_DATASETS = _REPO_ROOT / "datasets"

# Administrative head keywords -> coarse type
_ADMIN_HEADS: List[Tuple[str, str]] = [
    ("thị trấn", "WARD"),
    ("thị xã", "CITY"),
    ("thành phố", "CITY"),
    ("phường", "WARD"),
    ("xã", "WARD"),
    ("quận", "DISTRICT"),
    ("huyện", "DISTRICT"),
    ("tỉnh", "PROVINCE"),
    ("tp.", "CITY"),
    ("tp", "CITY"),
]

# Tourism head nouns (lowercased). Order matters: longer phrases first.
_TOURISM_HEADS: List[str] = [
    "khu du lịch", "vườn quốc gia", "khu di tích", "khu bảo tồn",
    "bảo tàng", "nhà thờ", "phố cổ", "công viên", "khu tưởng niệm",
    "chùa", "đền", "miếu", "đình", "lăng", "hồ", "thác", "núi",
    "chợ", "cầu", "thành", "dinh", "làng", "biển", "đảo", "suối",
]

# Words that should stop proper-noun consumption even if capitalized.
_STOP_TOKENS: Set[str] = {
    "Sau", "Trước", "Hiện", "Khi", "Nếu", "Và", "Hay", "Hoặc",
    "Là", "Của", "Ở", "Tại", "Cho", "Theo", "Với", "Từ",
}

# Context cues to disambiguate OLD vs NEW ward.
_NEW_CUES = ("sau sáp nhập", "sau khi sáp nhập", "tên mới", "hiện nay", "bây giờ", "đổi thành", "mới")
_OLD_CUES = ("trước sáp nhập", "trước khi sáp nhập", "tên cũ", "ban đầu", "trước đây", "cũ")


def _is_proper_token(tok: str) -> bool:
    """A name token: first alphabetic char is uppercase, not a stop word."""
    if not tok:
        return False
    if tok in _STOP_TOKENS:
        return False
    first = tok[0]
    return first.isalpha() and first.isupper()


class RuleBasedNER:
    """Dictionary + regex NER used as an offline fallback."""

    def __init__(self) -> None:
        self.old_wards: Set[str] = set()
        self.new_wards: Set[str] = set()
        self.provinces: Set[str] = set()
        self.districts: Set[str] = set()
        self.tourism_names: List[str] = []  # kept longest-first for greedy match
        self._load_dictionaries()

    # ------------------------------------------------------------------ loaders
    def _load_dictionaries(self) -> None:
        wards_csv = _DATASETS / "Wards.csv"
        if wards_csv.exists():
            try:
                with open(wards_csv, encoding="utf-8") as f:
                    for row in csv.DictReader(f):
                        if row.get("old_name"):
                            self.old_wards.add(row["old_name"].strip().lower())
                        if row.get("new_name"):
                            self.new_wards.add(row["new_name"].strip().lower())
                        if row.get("district_city"):
                            self.districts.add(row["district_city"].strip().lower())
            except Exception as e:  # pragma: no cover - defensive
                logger.warning(f"RuleBasedNER: could not read Wards.csv: {e}")

        prov_csv = _DATASETS / "Provinces.csv"
        if prov_csv.exists():
            try:
                with open(prov_csv, encoding="utf-8") as f:
                    for row in csv.DictReader(f):
                        for key in ("old_province_name", "new_province_name"):
                            if row.get(key):
                                self.provinces.add(row[key].strip().lower())
            except Exception as e:  # pragma: no cover
                logger.warning(f"RuleBasedNER: could not read Provinces.csv: {e}")

        self._load_tourism_names()
        logger.info(
            "RuleBasedNER ready: %d old wards, %d new wards, %d provinces, %d POIs",
            len(self.old_wards), len(self.new_wards), len(self.provinces),
            len(self.tourism_names),
        )

    def _load_tourism_names(self) -> None:
        xlsx = _DATASETS / "dulich.xlsx"
        if not xlsx.exists():
            return
        try:
            import openpyxl

            wb = openpyxl.load_workbook(xlsx, read_only=True, data_only=True)
            ws = wb.active
            rows = ws.iter_rows(values_only=True)
            next(rows, None)  # skip header (first column = Tên địa điểm)
            names: Set[str] = set()
            for row in rows:
                if row and row[0]:
                    name = str(row[0]).strip()
                    if len(name) >= 3:
                        names.add(name)
            # Greedy longest-first matching avoids partial hits.
            self.tourism_names = sorted(names, key=len, reverse=True)
            wb.close()
        except Exception as e:  # pragma: no cover
            logger.warning(f"RuleBasedNER: could not read dulich.xlsx: {e}")

    # ------------------------------------------------------------------- extract
    def extract_entities(self, text: str) -> List[Dict]:
        if not text or not text.strip():
            return []

        spans: List[Tuple[int, int, str, float]] = []  # start, end, label, score
        lower = text.lower()

        # 1) Tourism dictionary (exact substring, longest-first, non-overlapping)
        taken: List[Tuple[int, int]] = []

        def _overlaps(s: int, e: int) -> bool:
            return any(not (e <= ts or s >= te) for ts, te in taken)

        for name in self.tourism_names:
            idx = lower.find(name.lower())
            if idx != -1 and not _overlaps(idx, idx + len(name)):
                spans.append((idx, idx + len(name), "TOURISM", 0.90))
                taken.append((idx, idx + len(name)))

        # 2) Tourism via head noun + proper nouns (catches POIs not in the table)
        for s, e, val in self._scan_heads(text, _TOURISM_HEADS):
            if not _overlaps(s, e):
                spans.append((s, e, "TOURISM", 0.75))
                taken.append((s, e))

        # 3) Administrative entities
        for s, e, val, coarse in self._scan_admin_heads(text):
            if _overlaps(s, e):
                continue
            label = self._admin_label(val, coarse, lower)
            spans.append((s, e, label, 0.85))
            taken.append((s, e))

        spans.sort(key=lambda x: x[0])
        return [
            {
                "entity_group": label,
                "score": float(score),
                "word": text[s:e],
                "start": int(s),
                "end": int(e),
            }
            for (s, e, label, score) in spans
        ]

    # ------------------------------------------------------------------- helpers
    def _scan_heads(self, text: str, heads: List[str]) -> List[Tuple[int, int, str]]:
        """Find `<head> <ProperNoun...>` occurrences, return (start,end,value)."""
        results: List[Tuple[int, int, str]] = []
        lower = text.lower()
        for head in heads:
            for m in re.finditer(r"\b" + re.escape(head) + r"\b", lower):
                start = m.start()
                name_end = self._consume_proper_nouns(text, m.end())
                if name_end > m.end():
                    results.append((start, name_end, text[start:name_end]))
        return results

    def _scan_admin_heads(self, text: str) -> List[Tuple[int, int, str, str]]:
        results: List[Tuple[int, int, str, str]] = []
        lower = text.lower()
        for head, coarse in _ADMIN_HEADS:
            pattern = r"\b" + re.escape(head) + r"\b"
            for m in re.finditer(pattern, lower):
                start = m.start()
                name_end = self._consume_proper_nouns(text, m.end())
                if name_end > m.end():
                    results.append((start, name_end, text[start:name_end], coarse))
        return results

    @staticmethod
    def _consume_proper_nouns(text: str, pos: int, max_tokens: int = 4) -> int:
        """From char index `pos`, consume up to `max_tokens` proper-noun tokens."""
        end = pos
        count = 0
        n = len(text)
        i = pos
        while i < n and count < max_tokens:
            while i < n and text[i].isspace():
                i += 1
            if i >= n:
                break
            j = i
            while j < n and not text[j].isspace() and text[j] not in ",.?!;:()\"'":
                j += 1
            token = text[i:j]
            if _is_proper_token(token):
                end = j
                count += 1
                i = j
            else:
                break
        return end

    def _admin_label(self, value: str, coarse: str, lower_text: str) -> str:
        """Map a coarse admin type + dictionary lookup to the final label."""
        # Strip the leading head word to get the bare name for dictionary lookup.
        bare = re.sub(
            r"^(phường|xã|thị trấn|thị xã|thành phố|quận|huyện|tỉnh|tp\.?)\s+",
            "", value.strip(), flags=re.IGNORECASE,
        ).strip().lower()

        if coarse == "PROVINCE":
            return "PROVINCE"
        if coarse == "CITY":
            return "CITY"
        if coarse == "DISTRICT":
            return "DISTRICT"

        # WARD: decide OLD vs NEW
        in_old = bare in self.old_wards
        in_new = bare in self.new_wards
        if in_new and not in_old:
            return "NEW_WARD"
        if in_old and not in_new:
            return "OLD_WARD"
        # Ambiguous or unknown -> use surrounding context cues.
        if any(c in lower_text for c in _NEW_CUES) and not any(c in lower_text for c in _OLD_CUES):
            return "NEW_WARD"
        return "OLD_WARD"  # default: queries usually reference the historical name

    def get_info(self) -> Dict:
        return {
            "status": "available",
            "model_name_or_path": "rule_based_fallback",
            "labels": ["TOURISM", "OLD_WARD", "NEW_WARD", "PROVINCE", "DISTRICT", "CITY"],
            "aggregation_strategy": "rule_based",
            "dictionaries": {
                "old_wards": len(self.old_wards),
                "new_wards": len(self.new_wards),
                "provinces": len(self.provinces),
                "tourism_pois": len(self.tourism_names),
            },
        }