"""BM25 Score Explanation for RBSP.

Provides detailed, human-readable breakdowns of BM25 scoring decisions,
similar to Elasticsearch's explain API.  Useful for debugging search
relevance and understanding why a document ranked where it did.

The explanation is a recursive tree:  a top-level sum-of-terms explanation
contains per-term weight explanations, each of which breaks down into
IDF and TF-normalization components.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Dict, List


@dataclass
class Explanation:
    """A node in a score explanation tree.

    Attributes:
        score:       The numeric score at this level.
        description: Human-readable description of what this score represents.
        details:     Nested sub-explanations forming a recursive tree.
    """

    score: float
    description: str
    details: list[Explanation] = field(default_factory=list)

    def __str__(self) -> str:
        return format_explanation(self)


def format_explanation(explanation: Explanation, indent: int = 0) -> str:
    """Format an explanation tree as a readable indented string.

    Args:
        explanation: The root explanation node.
        indent:      Current indentation level (number of spaces).

    Returns:
        Multi-line string with indented score breakdown.
    """
    prefix = " " * indent
    lines = [f"{prefix}{explanation.score:.2f} = {explanation.description}"]
    for detail in explanation.details:
        lines.append(format_explanation(detail, indent + 2))
    return "\n".join(lines)


class ExplainableBM25:
    """BM25 scorer that produces detailed score explanations.

    Mirrors the BM25 formula from ``rbsp.core.bm25.BM25Scorer`` but
    returns an ``Explanation`` tree instead of a bare float.

    Args:
        k1: Term frequency saturation parameter (default 1.2).
        b:  Document length normalization parameter (default 0.75).
    """

    def __init__(self, k1: float = 1.2, b: float = 0.75) -> None:
        self.k1 = k1
        self.b = b

    def explain(
        self,
        query_terms: List[str],
        doc_term_freqs: Dict[str, int],
        doc_len: int,
        num_docs: int,
        avg_doc_len: float,
        doc_freqs: Dict[str, int],
    ) -> Explanation:
        """Produce a full BM25 score explanation.

        Args:
            query_terms:    List of query terms.
            doc_term_freqs: {term: count} for the candidate document.
            doc_len:        Total tokens in the document.
            num_docs:       Total documents in the corpus.
            avg_doc_len:    Average document length across the corpus.
            doc_freqs:      {term: number_of_documents_containing_term}.

        Returns:
            An ``Explanation`` tree with the total BM25 score at the root
            and per-term breakdowns as children.
        """
        term_explanations: List[Explanation] = []
        total_score = 0.0

        for term in query_terms:
            tf = doc_term_freqs.get(term, 0)
            df = doc_freqs.get(term, 0)

            # IDF: ln((N - df + 0.5) / (df + 0.5) + 1)
            idf_val = math.log((num_docs - df + 0.5) / (df + 0.5) + 1.0)
            idf_explanation = Explanation(
                score=idf_val,
                description=(
                    f"idf(docFreq={df}, numDocs={num_docs})"
                    f" = ln(({num_docs}-{df}+0.5)/({df}+0.5)+1)"
                ),
            )

            if tf <= 0:
                # Term not present in document
                term_explanations.append(Explanation(
                    score=0.0,
                    description=f"weight({term}) [BM25]: tf=0, no contribution",
                    details=[idf_explanation],
                ))
                continue

            # TF normalization
            numerator = tf * (self.k1 + 1.0)
            denominator = tf + self.k1 * (
                1.0 - self.b + self.b * doc_len / avg_doc_len
            )
            tf_norm = numerator / denominator

            tf_explanation = Explanation(
                score=tf_norm,
                description=(
                    f"tf_norm(tf={tf}, dl={doc_len}, avgdl={avg_doc_len:.0f})"
                    f" = ({tf}*({self.k1}+1))"
                    f"/({tf}+{self.k1}*(1-{self.b}+{self.b}*{doc_len}/{avg_doc_len:.0f}))"
                ),
            )

            term_score = idf_val * tf_norm
            total_score += term_score

            term_explanations.append(Explanation(
                score=term_score,
                description=f"weight({term}) [BM25]:",
                details=[idf_explanation, tf_explanation],
            ))

        return Explanation(
            score=total_score,
            description="sum of:",
            details=term_explanations,
        )
