"""
BIS SP-21 Hybrid Retrieval System
----------------------------------
Combines dense (FAISS + sentence-transformers) and sparse (BM25) search,
then re-ranks and deduplicates to return the top-5 unique IS standards.

Usage
-----
   # Index build (one-time, caches to data/processed/):
   python inference.py --build

   # Single query:
   python inference.py --query "Which standard covers 33 grade OPC cement?"

   # Batch from JSON file:
   python inference.py --input data/processed/public_test_set.json

   # Batch + write results JSON:
   python inference.py --input data/processed/public_test_set.json \
                          --output data/processed/retrieval_results.json
"""

from __future__ import annotations

import argparse
import json
import math
import re
import time
from pathlib import Path
from typing import Any

import faiss
import numpy as np
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
_ROOT = Path(__file__).resolve().parent
_CHUNKS_PATH = _ROOT / "data/processed/standards_chunks.json"
_STANDARDS_PATH = _ROOT / "data/processed/standards.json"
_EMBED_CACHE = _ROOT / "data/processed/embeddings.npy"
_INDEX_CACHE = _ROOT / "data/processed/faiss.index"

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
_MODEL_NAME = "all-MiniLM-L6-v2"
_TOP_K_DENSE = 10
_TOP_K_SPARSE = 10
_TOP_N_FINAL = 5
_SHORT_CHUNK_THRESHOLD = 40   # body words below this get a penalty
_SHORT_CHUNK_PENALTY = 0.15


# ---------------------------------------------------------------------------
# Text helpers
# ---------------------------------------------------------------------------
def _body_text(chunk_text: str) -> str:
    """Strip the leading 'IS XXXX: YYYY Title [Section]' prefix line."""
    parts = chunk_text.strip().split("\n", 1)
    return parts[1].strip() if len(parts) > 1 else parts[0]


def _tokenize(text: str) -> list[str]:
    """Lowercase word tokenizer for BM25."""
    return re.findall(r"[a-z0-9]+", text.lower())


def _bm25_doc(chunk: dict) -> list[str]:
    """
    Build the BM25 document for a chunk.

    Uses the full title from standards.json (stored in chunk["full_title"] by
    load_or_build) to avoid truncated-title misses.  Title is repeated ×4 so
    an exact title match dominates over body-text noise.
    """
    # full_title is injected by load_or_build; fall back to chunk title
    title = chunk.get("full_title") or chunk.get("title", "")
    title_tokens = _tokenize(title)
    kw_tokens = _tokenize(" ".join(chunk.get("keywords", [])))
    section_tokens = _tokenize(chunk.get("section", ""))
    text_tokens = _tokenize(_body_text(chunk.get("text", "")))
    return title_tokens * 4 + kw_tokens * 3 + section_tokens * 2 + text_tokens


def _norm_std_id(sid: str) -> str:
    return re.sub(r"\s+", " ", sid).strip().upper()


# ---------------------------------------------------------------------------
# Index builder
# ---------------------------------------------------------------------------
class RetrievalIndex:
    def __init__(
        self,
        chunks: list[dict],
        standards: list[dict],
        model: SentenceTransformer,
    ) -> None:
        self.chunks = chunks
        self.standards = standards
        self.model = model

        # Build lookup: standard_id → standard record
        self.std_lookup: dict[str, dict] = {
            _norm_std_id(s["standard_id"]): s for s in standards
        }

        # Build per-standard keyword set for boosting
        self.std_keywords: dict[str, set[str]] = {
            _norm_std_id(s["standard_id"]): set(_tokenize(" ".join(s.get("keywords", []))))
            for s in standards
        }

        # Dense index (FAISS)
        self.faiss_index: faiss.IndexFlatIP | None = None
        self.embeddings: np.ndarray | None = None

        # Sparse index (BM25)
        self.bm25: BM25Okapi | None = None
        self._bm25_docs: list[list[str]] = []

    # ------------------------------------------------------------------
    def build(self, use_cache: bool = True) -> None:
        self._build_dense(use_cache)
        self._build_sparse()

    def _build_dense(self, use_cache: bool) -> None:
        if use_cache and _EMBED_CACHE.exists() and _INDEX_CACHE.exists():
            print("Loading cached embeddings and FAISS index…")
            self.embeddings = np.load(str(_EMBED_CACHE))
            self.faiss_index = faiss.read_index(str(_INDEX_CACHE))
            return

        print(f"Encoding {len(self.chunks)} chunks with {_MODEL_NAME}…")
        texts = [c["text"] for c in self.chunks]
        emb = self.model.encode(
            texts,
            batch_size=64,
            show_progress_bar=True,
            normalize_embeddings=True,   # cosine via inner product
        )
        self.embeddings = emb.astype(np.float32)

        dim = self.embeddings.shape[1]
        self.faiss_index = faiss.IndexFlatIP(dim)
        self.faiss_index.add(self.embeddings)

        np.save(str(_EMBED_CACHE), self.embeddings)
        faiss.write_index(self.faiss_index, str(_INDEX_CACHE))
        print(f"FAISS index built: {self.faiss_index.ntotal} vectors, dim={dim}")

    def _build_sparse(self) -> None:
        print("Building BM25 index…")
        self._bm25_docs = [_bm25_doc(c) for c in self.chunks]
        self.bm25 = BM25Okapi(self._bm25_docs)
        print("BM25 index built.")


# ---------------------------------------------------------------------------
# Retrieval
# ---------------------------------------------------------------------------
class Retriever:
    def __init__(self, index: RetrievalIndex) -> None:
        self.idx = index

    def retrieve(self, query: str, top_n: int = _TOP_N_FINAL) -> list[dict]:
        t0 = time.perf_counter()

        query_tokens = _tokenize(query)

        # --- Dense retrieval ---
        q_emb = self.idx.model.encode(
            [query], normalize_embeddings=True
        ).astype(np.float32)
        dense_scores, dense_ids = self.idx.faiss_index.search(q_emb, _TOP_K_DENSE)
        dense_scores = dense_scores[0]
        dense_ids = dense_ids[0]

        # Normalise dense scores (already cosine, range ~[-1, 1] → shift to [0, 1])
        d_min, d_max = dense_scores.min(), dense_scores.max()
        d_range = d_max - d_min if d_max > d_min else 1.0
        dense_norm = {int(i): (s - d_min) / d_range for i, s in zip(dense_ids, dense_scores)}

        # --- Sparse retrieval ---
        bm25_raw = self.idx.bm25.get_scores(query_tokens)
        top_sparse_ids = np.argsort(bm25_raw)[::-1][:_TOP_K_SPARSE]
        top_sparse_scores = bm25_raw[top_sparse_ids]

        s_max = top_sparse_scores.max() if top_sparse_scores.max() > 0 else 1.0
        sparse_norm = {int(i): s / s_max for i, s in zip(top_sparse_ids, top_sparse_scores)}

        # --- Merge candidates ---
        candidate_ids = set(dense_norm) | set(sparse_norm)
        chunk_scores: dict[int, float] = {}
        for cid in candidate_ids:
            d = dense_norm.get(cid, 0.0)
            s = sparse_norm.get(cid, 0.0)
            chunk_scores[cid] = 0.6 * d + 0.4 * s   # weighted fusion

        # --- Re-ranking ---
        query_lower = query.lower()
        query_words = set(query_tokens)

        for cid, base in list(chunk_scores.items()):
            chunk = self.idx.chunks[cid]
            sid_norm = _norm_std_id(chunk["standard_id"])
            bonus = 0.0

            # Use the authoritative full title for all title-based signals
            full_title = chunk.get("full_title") or chunk.get("title", "")
            full_title_tokens = set(_tokenize(full_title))

            # Boost: keyword overlap with query
            kw_set = self.idx.std_keywords.get(sid_norm, set())
            kw_overlap = len(kw_set & query_words)
            if kw_overlap:
                bonus += 0.05 * min(kw_overlap, 4)

            # Boost: title word overlap with query (uses full, untruncated title)
            title_overlap = len(full_title_tokens & query_words)
            if title_overlap:
                bonus += 0.05 * min(title_overlap, 5)

            # Strong boost: majority of title words present in query — likely
            # the most on-point standard even if its chunk body is polluted.
            stop = {"and", "or", "for", "the", "of", "in", "a", "an", "to"}
            sig_title = full_title_tokens - stop
            sig_query = query_words - stop
            if sig_title and len(sig_title & sig_query) / len(sig_title) >= 0.6:
                bonus += 0.25

            # Boost: exact IS ID in query (user specifies a standard directly)
            if re.search(r'\bIS\s*\d+', query, re.IGNORECASE):
                for m in re.finditer(r'\bIS\s*\d+[\s:()A-Za-z\d]*:\s*\d{4}', query, re.IGNORECASE):
                    if _norm_std_id(m.group()) == sid_norm:
                        bonus += 0.20
                        break

            # Penalize very short chunks
            body_wc = len(_body_text(chunk.get("text", "")).split())
            if body_wc < _SHORT_CHUNK_THRESHOLD:
                bonus -= _SHORT_CHUNK_PENALTY

            chunk_scores[cid] = base + bonus

        # --- Group by standard_id, keep best chunk score ---
        std_best: dict[str, float] = {}
        std_chunk_repr: dict[str, dict] = {}
        for cid, score in chunk_scores.items():
            chunk = self.idx.chunks[cid]
            sid = chunk["standard_id"]
            if sid not in std_best or score > std_best[sid]:
                std_best[sid] = score
                std_chunk_repr[sid] = chunk

        # --- Sort and take top N ---
        ranked = sorted(std_best.items(), key=lambda x: x[1], reverse=True)[:top_n]

        results = []
        for sid, score in ranked:
            std_rec = self.idx.std_lookup.get(_norm_std_id(sid), {})
            results.append({
                "standard_id": sid,
                "title": std_rec.get("title", std_chunk_repr[sid].get("title", "")),
                "category": std_rec.get("category", std_chunk_repr[sid].get("category", "")),
                "score": round(float(score), 4),
                "matched_section": std_chunk_repr[sid].get("section", ""),
            })

        latency = time.perf_counter() - t0
        return results, latency


# ---------------------------------------------------------------------------
# Index load/build helper
# ---------------------------------------------------------------------------
def load_or_build(force_rebuild: bool = False) -> tuple[RetrievalIndex, Retriever]:
    with open(_CHUNKS_PATH, encoding="utf-8") as f:
        chunks = json.load(f)
    with open(_STANDARDS_PATH, encoding="utf-8") as f:
        standards = json.load(f)

    # Attach full title + keywords from standards.json to each chunk.
    # full_title ensures the BM25 document uses the authoritative (untruncated)
    # title from the structured record, not whatever ended up in the chunk prefix.
    std_map = {s["standard_id"]: s for s in standards}
    for c in chunks:
        rec = std_map.get(c["standard_id"], {})
        c["full_title"] = rec.get("title", c.get("title", ""))
        c["keywords"] = rec.get("keywords", [])

    print(f"Loaded {len(chunks)} chunks, {len(standards)} standards.")
    model = SentenceTransformer(_MODEL_NAME)
    index = RetrievalIndex(chunks, standards, model)
    index.build(use_cache=not force_rebuild)
    return index, Retriever(index)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def _format_result(
    query_id: str,
    query: str,
    results: list[dict],
    latency: float,
    expected_standards: list[str] | None = None,
) -> dict:
    out: dict[str, Any] = {
        "id": query_id,
        "query": query,
        "retrieved_standards": [r["standard_id"] for r in results],
        "details": results,
        "latency_seconds": round(latency, 4),
    }
    if expected_standards is not None:
        out["expected_standards"] = expected_standards
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description="BIS SP-21 Hybrid Retrieval")
    parser.add_argument("--build", action="store_true", help="Force rebuild of FAISS index")
    parser.add_argument("--query", type=str, help="Single query string")
    parser.add_argument("--input", type=str, help="JSON file with list of {id, query} objects")
    parser.add_argument("--output", type=str, help="Write JSON results to this file")
    args = parser.parse_args()

    index, retriever = load_or_build(force_rebuild=args.build)

    if args.query:
        results, latency = retriever.retrieve(args.query)
        out = _format_result("Q0", args.query, results, latency)
        print("\n" + "=" * 60)
        print(f"Query : {args.query}")
        print(f"Latency: {latency:.3f}s")
        print("\nTop results:")
        for i, r in enumerate(results, 1):
            print(f"  {i}. {r['standard_id']} — {r['title']}")
            print(f"     Category: {r['category']}  |  Section: {r['matched_section']}  |  Score: {r['score']}")
        if args.output:
            Path(args.output).write_text(json.dumps([out], indent=2, ensure_ascii=False), encoding="utf-8")
        return

    if args.input:
        with open(args.input, encoding="utf-8") as f:
            queries = json.load(f)

        all_results = []
        latencies = []
        for q in queries:
            qid = q.get("id", "?")
            qtext = q.get("query", "")
            results, latency = retriever.retrieve(qtext)
            latencies.append(latency)
            expected = q.get("expected_standards", [])
            out = _format_result(qid, qtext, results, latency, expected_standards=expected or None)
            all_results.append(out)
            hit = any(r["standard_id"] in expected for r in results)
            print(f"[{qid}] latency={latency:.3f}s  hit={hit}  retrieved={[r['standard_id'] for r in results]}")

        print(f"\nAvg latency: {sum(latencies)/len(latencies):.3f}s  |  Max: {max(latencies):.3f}s")

        # Simple Hit@5 eval
        hits = 0
        for q, out in zip(queries, all_results):
            expected = set(q.get("expected_standards", []))
            if expected & set(out["retrieved_standards"]):
                hits += 1
        print(f"Hit@5: {hits}/{len(queries)} = {hits/len(queries):.1%}")

        if args.output:
            Path(args.output).write_text(
                json.dumps(all_results, indent=2, ensure_ascii=False), encoding="utf-8"
            )
            print(f"Results written to {args.output}")
        return

    # Default: demo with one example query
    demo_query = (
        "Which standard specifies chemical and physical requirements "
        "for 33 grade Ordinary Portland Cement?"
    )
    results, latency = retriever.retrieve(demo_query)
    out = _format_result("DEMO", demo_query, results, latency)

    print("\n" + "=" * 60)
    print(f"Demo query : {demo_query}")
    print(f"Latency    : {latency:.3f}s")
    print("\nTop-5 retrieved standards:")
    for i, r in enumerate(results, 1):
        print(f"  {i}. {r['standard_id']} — {r['title']}")
        print(f"     Category : {r['category']}")
        print(f"     Section  : {r['matched_section']}")
        print(f"     Score    : {r['score']}")
    print("=" * 60)


if __name__ == "__main__":
    main()