modules/visualizations/explanations.py

"""Mathematical formula and explanation utilities."""

import streamlit as st


def show_bleu_formula():
    """Display BLEU score formula with explanation."""
    st.latex(r"""
    \text{BLEU} = \text{BP} \times \exp\left(\sum_{n=1}^{N} w_n \log p_n\right)
    """)
    
    st.markdown("""
    **Where:**
    - **BP** = Brevity Penalty (penalizes short candidates)
    - **p_n** = Modified precision for n-grams
    - **w_n** = Weight for n-gram precision (usually uniform)
    - **N** = Maximum n-gram size (typically 4)
    """)
    
    with st.expander("Brevity Penalty Formula"):
        st.latex(r"""
        \text{BP} = \begin{cases}
        1 & \text{if } c > r \\
        e^{(1-r/c)} & \text{if } c \leq r
        \end{cases}
        """)
        st.markdown("- **c** = candidate length, **r** = reference length")
    
    with st.expander("Modified N-gram Precision"):
        st.latex(r"""
        p_n = \frac{\sum_{\text{n-gram} \in \text{Cand}} \text{Count}_{\text{clip}}(\text{n-gram})}{\sum_{\text{n-gram} \in \text{Cand}} \text{Count}(\text{n-gram})}
        """)
        st.markdown("Count_clip = min(candidate count, reference count)")


def show_rogue_formula():
    """Display ROGUE score formula with explanation."""
    st.subheader("ROGUE-N (N-gram Based)")
    st.latex(r"""
    \text{ROGUE-N} = \frac{\sum_{S \in \{\text{References}\}} \sum_{\text{gram}_n \in S} \text{Count}_{\text{match}}(\text{gram}_n)}{\sum_{S \in \{\text{References}\}} \sum_{\text{gram}_n \in S} \text{Count}(\text{gram}_n)}
    """)
    st.markdown("**Focus: Recall** — how many reference n-grams were captured")
    
    st.subheader("ROGUE-L (Longest Common Subsequence)")
    st.latex(r"""
    \text{R}_{\text{lcs}} = \frac{LCS(X, Y)}{|X|}, \quad \text{P}_{\text{lcs}} = \frac{LCS(X, Y)}{|Y|}
    """)
    st.latex(r"""
    \text{F}_{\text{lcs}} = \frac{(1 + \beta^2) \text{R}_{\text{lcs}} \text{P}_{\text{lcs}}}{\text{R}_{\text{lcs}} + \beta^2 \text{P}_{\text{lcs}}}
    """)
    st.markdown("**LCS** finds longest sequence appearing in both (not necessarily consecutive)")


def show_perplexity_formula():
    """Display Perplexity formula with explanation."""
    st.latex(r"""
    \text{Perplexity} = \exp\left(-\frac{1}{N} \sum_{i=1}^{N} \log P(w_i | w_1 \ldots w_{i-1})\right)
    """)
    
    st.markdown("""
    **Interpretation:**
    - Perplexity = $2^{H}$ where $H$ is the cross-entropy
    - Can be thought of as "weighted average branching factor"
    - Lower is better (model is less "confused")
    """)
    
    with st.expander("Example Interpretation"):
        st.markdown("""
        | Perplexity | Meaning |
        |------------|---------|
        | 1 | Perfect prediction |
        | 10 | Model has ~10 choices at each step |
        | 100 | Model is very confused |
        | 1000 | Near random guessing |
        """)


def show_mrr_formula():
    """Display MRR formula with explanation."""
    st.latex(r"""
    \text{MRR} = \frac{1}{|Q|} \sum_{i=1}^{|Q|} \frac{1}{\text{rank}_i}
    """)
    
    st.markdown("""
    **Where:**
    - **|Q|** = Number of queries
    - **rank_i** = Position of correct answer for query i
    - If correct answer not in list: $\frac{1}{\text{rank}} = 0$
    """)
    
    with st.expander("Reciprocal Rank Examples"):
        st.markdown("""
        | Position | Reciprocal Rank |
        |------------|-----------------|
        | 1st | 1.0 |
        | 2nd | 0.5 |
        | 3rd | 0.33 |
        | 4th | 0.25 |
        | Not found | 0 |
        """)


def show_bert_score_formula():
    """Display BERT Score formula with explanation."""
    st.markdown("""
    BERT Score uses contextual embeddings from pre-trained BERT to compute similarity:
    """)
    
    st.latex(r"""
    \text{Similarity}(x_i, y_j) = \frac{\mathbf{x}_i^T \mathbf{y}_j}{||\mathbf{x}_i|| ||\mathbf{y}_j||}
    """)
    
    st.subheader("Greedy Matching for Precision/Recall")
    st.latex(r"""
    \text{P}_{\text{BERT}} = \frac{1}{|x|} \sum_{x_i \in x} \max_{y_j \in y} \mathbf{x}_i^T \mathbf{y}_j
    """)
    st.latex(r"""
    \text{R}_{\text{BERT}} = \frac{1}{|y|} \sum_{y_j \in y} \max_{x_i \in x} \mathbf{x}_i^T \mathbf{y}_j
    """)
    
    st.markdown("""
    **Key Idea:** Each token in one text is matched to the most similar token in the other text.
    - **Precision:** Average of best matches from candidate to reference
    - **Recall:** Average of best matches from reference to candidate
    """)


def show_metric_comparison_table():
    """Display comparison table of all metrics."""
    st.markdown("""
    | Metric | Type | Needs Reference | Best For | Range |
    |--------|------|-----------------|----------|-------|
    | **BLEU** | Lexical (Precision) | Yes | Machine Translation | 0-1 |
    | **ROGUE** | Lexical (Recall) | Yes | Summarization | 0-1 |
    | **Perplexity** | Model Confidence | No | Language Modeling | 1-∞ |
    | **MRR** | Ranking | Yes (answer) | QA/Retrieval | 0-1 |
    | **BERT Score** | Semantic | Yes | Paraphrase Detection | 0-1 |
    """)


def create_formula_expander(metric_name: str):
    """Create an expander with formula for a given metric.
    
    Args:
        metric_name: Name of the metric ('bleu', 'rogue', etc.)
    """
    with st.expander(f"📐 {metric_name.upper()} Formula & Explanation"):
        if metric_name.lower() == 'bleu':
            show_bleu_formula()
        elif metric_name.lower() == 'rogue':
            show_rogue_formula()
        elif metric_name.lower() == 'perplexity':
            show_perplexity_formula()
        elif metric_name.lower() == 'mrr':
            show_mrr_formula()
        elif metric_name.lower() == 'bert_score':
            show_bert_score_formula()
        else:
            st.write("Formula not available for this metric.")