query_planner.py

"""
GCAS Search Engine – Structured Query Planner
==============================================

Converts a raw student query (English / Hindi / Gujarati) into a
machine-readable QueryPlan with every filter needed for a direct
in-memory lookup — no embeddings required for ~90% of queries.

Pipeline
--------
  raw query
      ↓ normalizer (Hindi phrases → English, transliteration)
      ↓ resolve_college_in_query (keyword map: "M N College" → canonical)
      ↓ analyze_entities (district / taluka geo matching)
      ↓ _extract_* helpers (program, gender, category, medium, college_type…)
      → QueryPlan dataclass
"""
from __future__ import annotations

import logging
import re
from dataclasses import dataclass, field
from typing import List, Optional

import query_processor
from field_filter import detect_intent, get_preferred_tables

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# QueryPlan dataclass
# ---------------------------------------------------------------------------

@dataclass
class QueryPlan:
    """Everything needed to execute a structured data lookup."""

    # ── What the user wants ────────────────────────────────────────────────
    intent: str = "general"   # fees|hostel|cutoff|naac|contact|facilities|courses|general

    # ── Named entities (None = not specified) ─────────────────────────────
    college_name: Optional[str] = None      # canonical resolved name
    college_partial: Optional[str] = None   # raw words when not resolved by keyword map
    district: Optional[str] = None          # "Mehsana", "Surat" …
    university: Optional[str] = None        # substring match on UniversityName
    program_pattern: Optional[str] = None   # substring to match AdmissionName / Program

    # ── Attribute filters ─────────────────────────────────────────────────
    gender: Optional[str] = None            # "female" | "male"
    category: Optional[str] = None          # "general"|"sc"|"st"|"sebc"|"obc"|"ews"|"ph"
    medium: Optional[str] = None            # "English" | "Gujarati" | "Hindi"
    college_type: Optional[str] = None      # "Government" | "Self Finance" | "Grant in Aid"
    education_mode: Optional[str] = None    # "Girls-Only" | "Boys-Only"
    hostel_needed: Optional[str] = None     # "girls" | "boys" | "any"
    instate: bool = True                    # outstate flag for cutoff queries

    # ── Routing ───────────────────────────────────────────────────────────
    preferred_tables: List[str] = field(default_factory=list)
    scope: str = "specific"                 # "specific" | "list" | "aggregate"

    # ── Metadata for API response ─────────────────────────────────────────
    corrected_query: str = ""
    detected_language: str = "en"
    entity_corrections: List[dict] = field(default_factory=list)


# ---------------------------------------------------------------------------
# University alias map  (abbreviation → substring to match UniversityName)
# ---------------------------------------------------------------------------

_UNIVERSITY_ALIASES: dict = {
    r"\bgu\b": "Gujarat University",
    r"\bgtu\b": "Gujarat Technological",
    r"\bmsu\b": "Maharaja Sayajirao",
    r"\bhngu\b": "Hemchandracharya North Gujarat",
    r"\bvnsgu\b": "Veer Narmad South Gujarat",
    r"\bvnsgu\b": "Veer Narmad South Gujarat",
    r"\bspu\b": "Sardar Patel University",
    r"\bbknmu\b": "Bhakta Kavi Narsinh Mehta",
    r"\bkskvku\b": "Kutch University",
    r"\bbaou\b": "Babasaheb Ambedkar Open",
    r"\bcru\b": "Children University",
    r"\biite\b": "IITE",
    r"\bsssuv\b": "Somnath Sanskrit",
    r"\bgnlu\b": "Gujarat National Law",
    r"\bpdpu\b": "Pandit Deendayal Energy",
    r"\bddu\b": "Dharmsinh Desai",
    r"\baradhna\b": "Aradhna",
}

# ---------------------------------------------------------------------------
# Program alias map  (student shorthand → substring present in AdmissionName)
# ---------------------------------------------------------------------------

_PROGRAM_ALIASES: list = [
    # Most-specific first (longer patterns before shorter ones)
    (r"\bbba\s*llb\b|\bbbа\s*law\b",                    "B.B.A.-LL.B"),
    (r"\bba\s*llb\b|\bba\s*law\b",                      "B.A.-LL.B"),
    (r"\bbcom\s*llb\b|\bcommerce\s*law\b",               "B.COM. LL.B"),
    (r"\bllb\b|\blaw\b|\blegal\b",                       "LAW"),
    (r"\bba\s*b\.?ed\b|\bba\s*bed\b|\bba[-\s]bed\b",    "B.A-B.ED"),
    (r"\bbsc\s*b\.?ed\b|\bbsc[-\s]bed\b",                "B.SC-B.ED"),
    (r"\bb\.?ed\s*m\.?ed\b|\bbed[-\s]med\b",             "B.ED-M.ED"),
    (r"\bm\.?ed\b|\bmaster.*education\b",                "MASTER OF EDUCATION"),
    (r"\bb\.?ed\b|\bbachelor.*education\b",              "BACHELOR OF EDUCATION"),
    (r"\bm\.?b\.?a\b|\bmaster.*business\b",             "BUSINESS ADMINISTRATION"),
    (r"\bm\.?c\.?a\b|\bmaster.*computer.*appl",          "MASTER OF COMPUTER APPLICATION"),
    (r"\bm\.?sc\b|\bmaster.*science\b",                  "MASTER OF SCIENCE"),
    (r"\bm\.?com\b|\bmaster.*commerce\b",                "MASTER OF COMMERCE"),
    (r"\bm\.?a\b|\bmaster.*arts\b",                      "MASTER OF ARTS"),
    (r"\bm\.?s\.?w\b|\bmaster.*social.*work\b",         "MASTER OF SOCIAL WORK"),
    (r"\bm\.?lib\b|\bmaster.*library\b",                 "MASTER OF LIBRARY"),
    (r"\bb\.?c\.?a\b|\bbachelor.*computer.*appl",        "COMPUTER APPLICATION"),
    (r"\bb\.?b\.?a\b|\bbachelor.*business\b",            "BUSINESS ADMINISTRATION"),
    (r"\bb\.?com\b|\bbachelor.*commerce\b",              "COMMERCE"),
    (r"\bb\.?sc\b|\bbachelor.*science\b",                "BACHELOR OF SCIENCE"),
    (r"\bb\.?s\.?w\b|\bbachelor.*social.*work\b",        "SOCIAL WORK"),
    (r"\bb\.?lib\b|\bbachelor.*library\b",               "LIBRARY"),
    (r"\bnursing\b",                                      "NURSING"),
    (r"\bpharmac",                                        "PHARMAC"),
    (r"\bphysical.*educ|\bp\.?e\.?d\b",                  "PHYSICAL EDUCATION"),
    (r"\bjournalism\b|\bmass.*comm",                      "JOURNALISM"),
    (r"\bfine.*art\b|\bapplied.*art\b",                  "ART"),
    (r"\bdesign\b",                                       "DESIGN"),
    (r"\bdrama\b|\bperforming.*art",                      "PERFORMING"),
    (r"\bmusic\b",                                        "MUSIC"),
    (r"\bhome.*sci\b",                                    "HOME SCIENCE"),
    (r"\bsanskrit\b",                                     "SANSKRIT"),
    (r"\bhonours?\b|\bhonoure?d\b|\bnep\b",              "HONORS"),
    (r"\b(?:b\.?a\.?|bachelor.*arts?)\b",                "BACHELOR OF ARTS"),
    # Plain English words (after removing over-eager normalizer expansions)
    (r"\bcommerce\b",                                     "COMMERCE"),
    (r"\b(?:arts?)\b",                                    "ARTS"),
    (r"\bscience\b",                                      "SCIENCE"),
    (r"\bteaching\b",                                     "EDUCATION"),
    (r"\bcomputer\b",                                     "COMPUTER"),
]

# ---------------------------------------------------------------------------
# Gender keywords
# ---------------------------------------------------------------------------

_FEMALE_RE = re.compile(
    r"\b(girl|girls|female|ladies|lady|women|woman|"
    r"ladki|ladkiyon|ladkiyaan|mahila|stri)\b",
    re.IGNORECASE,
)
_MALE_RE = re.compile(
    r"\b(boy|boys|male|gents|gent|men|man|"
    r"ladka|ladkon|ladke|purush)\b",
    re.IGNORECASE,
)

# ---------------------------------------------------------------------------
# Category keywords
# ---------------------------------------------------------------------------

_CATEGORY_PATTERNS: list = [
    ("ph",      re.compile(r"\b(ph|divyang|handicap|disabled|pwd|pwbd)\b", re.IGNORECASE)),
    ("ews",     re.compile(r"\b(ews|economically\s+weaker)\b", re.IGNORECASE)),
    ("sebc",    re.compile(r"\b(sebc|obc|other\s+backward)\b", re.IGNORECASE)),
    ("sc",      re.compile(r"\b(sc|scheduled\s+caste|dalit)\b", re.IGNORECASE)),
    ("st",      re.compile(r"\b(st|scheduled\s+tribe|tribal|adivasi)\b", re.IGNORECASE)),
]

# ---------------------------------------------------------------------------
# Medium keywords
# ---------------------------------------------------------------------------

_MEDIUM_RE = re.compile(
    r"\b(english|gujarati|hindi|sanskrit)\s+medium\b",
    re.IGNORECASE,
)

# ---------------------------------------------------------------------------
# College type keywords
# ---------------------------------------------------------------------------

_COLLEGE_TYPE_MAP: list = [
    ("Government",    re.compile(r"\b(government|sarkari|govt)\b", re.IGNORECASE)),
    ("Grant in Aid",  re.compile(r"\b(grant.in.aid|grant|aided)\b", re.IGNORECASE)),
    ("Self Finance",  re.compile(r"\b(self.financ|private|private)\b", re.IGNORECASE)),
]

# ---------------------------------------------------------------------------
# Outstate keywords
# ---------------------------------------------------------------------------

_OUTSTATE_RE = re.compile(
    r"\b(outstate|out.state|outside gujarat|rajasthan|maharashtra|mp|"
    r"main.*rajasthan|rajasthan.*se|doosre state)\b",
    re.IGNORECASE,
)


# ---------------------------------------------------------------------------
# Extraction helpers
# ---------------------------------------------------------------------------

def _extract_university(query: str) -> Optional[str]:
    for pattern, canonical in _UNIVERSITY_ALIASES.items():
        if re.search(pattern, query, re.IGNORECASE):
            return canonical
    return None


def _extract_program(query: str) -> Optional[str]:
    for pattern, canonical in _PROGRAM_ALIASES:
        if re.search(pattern, query, re.IGNORECASE):
            return canonical
    return None


def _extract_gender(query: str) -> Optional[str]:
    if _FEMALE_RE.search(query):
        return "female"
    if _MALE_RE.search(query):
        return "male"
    return None


def _extract_category(query: str) -> Optional[str]:
    for cat, pattern in _CATEGORY_PATTERNS:
        if pattern.search(query):
            return cat
    return None


def _extract_medium(query: str) -> Optional[str]:
    m = _MEDIUM_RE.search(query)
    if m:
        return m.group(1).capitalize()
    return None


def _extract_college_type(query: str) -> Optional[str]:
    for ct, pattern in _COLLEGE_TYPE_MAP:
        if pattern.search(query):
            return ct
    return None


def _extract_hostel_needed(intent: str, gender: Optional[str]) -> Optional[str]:
    if intent != "hostel":
        return None
    if gender == "female":
        return "girls"
    if gender == "male":
        return "boys"
    return "any"


def _extract_education_mode(query: str, gender: Optional[str]) -> Optional[str]:
    if re.search(r"\b(girls.only|only.*girls|sirf.*ladki|mahila.*college|women.*only)\b",
                 query, re.IGNORECASE):
        return "Girls-Only"
    if re.search(r"\b(boys.only|only.*boys|sirf.*ladka)\b",
                 query, re.IGNORECASE):
        return "Boys-Only"
    return None


def _determine_scope(plan: "QueryPlan") -> str:
    """
    'specific'  → one college or one program+college → return 1-3 rows
    'list'      → district/university scoped → return top-N
    'aggregate' → count or average across many rows
    """
    if plan.college_name:
        return "specific"
    if plan.district or plan.university:
        return "list"
    if plan.program_pattern and not plan.district:
        return "list"
    return "list"


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

def build_query_plan(raw_query: str) -> QueryPlan:
    """
    Run the full NLP pipeline and return a QueryPlan ready for
    structured_search.lookup() or FAISS fallback.
    """
    # ── Step 1: existing normalisation + entity resolution ─────────────────
    processed = query_processor.process(raw_query)
    q = processed.final_query          # normalised, entities rewritten
    q_lower = q.lower()

    # ── Step 2: Intent ───────────────────────────────────────────────────
    intent = detect_intent(q)

    # ── Step 3: Entities from the processed pipeline ─────────────────────
    college_name: Optional[str] = None
    district: Optional[str] = None

    for e in processed.detected_entities:
        if e.entity_type == "college" and not college_name:
            college_name = e.matched
        elif e.entity_type in ("district", "city") and not district:
            # Only take actual district entities (not taluka) as district filter.
            # Talukas often appear inside college names (e.g. "VISNAGAR" in
            # "M. N. COLLEGE, VISNAGAR") and would cause false mismatches because
            # CollegeDistrict stores the district ("Mehsana"), not the taluka.
            district = e.matched

    # ── Step 4: Additional entity extractions ────────────────────────────
    university   = _extract_university(q_lower)
    program_pat  = _extract_program(q_lower)
    gender       = _extract_gender(q_lower)
    category     = _extract_category(q_lower)
    medium       = _extract_medium(q)
    college_type = _extract_college_type(q_lower)
    instate      = not bool(_OUTSTATE_RE.search(raw_query))
    edu_mode     = _extract_education_mode(q_lower, gender)
    hostel_need  = _extract_hostel_needed(intent, gender)

    # ── Step 5: Table routing ─────────────────────────────────────────────
    preferred_tables = get_preferred_tables(intent) or []

    # ── Step 6: Entity corrections for API response ───────────────────────
    entity_corrections = [
        {
            "original_span":  e.query_span,
            "corrected_to":   e.matched,
            "entity_type":    e.entity_type,
            "match_score":    e.score,
            "method":         e.method,
        }
        for e in processed.detected_entities
        if e.query_span.lower() != e.matched.lower()
    ]

    plan = QueryPlan(
        intent           = intent,
        college_name     = college_name,
        district         = district,
        university       = university,
        program_pattern  = program_pat,
        gender           = gender,
        category         = category,
        medium           = medium,
        college_type     = college_type,
        education_mode   = edu_mode,
        hostel_needed    = hostel_need,
        instate          = instate,
        preferred_tables = preferred_tables,
        corrected_query  = q,
        detected_language = processed.detected_language,
        entity_corrections = entity_corrections,
    )
    plan.scope = _determine_scope(plan)

    logger.info(
        "[plan] intent=%s college=%r district=%r uni=%r prog=%r "
        "gender=%s cat=%s medium=%s type=%s scope=%s",
        plan.intent, plan.college_name, plan.district, plan.university,
        plan.program_pattern, plan.gender, plan.category,
        plan.medium, plan.college_type, plan.scope,
    )

    return plan