"""Cross-encoder reranking + sentence-window context expansion, as used by
Notebook 5.

`Reranker.rerank` re-scores a candidate list with a cross-encoder and keeps
the top `top_n`. `relevance_score` sigmoid-normalizes those scores into a
0-1 "context relevance" proxy. `sentence_window_expand` widens a chunk to
include `window` sentences of surrounding context from its source document.
"""

import re

import numpy as np
import pandas as pd
from sentence_transformers import CrossEncoder

_SENTENCE_SPLIT_RE = re.compile(r"(?<=[.!?])\s+")


def sigmoid(x: float) -> float:
    return 1 / (1 + np.exp(-x))


def split_sentences(text: str) -> list:
    return [s.strip() for s in _SENTENCE_SPLIT_RE.split(text) if s.strip()]


def build_doc_sentences(docs_df: pd.DataFrame) -> dict:
    """Maps doc_id -> list of sentences, used by `sentence_window_expand`."""
    return {row["doc_id"]: split_sentences(row["text"]) for _, row in docs_df.iterrows()}


def relevance_score(reranked: list) -> float:
    """Mean sigmoid-normalized cross-encoder relevance of the reranked chunks."""
    if not reranked:
        return 0.0
    return float(np.mean([sigmoid(s) for _, s in reranked]))


class Reranker:
    def __init__(self, model_name: str):
        self.model_name = model_name
        self.model = CrossEncoder(model_name)

    def rerank(self, query: str, candidate_idxs: list, chunks_df: pd.DataFrame, top_n: int) -> list:
        if not candidate_idxs:
            return []
        pairs = [(query, chunks_df.iloc[idx]["text"]) for idx, _ in candidate_idxs]
        scores = self.model.predict(pairs)
        order = np.argsort(scores)[::-1][:top_n]
        return [(candidate_idxs[i][0], float(scores[i])) for i in order]


def _locate_in_sentences(sentences: list, target: str, from_end: bool = False) -> int | None:
    """Find index of target in sentences list.

    Tries exact match first. If that fails, uses a 40-char anchor from the
    start (or end when from_end=True) of target to handle the common case
    where recursive chunking cuts a sentence in half — the chunk boundary
    lands mid-sentence so target is a fragment not present verbatim in the
    document sentence list.
    """
    for i, s in enumerate(sentences):
        if s == target:
            return i
    anchor = target[-40:].strip() if from_end else target[:40].strip()
    if len(anchor) < 10:
        return None
    for i, s in enumerate(sentences):
        if anchor in s or s in target:
            return i
    return None


def sentence_window_expand(chunk_idx: int, chunks_df: pd.DataFrame, doc_sentences: dict, window: int = 1):
    """Returns (expanded_text, (start, end)) -- range is None if the chunk's
    sentences couldn't be located in the source document (falls back to the
    raw chunk text)."""
    chunk = chunks_df.iloc[chunk_idx]
    sentences = doc_sentences[chunk["doc_id"]]
    chunk_sentences = split_sentences(chunk["text"])
    if not chunk_sentences:
        return chunk["text"], None

    first_idx = _locate_in_sentences(sentences, chunk_sentences[0], from_end=False)
    last_idx  = _locate_in_sentences(sentences, chunk_sentences[-1], from_end=True)
    if first_idx is None or last_idx is None:
        return chunk["text"], None

    start = max(0, first_idx - window)
    end = min(len(sentences), last_idx + 1 + window)
    expanded = " ".join(sentences[start:end])
    return expanded, (start, end)