import spaces
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_NAME = "cross-encoder/nli-deberta-v3-small"
OPTION_LETTERS = ["A", "B", "C", "D", "E"]
MAX_LENGTH = 256

# The model is loaded on the first prediction rather than at import time, so the
# Space finishes booting immediately and a slow download can never crash startup.
_MODEL_CACHE = {}


def load_model():
    """Load tokenizer and model once and keep them in memory."""
    if "model" not in _MODEL_CACHE:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
        model.eval()

        # NLI checkpoints do not agree on label ordering, so the entailment
        # column is read from the config instead of being hard-coded. A wrong
        # index would silently invert the ranking.
        entailment_index = None
        for label, index in model.config.label2id.items():
            if label.lower().startswith("entail"):
                entailment_index = index
        if entailment_index is None:
            entailment_index = model.config.num_labels - 1

        _MODEL_CACHE["tokenizer"] = tokenizer
        _MODEL_CACHE["model"] = model
        _MODEL_CACHE["entailment_index"] = entailment_index

    return (
        _MODEL_CACHE["tokenizer"],
        _MODEL_CACHE["model"],
        _MODEL_CACHE["entailment_index"],
    )


@spaces.GPU(duration=30)
def score_options(prompt, options):
    """Return one entailment probability per option, in A-E order."""
    tokenizer, model, entailment_index = load_model()
    model.to("cuda")

    premises = [prompt] * len(options)
    encoded = tokenizer(
        premises,
        options,
        truncation=True,
        max_length=MAX_LENGTH,
        padding=True,
        return_tensors="pt",
    ).to("cuda")

    with torch.no_grad():
        logits = model(**encoded).logits


def rank_options(option_scores):
    """Return option letters sorted from highest score to lowest."""
    order = sorted(range(len(option_scores)), key=lambda i: option_scores[i], reverse=True)
    return [OPTION_LETTERS[i] for i in order]


def format_results(prompt, options, option_scores, ranked_letters):
    """Build the markdown result table shown to the user."""
    score_by_letter = dict(zip(OPTION_LETTERS, option_scores))
    text_by_letter = dict(zip(OPTION_LETTERS, options))

    lines = ["| Rank | Option | Answer text | Entailment score |", "|---|---|---|---|"]
    for position, letter in enumerate(ranked_letters[:3], start=1):
        answer_text = text_by_letter[letter]
        if len(answer_text) > 90:
            answer_text = answer_text[:90] + "..."
        lines.append(
            f"| {position} | **{letter}** | {answer_text} | {score_by_letter[letter]:.4f} |"
        )

    remaining = ", ".join(
        f"{letter} ({score_by_letter[letter]:.4f})" for letter in ranked_letters[3:]
    )
    lines.append("")
    lines.append(f"Not in the top three: {remaining}")
    return "\n".join(lines)


def predict(prompt, option_a, option_b, option_c, option_d, option_e):
    """Validate the inputs, score the options and return the top three."""
    prompt = (prompt or "").strip()
    options = [(text or "").strip() for text in
               [option_a, option_b, option_c, option_d, option_e]]

    if not prompt:
        return "Please enter a question before predicting.", ""

    missing = [OPTION_LETTERS[i] for i, text in enumerate(options) if not text]
    if missing:
        return f"Please fill in every option. Still empty: {', '.join(missing)}.", ""

    try:
        option_scores = score_options(prompt, options)
    except Exception as error:  # keeps the app alive if the model fails to load
        return f"Something went wrong while running the model: {error}", ""

    ranked_letters = rank_options(option_scores)
    results_table = format_results(prompt, options, option_scores, ranked_letters)
    submission_string = " ".join(ranked_letters[:3])
    return results_table, submission_string


EXAMPLE_ONE = [
    "Which property of a material determines how much it resists the flow of "
    "electric current?",
    "Resistivity, an intrinsic property that depends on the material itself",
    "Mass, since heavier materials always conduct electricity less well",
    "Volume, because a larger object always has a larger current",
    "Colour, as darker materials absorb more electrical energy",
    "Density, which fixes the number of free electrons in every metal",
]

EXAMPLE_TWO = [
    "What is the main reason a catalyst increases the rate of a chemical reaction?",
    "It provides an alternative pathway with a lower activation energy",
    "It increases the total energy released by the reaction",
    "It shifts the position of equilibrium towards the products",
    "It raises the temperature of the reaction mixture directly",
    "It permanently bonds to the reactants and is consumed by them",
]

DESCRIPTION = """
# Smart MCQ Solver

Type a multiple-choice question and its five options. The model ranks all five and
returns the **top three in order**, which is the submission format used by the
Kaggle *Smart MCQ Solver Challenge* (metric: MAP@3).

**What is running here:** `cross-encoder/nli-deberta-v3-small` used zero-shot.
The question is the premise, each option is a hypothesis, and options are ranked
by entailment probability. This is Model 2 from the project. It is not the
fine-tuned eight-member DeBERTa-v3 ensemble that produced the 0.736 leaderboard
score, so expect weaker answers than that number suggests.

Devang Khirsaria · 23f1001439 · BS in Data Science and Applications, IIT Madras
"""

with gr.Blocks() as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column():
            prompt_box = gr.Textbox(
                label="Question",
                placeholder="Type the question here",
                lines=3,
                value=EXAMPLE_ONE[0],
            )
            option_boxes = []
            for index, letter in enumerate(OPTION_LETTERS):
                option_boxes.append(
                    gr.Textbox(
                        label=f"Option {letter}",
                        placeholder=f"Text of option {letter}",
                        lines=2,
                        value=EXAMPLE_ONE[index + 1],
                    )
                )

            with gr.Row():
                predict_button = gr.Button("Predict top 3", variant="primary")
                clear_button = gr.Button("Clear")

            with gr.Row():
                example_one_button = gr.Button("Load physics example")
                example_two_button = gr.Button("Load chemistry example")

        with gr.Column():
            results_output = gr.Markdown(
                label="Ranked predictions",
                value="Press **Predict top 3** to run the model. "
                "The first run downloads the model and takes about a minute; "
                "later runs take a couple of seconds.",
            )
            submission_output = gr.Textbox(
                label="Kaggle submission format",
                info="Exactly the three letters you would write in the Prediction column.",
            )

    all_inputs = [prompt_box] + option_boxes

    predict_button.click(
        fn=predict,
        inputs=all_inputs,
        outputs=[results_output, submission_output],
    )
    clear_button.click(
        fn=lambda: ["", "", "", "", "", ""],
        inputs=None,
        outputs=all_inputs,
    )
    example_one_button.click(
        fn=lambda: EXAMPLE_ONE,
        inputs=None,
        outputs=all_inputs,
    )
    example_two_button.click(
        fn=lambda: EXAMPLE_TWO,
        inputs=None,
        outputs=all_inputs,
    )

if __name__ == "__main__":
    demo.launch()