import math
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# ---- Model config ----
MODEL_NAME = "gpt2"           # e.g., "distilgpt2", "gpt2", "gpt2-medium"
device = "cuda" if torch.cuda.is_available() else "cpu"

tok = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)
model.eval()

def seq_logprob(context: str, candidate: str, assume_leading_space: bool, show_topk: int):
    """
    Return (total_logprob, step_detail_text, token_list)
    Computes P(candidate | context) via chain rule over the candidate tokens.
    """
    if not context.strip():
        return None, "Please enter context.", []

    # Helpful for GPT-2 BPE “word starts”
    cand_text = (" " + candidate) if assume_leading_space else candidate

    with torch.no_grad():
        ctx_ids = tok.encode(context, return_tensors="pt").to(device)
        cand_ids = tok.encode(cand_text, add_special_tokens=False)

        if len(cand_ids) == 0:
            return None, "Candidate tokenized to empty sequence (check spacing).", []

        total_logprob = 0.0
        step_lines = []
        input_ids = ctx_ids
        token_texts = []

        for i, t_id in enumerate(cand_ids):
            outputs = model(input_ids=input_ids)
            logits = outputs.logits[:, -1, :]
            logprobs = torch.log_softmax(logits, dim=-1)
            token_logprob = logprobs[0, t_id].item()
            total_logprob += token_logprob

            tok_str = tok.decode([t_id])
            token_texts.append(tok_str)

            if show_topk > 0:
                topk_vals, topk_idx = torch.topk(logprobs, k=min(show_topk, logprobs.shape[-1]), dim=-1)
                tops = ", ".join([f"{repr(tok.decode([int(idx)]))}:{val.item():.2f}"
                                  for idx, val in zip(topk_idx[0], topk_vals[0])])
                step_lines.append(
                    f"Step {i+1}: token={repr(tok_str)}  logprob={token_logprob:.6f}  "
                    f"prob={math.exp(token_logprob):.6e}\n  top-{show_topk}: {tops}"
                )
            else:
                step_lines.append(
                    f"Step {i+1}: token={repr(tok_str)}  logprob={token_logprob:.6f}  "
                    f"prob={math.exp(token_logprob):.6e}"
                )

            # teacher-forcing: append the true token to continue conditioning
            input_ids = torch.cat([input_ids, torch.tensor([[t_id]], device=device)], dim=1)

        detail_text = "\n".join(step_lines)
        return total_logprob, detail_text, token_texts

def compare_candidates(context, cand1, cand2, assume_space, topk):
    # Basic input checks
    errs = []
    if not context.strip():
        errs.append("Please enter a context.")
    if not cand1.strip():
        errs.append("Please enter Candidate A.")
    if not cand2.strip():
        errs.append("Please enter Candidate B.")
    if errs:
        return (
            f"<div style='color:#b00020;font-weight:600'>{' '.join(errs)}</div>",
            "", "", "", "", ""
        )

    # Compute log-probs
    logp1, details1, toks1 = seq_logprob(context, cand1, assume_space, topk)
    logp2, details2, toks2 = seq_logprob(context, cand2, assume_space, topk)

    if logp1 is None or logp2 is None:
        return (
            "<div style='color:#b00020;font-weight:600'>Tokenization error. Check inputs.</div>",
            details1, details2, "", "", ""
        )

    # Summaries for each candidate
    def make_summary(label, cand, logp, toks):
        seq_prob = math.exp(logp)
        return (
            f"**{label}**: {cand}\n\n"
            f"Tokenization: {toks}\n"
            f"Total logprob: {logp:.6f}\n"
            f"Sequence probability: {seq_prob:.6e}"
        )

    summary1 = make_summary("Candidate A", cand1, logp1, toks1)
    summary2 = make_summary("Candidate B", cand2, logp2, toks2)

    # Ratio, odds, and winner
    log_odds = logp1 - logp2                # log(P(A)/P(B))
    # Cap extreme ratios for display; still show exact log-odds
    try:
        ratio = math.exp(log_odds)
        ratio_str = f"{ratio:.6e}"
    except OverflowError:
        ratio_str = "∞ (overflow)"
    winner = "Candidate A" if logp1 > logp2 else ("Tie" if abs(log_odds) < 1e-12 else "Candidate B")

    if winner == "Candidate A":
        win_color = "#166534"  # green
    elif winner == "Candidate B":
        win_color = "#1d4ed8"  # blue
    else:
        win_color = "#92400e"  # amber (tie)

    headline = (
        f"<div style='padding:14px;border-radius:12px;background:#f8fafc;"
        f"border:1px solid #e2e8f0;margin-bottom:10px'>"
        f"<div style='font-size:20px;font-weight:800;color:{win_color};'>Winner: {winner}</div>"
        f"<div style='margin-top:6px;font-size:16px;'>"
        f"Odds (A/B) = <b>{ratio_str}</b> &nbsp;|&nbsp; "
        f"log-odds = <b>{log_odds:.6f}</b>"
        f"</div>"
        f"<div style='margin-top:6px;color:#475569'>"
        f"(Odds &gt; 1 means A is more probable; &lt; 1 means B is more probable.)"
        f"</div></div>"
    )

    return headline, summary1, details1, summary2, details2, ""

with gr.Blocks(title="Two-Candidate Next-Token Probability Comparator") as demo:
    gr.Markdown(
        "# Two-Candidate Next-Word/Token Probability\n"
        "Given a **context**, compare the conditional probabilities of **two candidate continuations**.\n"
        "- Uses a pretrained causal LM (default: GPT-2). No fine-tuning.\n"
        "- Works at the **token** level; multi-token “words” are handled via the chain rule.\n"
        "- The **Winner** is the higher-probability candidate; we also show the **odds ratio (A/B)** and log-odds."
    )

    with gr.Row():
        context = gr.Textbox(label="Context (prompt)", lines=6, placeholder="Paste your prior text here...")

    with gr.Row():
        cand1 = gr.Textbox(label="Candidate A (follow-up)")
        cand2 = gr.Textbox(label="Candidate B (follow-up)")

    with gr.Row():
        assume_space = gr.Checkbox(
            value=True,
            label="Assume leading space before candidates (helps align with word starts in GPT-2 tokenization)"
        )
        topk = gr.Slider(0, 20, value=5, step=1, label="Show top-k alternatives (per token step)")

    btn = gr.Button("Compare")
    winner_html = gr.HTML()
    summary1 = gr.Markdown()
    details1 = gr.Textbox(label="Candidate A — step-by-step", lines=10)
    summary2 = gr.Markdown()
    details2 = gr.Textbox(label="Candidate B — step-by-step", lines=10)
    _hidden = gr.Textbox(visible=False)

    btn.click(
        fn=compare_candidates,
        inputs=[context, cand1, cand2, assume_space, topk],
        outputs=[winner_html, summary1, details1, summary2, details2, _hidden]
    )

demo.launch()