src/bm25.py

"""
src/bm25.py  —  BM25 keyword retrieval
Uses LangChain's BM25Retriever with the custom tokenizer from utils.py.

Document schema (one LangChain Document per product):
  page_content : text BM25 scores against =
                   title + features + description + categories +
                   details (flattened) + store + top-k review titles & texts
  metadata     : structured fields for display in app.py
                   (parent_asin, title, main_category, price, store,
                    categories, features, description, details, top_reviews)

Data source expected: HuggingFace Dataset objects as loaded in
  milestone1_exploration.ipynb via load_dataset("McAuley-Lab/Amazon-Reviews-2023", ...)
  OR the saved .jsonl subsets in data/raw/.
"""

import json
import pickle
from pathlib import Path
from typing import Any
import sys
from datasets import Dataset
from langchain_community.retrievers import BM25Retriever
from langchain_core.documents import Document
ROOT_FOLDER = Path(__file__).resolve().parent.parent

sys.path.append(str(ROOT_FOLDER))
from src.utils import simple_tokenize, extract_image
from src.eda_helpers import get_best_reviews


# ── field helpers ─────────────────────────────────────────────────────────────

def _coerce_str(value: Any) -> str:
    """Safely flatten any metadata field to a plain string."""
    if value is None:
        return ""
    if isinstance(value, list):
        return " ".join(_coerce_str(v) for v in value)
    if isinstance(value, dict):
        return " ".join(f"{k} {_coerce_str(v)}" for k, v in value.items())
    s = str(value)
    # treat the literal string "None" as empty
    return "" if s.strip().lower() == "none" else s


def _parse_details(details: Any) -> dict:
    """
    'details' in this dataset is stored as a JSON string, e.g.:
        '{"Brand": "Luzianne", "Item Form": "Ground", ...}'
    Parse it safely; return an empty dict on failure.
    """
    if not details:
        return {}
    if isinstance(details, dict):
        return details
    try:
        return json.loads(str(details))
    except (json.JSONDecodeError, TypeError):
        return {}


def _parse_price(price: Any) -> float | None:
    """price can be a float, an int, or the string 'None'."""
    if price is None:
        return None
    try:
        v = float(price)
        return None if v != v else v   # NaN guard
    except (ValueError, TypeError):
        return None


# ── review selection ──────────────────────────────────────────────────────────

def get_top_reviews(
    reviews_dataset_dict,
    parent_asin: str,
    k: int = 5,
) -> list[dict]:
    """
    Select the top-k reviews for a product using get_best_reviews() from
    eda_helpers.py (weighted score: helpful_vote 50%, verified_purchase 30%,
    rating extremity 20%).

    Parameters
    ----------
    reviews_dataset_dict : the full reviews DatasetDict (raw_reviews) —
                           NOT the pre-selected 'full' split, because
                           get_best_reviews() selects 'full' internally.
    parent_asin          : product identifier
    k                    : number of reviews to return

    Returns
    -------
    List of dicts with keys: title, text, rating, helpful_vote
    """
    result = get_best_reviews(reviews_dataset_dict, parent_asin, top_k=k)

    # get_best_reviews returns (total_count, Dataset) when top_k is set,
    # or a bare Dataset with 0 rows when no reviews are found.
    if isinstance(result, tuple):
        _, matched = result
    else:
        matched = result

    if len(matched) == 0:
        return []

    return [
        {
            "title":        row.get("title", "") or "",
            "text":         row.get("text", "") or "",
            "rating":       row.get("rating"),
            "helpful_vote": row.get("helpful_vote", 0),
        }
        for row in matched
    ]


# ── document construction ─────────────────────────────────────────────────────

def format_review(review: dict) -> str:
    """Format a single review the same way as in the notebook."""
    return (
        f"Review (Rating: {review['rating']}): "
        f"{review['title']}. "
        f"{review['text']}\n    "
    )


def build_page_content(product: dict, top_reviews: list[dict]) -> str:
    """
    Build the page_content string that BM25 will index.
    Mirrors the create_document() structure in milestone1_exploration.ipynb.
    """
    title       = _coerce_str(product.get("title"))
    description = " ".join(product.get("description") or [])
    features    = "\n".join(product.get("features") or [])
    categories  = " > ".join(product.get("categories") or [])
    store       = _coerce_str(product.get("store"))
    details     = _parse_details(product.get("details"))
    details_str = " ".join(f"{k}: {v}" for k, v in details.items())

    review_lines = "".join(format_review(r) for r in top_reviews)
    n_reviews    = len(top_reviews)

    return f"""Product: {title}
Category: {categories}
Store: {store}

Features:
    {features}

Description:
    {description}

Details:
    {details_str}

Top Reviews (showing {n_reviews}):
    {review_lines}"""

def build_document(product: dict, top_reviews: list[dict]) -> Document | None:
    """
    Build one LangChain Document for a single product row from the metadata Dataset.
    Returns None if there is no indexable text.
    """
    page_content = build_page_content(product, top_reviews)
    if not page_content.strip():
        return None

    details_dict = _parse_details(product.get("details"))

    metadata = {
        "parent_asin":    product.get("parent_asin", ""),
        "title":          _coerce_str(product.get("title")),
        "main_category":  _coerce_str(product.get("main_category")),
        "price":          _parse_price(product.get("price")),
        "store":          _coerce_str(product.get("store")),
        "categories":     _coerce_str(product.get("categories")),
        "features":       _coerce_str(product.get("features")),
        "description":    _coerce_str(product.get("description")),
        "details":        details_dict,
        "average_rating": product.get("average_rating"),
        "rating_number":  product.get("rating_number"),
        "image":      extract_image(product),
        "top_reviews":    top_reviews,
    }

    return Document(page_content=page_content, metadata=metadata)


def pregroup_reviews(
    reviews_dataset_dict,
    max_reviews_per_product: int = 5,
) -> dict:
    """
    Pre-group top-k reviews per product using DuckDB for efficient scoring
    and ranking — never loads all 14M reviews into Python memory at once.

    Uses a single SQL query with ROW_NUMBER() to rank reviews per product
    by the same weighted score as eda_helpers.get_best_reviews():
        helpful_vote 50% (log-scaled) + verified_purchase 30% + rating extremity 20%
    """
    import duckdb

    print("Pre-grouping reviews via DuckDB (memory-efficient) ...")
    arrow_table = reviews_dataset_dict["full"].data.table

    k = max_reviews_per_product
    query = f"""
        WITH scored AS (
            SELECT
                parent_asin,
                title,
                text,
                rating,
                helpful_vote,
                verified_purchase,
                (
                    0.5 * (LN(1 + GREATEST(COALESCE(helpful_vote, 0), 0)))
                    + 0.3 * (CASE WHEN verified_purchase THEN 1.0 ELSE 0.0 END)
                    + 0.2 * (ABS(COALESCE(rating, 3.0) - 3.0) / 2.0)
                ) AS score
            FROM arrow_table
            WHERE parent_asin IS NOT NULL AND parent_asin != ''
        ),
        ranked AS (
            SELECT *,
                ROW_NUMBER() OVER (
                    PARTITION BY parent_asin
                    ORDER BY score DESC
                ) AS rn
            FROM scored
        )
        SELECT parent_asin, title, text, rating, helpful_vote
        FROM ranked
        WHERE rn <= {k}
        ORDER BY parent_asin, rn
    """

    rows = duckdb.query(query).fetchall()
    cols = ["parent_asin", "title", "text", "rating", "helpful_vote"]

    result = {}
    for row in rows:
        r           = dict(zip(cols, row))
        asin        = r.pop("parent_asin")
        result.setdefault(asin, []).append(r)

    print(f"  {len(result):,} unique parent_asins grouped")
    print("  pre-grouping done")
    return result


def build_documents(
    metadata_dataset: Dataset,
    reviews_dataset_dict,
    max_products: int | None = None,
    max_reviews_per_product: int = 5,
    reviews_lookup: dict | None = None,
) -> list[Document]:
    """
    Build one LangChain Document per product.

    Pass reviews_lookup (from pregroup_reviews) to skip per-product DuckDB
    queries entirely — much faster for large datasets.
    """
    total = len(metadata_dataset)
    n     = min(total, max_products) if max_products else total
    print(f"Building documents for {n} products ...")

    docs = []
    for i in range(n):
        product     = metadata_dataset[i]
        parent_asin = product.get("parent_asin", "")

        if reviews_lookup is not None:
            top_reviews = reviews_lookup.get(parent_asin, [])[:max_reviews_per_product]
        else:
            top_reviews = get_top_reviews(
                reviews_dataset_dict, parent_asin, k=max_reviews_per_product
            )

        doc = build_document(product, top_reviews)
        if doc:
            docs.append(doc)

        if (i + 1) % 500 == 0:
            print(f"  ... {i + 1}/{n} products processed")

    print(f"  -> {len(docs)} documents built (skipped {n - len(docs)} empty)")
    return docs


# ── index build & persist ─────────────────────────────────────────────────────

def build_and_save(
    documents: list[Document],
    index_path: str | Path = "data/processed/bm25_index.pkl",
    corpus_path: str | Path = "data/processed/bm25_corpus.pkl",
) -> BM25Retriever:
    """
    Build a BM25Retriever from documents, then pickle both the
    tokenized corpus and the retriever to disk.

    Parameters
    ----------
    documents   : output of build_documents()
    index_path  : e.g. 'data/processed/bm25_index.pkl'
    corpus_path : e.g. 'data/processed/bm25_corpus.pkl'

    Returns
    -------
    The fitted BM25Retriever instance.
    """
    index_path  = Path(index_path)
    corpus_path = Path(corpus_path)
    index_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"Fitting BM25 index over {len(documents)} documents …")
    retriever = BM25Retriever.from_documents(
        documents,
        preprocess_func=simple_tokenize,
    )

    # Save tokenized corpus separately — useful for inspection in the notebook
    tokenized_corpus = [simple_tokenize(doc.page_content) for doc in documents]
    with open(corpus_path, "wb") as f:
        pickle.dump(tokenized_corpus, f)
    print(f"Tokenized corpus saved → {corpus_path}")

    with open(index_path, "wb") as f:
        pickle.dump(retriever, f)
    print(f"BM25 index saved       → {index_path}")

    return retriever


# ── load ──────────────────────────────────────────────────────────────────────

def load(index_path: str | Path = "data/processed/bm25_index.pkl") -> BM25Retriever:
    """
    Load a previously saved BM25Retriever from disk.
    Call this in app.py instead of rebuilding every time.
    """
    index_path = Path(index_path)
    if not index_path.exists():
        raise FileNotFoundError(
            f"BM25 index not found at '{index_path}'.\n"
            "Run build_and_save() from your notebook first."
        )
    with open(index_path, "rb") as f:
        retriever = pickle.load(f)
    print(f"BM25 index loaded ← {index_path}")
    return retriever


# ── search ────────────────────────────────────────────────────────────────────

def search(
    retriever: BM25Retriever,
    query: str,
    top_k: int = 3,
) -> list[dict]:
    """
    Search the BM25Retriever for a query, returning metadata of top-k results.
    
    Performs a BM25 keyword search on the indexed documents. Tokenizes the query
    using the same tokenizer as the index, computes BM25 scores for all documents,
    and returns structured metadata (including score) for the top-k matches.
    """
    retriever.k = top_k
    
    # Tokenize query the same way the index was built
    tokenized_query = simple_tokenize(query)
    
    # Get raw BM25 scores for ALL documents
    scores = retriever.vectorizer.get_scores(tokenized_query)  # np.ndarray, len = n_docs
    
    # Get indices of top-k scores
    top_indices = sorted(
        range(len(scores)),
        key=scores.__getitem__,
        reverse=True
    )[:top_k]

    # Collect metadata for top results, including score
    return [
        {**retriever.docs[i].metadata, "score": scores[i]}
        for i in top_indices
    ]


# ── notebook entry point ──────────────────────────────────────────────────────

def build_from_hf_datasets(
    metadata_dataset: Dataset,
    reviews_dataset_dict,
    index_path: str | Path = "data/processed/tokenisation/bm25_index.pkl",
    corpus_path: str | Path = "data/processed/tokenisation/bm25_corpus.pkl",
    max_products: int | None = None,
    max_reviews_per_product: int = 5,
) -> BM25Retriever:
    """
    End-to-end helper to call from milestone1_exploration.ipynb.

    Example usage in the notebook:
    --------------------------------
    from src.bm25 import build_from_hf_datasets, load, search

    retriever = build_from_hf_datasets(
        metadata_dataset=raw_metadata['full'],   
        reviews_dataset_dict=raw_reviews,        
        max_products=500,                        
    )

    # Later in app.py — just load the saved index:
    # retriever = load("data/processed/bm25_index.pkl")
    # results   = search(retriever, "something sweet for a cheese board")
    """
    reviews_lookup = pregroup_reviews(reviews_dataset_dict, max_reviews_per_product)
    docs = build_documents(
        metadata_dataset,
        reviews_dataset_dict,
        max_products=max_products,
        max_reviews_per_product=max_reviews_per_product,
        reviews_lookup=reviews_lookup,
    )
    return build_and_save(docs, index_path=index_path, corpus_path=corpus_path)


def build_from_hf_datasets_batched(
    metadata_dataset: Dataset,
    reviews_dataset_dict,
    index_path: str | Path = "data/processed/tokenisation/bm25_index.pkl",
    corpus_path: str | Path = "data/processed/tokenisation/bm25_corpus.pkl",
    batch_size: int = 2000,
    max_reviews_per_product: int = 5,
    max_products: int | None = None,
) -> BM25Retriever:
    """
    Memory-safe version of build_from_hf_datasets — builds documents in
    batches to avoid OOM kernel crashes on large datasets.

    Checkpoints completed batches to data/processed/checkpoints/ after each
    batch, so if the kernel dies mid-run you can resume from the last
    completed batch instead of starting over.

    Example usage in the notebook:
    --------------------------------
    retriever = build_from_hf_datasets_batched(
        metadata_dataset=raw_metadata['full'],
        reviews_dataset_dict=raw_reviews,
        batch_size=5000,
        max_reviews_per_product=3,
        max_products=60000,   # None = use all
    )
    """
    index_path  = Path(index_path)
    corpus_path = Path(corpus_path)

    # checkpoint folder lives next to the index
    checkpoint_dir = index_path.parent / "checkpoints"
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    total = min(len(metadata_dataset), max_products) if max_products else len(metadata_dataset)

    # find resume point — checkpoints named docs_0.pkl, docs_2000.pkl, ...
    existing = sorted(checkpoint_dir.glob("docs_*.pkl"))
    if existing:
        last_ckpt    = existing[-1]
        resume_start = int(last_ckpt.stem.split("_")[1]) + batch_size
        print(f"Resuming from product {resume_start} "
              f"({len(existing)} checkpoint(s) found)")
        all_docs = []
        for ckpt in existing:
            with open(ckpt, "rb") as f:
                all_docs.extend(pickle.load(f))
        print(f"  loaded {len(all_docs)} docs from checkpoints")
    else:
        resume_start = 0
        all_docs     = []
        print(f"Starting fresh — {total} products to process")

    # pre-group all reviews once
    reviews_lookup = pregroup_reviews(reviews_dataset_dict, max_reviews_per_product)

    # batch loop
    for start in range(resume_start, total, batch_size):
        end   = min(start + batch_size, total)
        print(f"\nBatch {start}-{end} of {total} ...")

        batch      = metadata_dataset.select(range(start, end))
        batch_docs = build_documents(
            batch,
            reviews_dataset_dict,
            max_products=None,
            max_reviews_per_product=max_reviews_per_product,
            reviews_lookup=reviews_lookup,
        )
        all_docs.extend(batch_docs)

        # save checkpoint for this batch
        ckpt_path = checkpoint_dir / f"docs_{start}.pkl"
        with open(ckpt_path, "wb") as f:
            pickle.dump(batch_docs, f)
        print(f"  checkpoint saved -> {ckpt_path.name}")
        print(f"  cumulative docs  : {len(all_docs)}")

    # build final index
    print(f"\nAll batches done - {len(all_docs)} total documents.")
    retriever = build_and_save(all_docs, index_path=index_path, corpus_path=corpus_path)

    # clean up checkpoints now that final index is safely written
    for ckpt in checkpoint_dir.glob("docs_*.pkl"):
        ckpt.unlink()
    print("Checkpoints cleaned up.")

    return retriever