structured_search.py

"""
GCAS Search Engine – Structured In-Memory Lookup
=================================================

Replaces FAISS for ~90% of queries by scanning the in-memory data_store
with the structured filters extracted in query_planner.QueryPlan.

Why this is better than FAISS for this dataset
-----------------------------------------------
• The data is a structured Excel DB — not free text. Every query in the
  sample taxonomy is a WHERE-clause lookup (college + intent + filters).
• FAISS uses semantic embeddings which are dominated by "for girls / mahila"
  style phrases, causing wrong colleges to rank first.
• A Python list scan of 44,500 rows completes in < 80 ms — fast enough for
  a voice agent with no embedding roundtrip.

Lookup strategy
---------------
1. If college_name is resolved → hard-match on CollegeName.
2. Else if district / university / program → soft scan with those filters.
3. Sort/rank results:
   - college-specific: order by program name (alphabetical, fees asc for fee intent)
   - list queries: Government > Grant-in-Aid > Self-Finance, then alphabetical
4. Returns list of {table, row_index, score, data} dicts — same shape as
   indexer.search() so the rest of the pipeline is unchanged.
"""
from __future__ import annotations

import logging
import re
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Column name normalisers (handles naming inconsistencies across tables)
# ---------------------------------------------------------------------------

def _get(row: Dict, *keys: str, default: str = "") -> str:
    for k in keys:
        v = row.get(k)
        if v is not None:
            s = str(v).strip()
            if s and s.lower() not in ("nan", "none", "nat"):
                return s
    return default


def _college_name(row: Dict) -> str:
    return _get(row, "CollegeName", "College")


def _district(row: Dict) -> str:
    return _get(row, "CollegeDistrict", "DistrictName")


def _university(row: Dict) -> str:
    return _get(row, "UniversityName", "University")


def _program(row: Dict) -> str:
    return _get(row, "AdmissionName", "Program")


def _college_type(row: Dict) -> str:
    # CollegeMaster uses InstituteTypeForCollegeName
    # IntakeMaster uses CollegeProgramTypeName (Grant in Aid - Regular / Self Finance / Government)
    return _get(row, "InstituteTypeForCollegeName", "CollegeProgramTypeName")


def _medium(row: Dict) -> str:
    return _get(row, "MediumName", "Medium")


def _education_mode(row: Dict) -> str:
    return _get(row, "EducationMode")


# ---------------------------------------------------------------------------
# College-name matching
# ---------------------------------------------------------------------------

def _college_matches(row: Dict, college_name: str) -> bool:
    """
    Match a resolved canonical college name against a data-store row.
    Handles:
      - exact match
      - row name starts with / contains the canonical name
      - canonical name is a prefix of the row name (e.g. row has suffix like
        "B.COM. (ENGLISH MEDIUM) ADDITIONAL DIVISION SELF FINANCE (FOR WOMEN'S)")
    """
    cn = college_name.strip().lower()
    rn = _college_name(row).strip().lower()
    if not rn:
        return False
    return (
        rn == cn
        or rn.startswith(cn)
        or cn.startswith(rn)
        or cn in rn
    )


# ---------------------------------------------------------------------------
# College-type matching (handles "Grant in Aid - Regular" → "Grant in Aid")
# ---------------------------------------------------------------------------

_COLLEGE_TYPE_NORM: Dict[str, str] = {
    "government": "government",
    "grant in aid": "grant",
    "grant-in-aid": "grant",
    "aided": "grant",
    "self finance": "self",
    "self-finance": "self",
    "self financing": "self",
    "private": "self",
    "university department": "university",
    "constituent": "constituent",
    "diet": "diet",
}

def _type_matches(row: Dict, college_type: str) -> bool:
    ct_raw = _college_type(row).lower()
    wanted = college_type.lower()
    # normalise both sides
    ct_norm  = next((v for k, v in _COLLEGE_TYPE_NORM.items() if k in ct_raw), ct_raw)
    want_norm = next((v for k, v in _COLLEGE_TYPE_NORM.items() if k in wanted), wanted)
    return want_norm in ct_norm or ct_norm in want_norm


# ---------------------------------------------------------------------------
# Hostel type matching
# ---------------------------------------------------------------------------

def _hostel_matches(row: Dict, hostel_needed: str) -> bool:
    if hostel_needed == "girls":
        return _get(row, "GirlsHostel").lower() == "yes"
    if hostel_needed == "boys":
        return _get(row, "BoysHostel").lower() == "yes"
    # "any" — at least one hostel exists
    return (
        _get(row, "GirlsHostel").lower() == "yes"
        or _get(row, "BoysHostel").lower() == "yes"
    )


# ---------------------------------------------------------------------------
# NAAC filter
# ---------------------------------------------------------------------------

def _naac_matches(row: Dict, query: str) -> bool:
    is_naac = _get(row, "IsNAAC").lower()
    if is_naac != "yes":
        return False
    # Optional: grade filter (e.g. "A+ grade")
    grade_m = re.search(r"\bnaac\s+(a\+?|b\+?|c)\b", query, re.IGNORECASE)
    if grade_m:
        wanted_grade = grade_m.group(1).upper()
        actual_grade = _get(row, "NAACGrade").upper()
        return actual_grade.startswith(wanted_grade)
    return True


# ---------------------------------------------------------------------------
# Row matcher — applies all plan filters
# ---------------------------------------------------------------------------

def _row_matches(
    row: Dict,
    college_name: Optional[str],
    district: Optional[str],
    university: Optional[str],
    program_pattern: Optional[str],
    gender: Optional[str],
    medium: Optional[str],
    college_type: Optional[str],
    education_mode: Optional[str],
    hostel_needed: Optional[str],
    intent: str,
    normalized_query: str,
) -> bool:

    # ── College name (hard match when resolved) ───────────────────────────
    if college_name:
        if not _college_matches(row, college_name):
            return False

    # ── District ──────────────────────────────────────────────────────────
    # Skip district filter when college_name is already resolved — the college
    # name uniquely identifies the college and is more specific than district.
    # (Filtering by district would incorrectly drop colleges in a taluka whose
    #  name matches the query but whose CollegeDistrict differs from the taluka.)
    if district and not college_name:
        dist = _district(row).lower()
        if not dist:
            return False
        if district.lower() not in dist and dist not in district.lower():
            return False

    # ── University ────────────────────────────────────────────────────────
    if university:
        uni = _university(row).lower()
        if university.lower() not in uni:
            return False

    # ── Program ───────────────────────────────────────────────────────────
    if program_pattern:
        prog = _program(row).lower()
        if not prog:
            return False
        if program_pattern.lower() not in prog:
            return False

    # ── College type ──────────────────────────────────────────────────────
    if college_type:
        if not _type_matches(row, college_type):
            return False

    # ── Medium ────────────────────────────────────────────────────────────
    if medium:
        med = _medium(row).lower()
        if medium.lower() not in med:
            return False

    # ── Education mode (Girls-Only / Boys-Only) ───────────────────────────
    if education_mode:
        em = _education_mode(row).lower()
        if education_mode.lower() not in em:
            return False

    # ── Hostel availability ───────────────────────────────────────────────
    if hostel_needed:
        if not _hostel_matches(row, hostel_needed):
            return False

    # ── NAAC filter ───────────────────────────────────────────────────────
    if intent == "naac":
        if not _naac_matches(row, normalized_query):
            return False

    return True


# ---------------------------------------------------------------------------
# Scoring / ranking
# ---------------------------------------------------------------------------

_TYPE_RANK = {
    "government": 0,
    "grant": 1,
    "constituent": 2,
    "university": 3,
    "diet": 4,
    "self": 5,
}

def _rank_score(row: Dict, intent: str, gender: Optional[str]) -> float:
    """
    A deterministic score in [0, 1] for sorting list-query results.
    Higher = ranked first.
    """
    ct_raw = _college_type(row).lower()
    ct_norm = next((v for k, v in _COLLEGE_TYPE_NORM.items() if k in ct_raw), "self")
    base = 1.0 - (_TYPE_RANK.get(ct_norm, 5) / 10.0)

    # For fee queries: prefer rows that have non-zero female fee when gender=female
    if intent == "fees" and gender == "female":
        fee_str = _get(row,
                       "Programme Fees for Female Student (Annual/ Two Semesters",
                       "fee_girls")
        try:
            if float(fee_str) > 0:
                base += 0.05
        except (ValueError, TypeError):
            pass

    return min(base, 0.99)


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------

def lookup(
    plan: Any,           # QueryPlan (imported lazily to avoid circular imports)
    data_store: Dict[str, List[Dict]],
    max_rows: int = 200,
) -> List[Dict[str, Any]]:
    """
    Scan data_store with structured filters from QueryPlan.

    Returns a list of candidate dicts:
      {table, row_index, score, data}
    Same shape as indexer.search() so the rest of the pipeline works unchanged.

    Returns [] when no entities are specified (caller should fall back to FAISS).
    """
    # Guard: don't run a full-table scan with zero filters (return to FAISS)
    has_filter = any([
        plan.college_name,
        plan.district,
        plan.university,
        plan.program_pattern,
        plan.college_type,
        plan.education_mode,
        plan.hostel_needed,
        (plan.intent == "naac"),
    ])
    if not has_filter:
        logger.debug("[structured] No filters — deferring to FAISS.")
        return []

    tables = plan.preferred_tables or list(data_store.keys())
    results: List[Dict[str, Any]] = []

    for table_name in tables:
        rows = data_store.get(table_name, [])
        for idx, row in enumerate(rows):
            if _row_matches(
                row             = row,
                college_name    = plan.college_name,
                district        = plan.district,
                university      = plan.university,
                program_pattern = plan.program_pattern,
                gender          = plan.gender,
                medium          = plan.medium,
                college_type    = plan.college_type,
                education_mode  = plan.education_mode,
                hostel_needed   = plan.hostel_needed,
                intent          = plan.intent,
                normalized_query = plan.corrected_query,
            ):
                score = _rank_score(row, plan.intent, plan.gender)
                results.append({
                    "table":     table_name,
                    "row_index": idx,
                    "score":     score,
                    "data":      row,
                })
                if len(results) >= max_rows:
                    break
        if len(results) >= max_rows:
            break

    if results:
        # Sort: highest score first, then college name alphabetically
        results.sort(key=lambda c: (-c["score"], _college_name(c["data"]).lower()))
        logger.info(
            "[structured] Found %d rows | college=%r district=%r uni=%r prog=%r "
            "gender=%s intent=%s",
            len(results), plan.college_name, plan.district,
            plan.university, plan.program_pattern, plan.gender, plan.intent,
        )
    else:
        logger.info(
            "[structured] 0 rows matched — will fall back to FAISS | "
            "college=%r district=%r uni=%r prog=%r",
            plan.college_name, plan.district, plan.university, plan.program_pattern,
        )

    return results