"""Nap du lieu da build va cung cap tra cuu / truy xuat ngu canh.

Tat ca file nam trong space/data/, sinh boi build_data.py. Chi muc TF-IDF duoc
dung mot lan luc khoi dong (1.322 doan ngu canh, duoi mot giay).
"""

import json
import math
import random
import re
import unicodedata
from pathlib import Path
from typing import Dict, List, Optional

DATA = Path(__file__).parent / "data"

MISSING = (
    "Khong co du lieu trong space/data. Chay truoc:\n"
    '  python build_data.py "<duong-dan>/Thom Paper"'
)


def _load(name: str, default=None):
    path = DATA / name
    if not path.exists():
        return default
    return json.loads(path.read_text(encoding="utf-8"))


dataset: List[dict] = _load("dataset.json", [])
facets: dict = _load("facets.json", {})
stats_cache: dict = _load("stats.json", {})
lessons: list = _load("lessons.json", [])
metrics: dict = _load("metrics.json", {})
testset: List[dict] = _load("testset.json", [])
error_groups: list = _load("errors.json", [])
build_info: dict = _load("built.json", {})

has_data = bool(dataset)

by_id = {it["id"]: it for it in dataset}
test_by_id = {it["id"]: it for it in testset}

WORD = re.compile(r"[^\w\s]", re.UNICODE)


def fold(text: str) -> List[str]:
    text = unicodedata.normalize("NFC", (text or "").lower()).replace("_", " ")
    return WORD.sub(" ", text).split()


# --------------------------------------------- chi muc ngu canh cho truy xuat
contexts: List[dict] = []
_seen = set()
for _it in dataset:
    if _it["contextId"] in _seen:
        continue
    _seen.add(_it["contextId"])
    contexts.append(
        {
            "contextId": _it["contextId"],
            "context": _it["context"],
            "grade": _it["grade"],
            "topic": _it["topic"],
            "lesson": _it["lesson"],
            "page": _it["page"],
            "section": _it["section"],
        }
    )

_df: Dict[str, int] = {}
for _c in contexts:
    for _term in set(fold(_c["context"])):
        _df[_term] = _df.get(_term, 0) + 1

_N = len(contexts) or 1


def _idf(term: str) -> float:
    return math.log(1 + _N / (1 + _df.get(term, 0)))


for _c in contexts:
    _tf: Dict[str, int] = {}
    for _term in fold(_c["context"]):
        _tf[_term] = _tf.get(_term, 0) + 1
    _weights = {t: (1 + math.log(n)) * _idf(t) for t, n in _tf.items()}
    _c["_weights"] = _weights
    _c["_norm"] = math.sqrt(sum(w * w for w in _weights.values())) or 1.0

context_count = len(contexts)


def _public(c: dict, score: Optional[float] = None) -> dict:
    out = {k: v for k, v in c.items() if not k.startswith("_")}
    if score is not None:
        out["score"] = round(score, 4)
    return out


def retrieve_contexts(
    question: str,
    grade=None,
    lesson: Optional[str] = None,
    topic: Optional[str] = None,
    limit: int = 3,
) -> List[dict]:
    """Tim doan ngu canh khop nhat voi cau hoi, co the gioi han theo lop / bai."""
    terms = fold(question)
    if not terms or not contexts:
        return []

    tf: Dict[str, int] = {}
    for term in terms:
        tf[term] = tf.get(term, 0) + 1
    q_weights = {t: (1 + math.log(n)) * _idf(t) for t, n in tf.items()}
    q_norm = math.sqrt(sum(w * w for w in q_weights.values())) or 1.0

    grade_num = int(grade) if grade not in (None, "") else None

    def pick(use_lesson: bool, use_topic: bool):
        return [
            c
            for c in contexts
            if (grade_num is None or c["grade"] == grade_num)
            and (not use_lesson or not lesson or c["lesson"] == lesson)
            and (not use_topic or not topic or c["topic"] == topic)
        ]

    # Noi long dan: bai hoc -> chu de -> ca lop. KHONG BAO GIO vuot ra ngoai lop
    # da chon, vi tra dap an lop 4 cho cau hoi lop 3 la sai hon la khong tra loi.
    candidates = pick(True, True) or pick(False, True) or pick(False, False)
    if not candidates:
        return []

    scored = []
    for c in candidates:
        dot = 0.0
        weights = c["_weights"]
        for term, w in q_weights.items():
            cw = weights.get(term)
            if cw:
                dot += w * cw
        if dot > 0:
            scored.append((dot / (q_norm * c["_norm"]), c))

    if not scored:
        return [_public(candidates[0], 0.0)] if candidates else []

    scored.sort(key=lambda pair: pair[0], reverse=True)
    return [_public(c, score) for score, c in scored[:limit]]


# -------------------------------------------------------------------- tra cuu
def search_dataset(
    q: Optional[str] = None,
    grade=None,
    topic: Optional[str] = None,
    lesson: Optional[str] = None,
    difficulty: Optional[str] = None,
    question_type: Optional[str] = None,
    split: Optional[str] = None,
    limit: int = 25,
    offset: int = 0,
) -> dict:
    limit = max(1, min(int(limit or 25), 200))
    offset = max(0, int(offset or 0))
    needle = (q or "").strip().lower()
    grade_num = int(grade) if grade not in (None, "") else None

    matched = []
    for it in dataset:
        if grade_num is not None and it["grade"] != grade_num:
            continue
        if topic and it["topic"] != topic:
            continue
        if lesson and it["lesson"] != lesson:
            continue
        if difficulty and it["difficulty"] != difficulty:
            continue
        if question_type and it["questionType"] != question_type:
            continue
        if split and it["split"] != split:
            continue
        if needle:
            hay = f"{it['id']} {it['question']} {it['answer']} {it['context']} {it['lesson']}".lower()
            if needle not in hay:
                continue
        matched.append(it)

    page = [
        {**it, "hasPredictions": it["split"] == "test" and it["id"] in test_by_id}
        for it in matched[offset : offset + limit]
    ]
    return {"total": len(matched), "offset": offset, "limit": limit, "items": page}


def sample(grade=None) -> Optional[dict]:
    grade_num = int(grade) if grade not in (None, "") else None
    pool = [t for t in testset if grade_num is None or t["grade"] == grade_num] if grade_num else testset
    return random.choice(pool) if pool else None


def dataset_stats() -> dict:
    return stats_cache or {"total": len(dataset), "contexts": context_count}


# ------------------------------- chi muc cau hoi cua tap test (goi y thay the)
_q_df: Dict[str, int] = {}
for _t in testset:
    for _term in set(fold(_t["question"])):
        _q_df[_term] = _q_df.get(_term, 0) + 1

_QN = len(testset) or 1


def _q_idf(term: str) -> float:
    return math.log(1 + _QN / (1 + _q_df.get(term, 0)))


_q_index = []
for _t in testset:
    _tf: Dict[str, int] = {}
    for _term in fold(_t["question"]):
        _tf[_term] = _tf.get(_term, 0) + 1
    _w = {t: (1 + math.log(n)) * _q_idf(t) for t, n in _tf.items()}
    _q_index.append((_t, _w, math.sqrt(sum(v * v for v in _w.values())) or 1.0))


def nearest_questions(question: str, grade=None, limit: int = 3) -> List[dict]:
    """Cau hoi trong tap test gan nghia nhat — dung khi chua co checkpoint."""
    terms = fold(question)
    if not terms or not _q_index:
        return []

    tf: Dict[str, int] = {}
    for term in terms:
        tf[term] = tf.get(term, 0) + 1
    q_weights = {t: (1 + math.log(n)) * _q_idf(t) for t, n in tf.items()}
    q_norm = math.sqrt(sum(w * w for w in q_weights.values())) or 1.0

    grade_num = int(grade) if grade not in (None, "") else None
    scored = []
    for item, weights, norm in _q_index:
        if grade_num is not None and item["grade"] != grade_num:
            continue
        dot = 0.0
        for term, w in q_weights.items():
            iw = weights.get(term)
            if iw:
                dot += w * iw
        if dot > 0:
            scored.append((dot / (q_norm * norm), item))

    scored.sort(key=lambda pair: pair[0], reverse=True)
    return [
        {
            "id": it["id"],
            "question": it["question"],
            "grade": it["grade"],
            "lesson": it["lesson"],
            "gold": it["gold"],
            "score": round(score, 4),
        }
        for score, it in scored[:limit]
    ]


# ------------------------------- tra dap an chuan trong sach theo cau hoi
# Thuc nghiem chi sinh du doan cho 1.073 cau test, nhung ca 10.961 cau deu co
# dap an chuan. Khi chua co checkpoint, it nhat tra duoc dap an trong sach.
_by_question: Dict[str, List[dict]] = {}
for _it in dataset:
    _by_question.setdefault(" ".join(fold(_it["question"])), []).append(_it)


def lookup_question(question: str, grade=None) -> Optional[dict]:
    """Tim cau hoi trung khop trong bo du lieu (ca train/dev/test)."""
    hits = _by_question.get(" ".join(fold(question)))
    if not hits:
        return None
    grade_num = int(grade) if grade not in (None, "") else None
    if grade_num is None:
        return hits[0]
    # Cung mot cau hoi co the xuat hien o nhieu lop voi dap an khac nhau —
    # chi nhan cau cua dung lop dang chon.
    for it in hits:
        if it["grade"] == grade_num:
            return it
    return None