"""Offline hybrid BM25/character-TF-IDF retrieval over the public book corpus."""

from __future__ import annotations

import json
import math
import re
import unicodedata
from collections import Counter
from pathlib import Path
from typing import Any, Iterable

from sklearn.feature_extraction.text import TfidfVectorizer


RESOURCE_DIR = Path(__file__).resolve().parent

FAMILY_KEYWORDS = {
    "writing_system": (
        "alphabet", "braille", "character", "decipher", "glyph", "letter",
        "orthography", "script", "symbol", "writing system",
    ),
    "phonetics": (
        "accent", "consonant", "metre", "phonetic", "pronounce", "rhyme",
        "stress", "syllable", "tone", "vowel",
    ),
    "phonology": (
        "alternation", "correspondence", "sound change", "sound rule",
        "underlying form",
    ),
    "noun_morphology": (
        "case", "gender", "noun", "plural", "possession", "possessive",
        "singular",
    ),
    "verb_morphology": (
        "agreement", "aspect", "conjug", "object", "person", "subject",
        "tense", "verb",
    ),
    "syntax": ("clause", "focus", "sentence", "subject", "object", "word order"),
    "number_system": ("arithmetic", "digit", "number", "numeral", "numerals"),
    "kinship_orientation": (
        "brother", "daughter", "direction", "east", "family", "father",
        "kinship", "map", "mother", "north", "sister", "son", "south", "west",
    ),
    "matching": ("correspondence", "match", "matching", "random order"),
    "translation": ("translate", "translation"),
    "fill_blanks": ("blank", "fill in", "gap", "missing"),
}

TOPIC_FAMILIES = {
    "writing systems and script decipherment": {"writing_system"},
    "phonetics, stress, tone, and versification": {"phonetics"},
    "phonological rules and sound correspondences": {"phonology", "phonetics"},
    "noun morphology and noun phrases": {"noun_morphology"},
    "verb morphology and argument structure": {"verb_morphology"},
    "syntax, word order, focus, and alignment": {"syntax"},
    "semantics and graph-based matching": {"matching"},
    "number systems": {"number_system"},
    "orientation, kinship, and other structural problems": {"kinship_orientation"},
    "general problem-solving methodology": set(FAMILY_KEYWORDS),
}


def normalize(text: Any) -> str:
    return " ".join(unicodedata.normalize("NFKC", str(text)).casefold().split())


def tokenize(text: Any) -> list[str]:
    return re.findall(r"[^\W_]+", normalize(text), flags=re.UNICODE)


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    with path.open("r", encoding="utf-8") as handle:
        return [json.loads(line) for line in handle if line.strip()]


def infer_task_families(context: str, query: str) -> list[str]:
    """Infer likely families from the problem itself, never from a gold answer."""
    text = normalize(context + "\n" + query)
    scored: list[tuple[int, str]] = []
    for family, keywords in FAMILY_KEYWORDS.items():
        score = sum(1 for keyword in keywords if keyword in text)
        if score:
            scored.append((score, family))
    return [family for _, family in sorted(scored, key=lambda item: (-item[0], item[1]))]


class BM25:
    def __init__(self, documents: Iterable[str], k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self.tokens = [tokenize(document) for document in documents]
        self.lengths = [len(tokens) for tokens in self.tokens]
        self.average_length = sum(self.lengths) / max(len(self.lengths), 1)
        self.term_frequencies = [Counter(tokens) for tokens in self.tokens]
        document_frequency: Counter[str] = Counter()
        for tokens in self.tokens:
            document_frequency.update(set(tokens))
        document_count = len(self.tokens)
        self.idf = {
            term: math.log(
                1.0 + (document_count - frequency + 0.5) / (frequency + 0.5)
            )
            for term, frequency in document_frequency.items()
        }

    def scores(self, query: str) -> list[float]:
        query_terms = Counter(tokenize(query))
        scores: list[float] = []
        for frequencies, length in zip(self.term_frequencies, self.lengths):
            score = 0.0
            normalization = self.k1 * (
                1.0 - self.b + self.b * length / max(self.average_length, 1.0)
            )
            for term, query_frequency in query_terms.items():
                frequency = frequencies.get(term, 0)
                if not frequency:
                    continue
                score += (
                    self.idf.get(term, 0.0)
                    * frequency
                    * (self.k1 + 1.0)
                    / (frequency + normalization)
                    * min(query_frequency, 3)
                )
            scores.append(score)
        return scores


class CharNgramTFIDF:
    """Cosine similarity over normalized character 3-5 grams."""

    def __init__(self, documents: Iterable[str]):
        self.vectorizer = TfidfVectorizer(
            analyzer="char_wb",
            ngram_range=(3, 5),
            preprocessor=normalize,
            lowercase=False,
            sublinear_tf=True,
            norm="l2",
        )
        self.matrix = self.vectorizer.fit_transform(documents)

    def scores(self, query: str) -> list[float]:
        query_vector = self.vectorizer.transform([query])
        return (self.matrix @ query_vector.T).toarray().ravel().tolist()


def _clip(text: Any, limit: int) -> str:
    value = str(text or "").strip()
    if len(value) <= limit:
        return value
    clipped = value[:limit].rsplit("\n", 1)[0].rstrip()
    if len(clipped) < limit // 2:
        clipped = value[:limit].rsplit(" ", 1)[0].rstrip()
    return clipped + "\n[excerpt truncated]"


class BookRetriever:
    """Retrieve methods and analogous examples from langsci/420 only."""

    def __init__(self, resource_dir: Path = RESOURCE_DIR):
        self.methods = read_jsonl(resource_dir / "book_methods.jsonl")
        self.examples = [
            row
            for row in read_jsonl(resource_dir / "book_examples.jsonl")
            if row.get("string_only_usable", False)
        ]
        self.method_bm25 = BM25(row["retrieval_text"] for row in self.methods)
        self.example_bm25 = BM25(row["retrieval_text"] for row in self.examples)
        self.method_char_tfidf = CharNgramTFIDF(
            row["retrieval_text"] for row in self.methods
        )
        self.example_char_tfidf = CharNgramTFIDF(
            row["retrieval_text"] for row in self.examples
        )

    @staticmethod
    def _topic_bonus(row: dict[str, Any], families: set[str]) -> float:
        return 1.25 * len(
            families & TOPIC_FAMILIES.get(str(row.get("topic", "")), set())
        )

    def retrieve(
        self,
        row: dict[str, Any],
        top_methods: int = 2,
        top_examples: int = 2,
        char_tfidf_weight: float = 3.0,
    ) -> dict[str, Any]:
        context = str(row.get("context", ""))
        query = str(row.get("query", ""))
        test_text = context + "\n" + query
        families = infer_task_families(context, query)
        family_set = set(families)
        expanded_query = test_text + "\n" + " ".join(families)

        ranked_examples: list[tuple[float, dict[str, Any]]] = []
        example_bm25_scores = self.example_bm25.scores(expanded_query)
        example_char_scores = self.example_char_tfidf.scores(expanded_query)
        for bm25_score, char_score, example in zip(
            example_bm25_scores, example_char_scores, self.examples
        ):
            adjusted = (
                bm25_score
                + char_tfidf_weight * char_score
                + self._topic_bonus(example, family_set)
            )
            if example.get("solution_tier") == "detailed_worked":
                adjusted += 0.35
            ranked_examples.append((adjusted, example))
        ranked_examples.sort(key=lambda item: (-item[0], item[1]["id"]))
        selected_examples = ranked_examples[:top_examples]

        linked_method_ids = {
            method_id
            for _, example in ranked_examples[:max(top_examples * 3, 5)]
            for method_id in example.get("method_ids", [])
        }
        ranked_methods: list[tuple[float, dict[str, Any]]] = []
        method_bm25_scores = self.method_bm25.scores(expanded_query)
        method_char_scores = self.method_char_tfidf.scores(expanded_query)
        for bm25_score, char_score, method in zip(
            method_bm25_scores, method_char_scores, self.methods
        ):
            adjusted = (
                bm25_score
                + char_tfidf_weight * char_score
                + self._topic_bonus(method, family_set)
            )
            if method["id"] in linked_method_ids:
                adjusted += 0.75
            ranked_methods.append((adjusted, method))
        ranked_methods.sort(key=lambda item: (-item[0], item[1]["id"]))

        return {
            "inferred_task_families": families,
            "char_tfidf_weight": char_tfidf_weight,
            "methods": [method for _, method in ranked_methods[:top_methods]],
            "examples": [example for _, example in selected_examples],
        }

    def format_for_prompt(self, result: dict[str, Any], max_chars: int = 16000) -> str:
        blocks = [
            "BOOK REFERENCE MATERIAL\n"
            "Use it for transferable methods and analogies. Do not copy a conclusion "
            "unless it is supported by the current problem."
        ]
        families = result.get("inferred_task_families", [])
        if families:
            blocks.append("Likely families inferred from the current text: " + ", ".join(families))

        for method in result["methods"]:
            blocks.append(
                f"METHOD — {method['section_title']}\n" + _clip(method.get("text"), 2800)
            )
        for example in result["examples"]:
            blocks.append(
                "\n".join(
                    [
                        f"ANALOGOUS WORKED EXAMPLE — {example.get('language', 'unknown language')}",
                        "Problem context:",
                        _clip(example.get("context"), 2400),
                        "Problem query:",
                        _clip(example.get("query"), 1200),
                        "Worked solution:",
                        _clip(example.get("reasoning_trace"), 3600),
                    ]
                )
            )

        output = "\n\n".join(blocks)
        if len(output) > max_chars:
            output = output[:max_chars].rsplit("\n", 1)[0].rstrip()
            output += "\n[retrieved material truncated]"
        return output