import os
import re

import torch
import gradio as gr
import spaces
import wandb
import faiss
import pandas as pd

from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForMultipleChoice


# ============================================================
# CONFIG
# ============================================================

ENTITY = "24f3000211-dl-genai-project"
PROJECT = "24f3000211-t22026"

MODEL_ARTIFACT = (
    "24f3000211-dl-genai-project/"
    "24f3000211-t22026/"
    "deberta-small-rag-2:v0"
)

BASE_MODEL = "microsoft/deberta-v3-small"

MODEL_FILE = "deberta-small-rag-2.pth"

EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

INDEX_FILE = "rag_index.faiss"
RAG_DATA_FILE = "rag_train.pkl"

# Final RAG model configuration
MAX_LENGTH = 368

# Number of examples actually inserted into the RAG prompt
TOP_K = 3

# Number of FAISS candidates searched before filtering
RETRIEVAL_TOP_N = 20

# Minimum similarity required
SIMILARITY_THRESHOLD = 0.55


# ============================================================
# DEVICE
# ============================================================

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

print("Initial device:", DEVICE)


# ============================================================
# DOWNLOAD TRAINED DEBERTA CHECKPOINT FROM W&B
# ============================================================

print("Downloading DeBERTa model from W&B...")

wandb.login(
    key=os.environ["WANDB_API_KEY"]
)

run = wandb.init(
    entity=ENTITY,
    project=PROJECT,
    job_type="huggingface-inference"
)

artifact = run.use_artifact(
    MODEL_ARTIFACT,
    type="model"
)

artifact_dir = artifact.download()

run.finish()

checkpoint_path = os.path.join(
    artifact_dir,
    MODEL_FILE
)

print("Checkpoint:", checkpoint_path)


# ============================================================
# LOAD TOKENIZER
# ============================================================

print("Loading tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(
    BASE_MODEL
)


# ============================================================
# CREATE DEBERTA MULTIPLE-CHOICE MODEL
# ============================================================

print("Creating DeBERTa model...")

model = AutoModelForMultipleChoice.from_pretrained(
    BASE_MODEL
)


# ============================================================
# LOAD TRAINED WEIGHTS
# ============================================================

print("Loading trained weights...")

checkpoint = torch.load(
    checkpoint_path,
    map_location="cpu"
)


# ------------------------------------------------------------
# Support different checkpoint formats
# ------------------------------------------------------------

if isinstance(checkpoint, dict):

    if "model_state_dict" in checkpoint:

        state_dict = checkpoint["model_state_dict"]

    elif "state_dict" in checkpoint:

        state_dict = checkpoint["state_dict"]

    else:

        state_dict = checkpoint

else:

    state_dict = checkpoint.state_dict()


# ------------------------------------------------------------
# Remove DataParallel prefix if present
# ------------------------------------------------------------

state_dict = {
    (
        key.replace("module.", "", 1)
        if key.startswith("module.")
        else key
    ): value

    for key, value in state_dict.items()
}


# ------------------------------------------------------------
# Load checkpoint
# ------------------------------------------------------------

model.load_state_dict(
    state_dict,
    strict=True
)


# ------------------------------------------------------------
# Verify important trained head
# ------------------------------------------------------------

required_keys = [
    "classifier.weight",
    "classifier.bias",
    "pooler.dense.weight",
    "pooler.dense.bias"
]

missing_required = [
    key
    for key in required_keys
    if key not in state_dict
]

if missing_required:

    raise RuntimeError(
        "Required trained weights are missing: "
        + str(missing_required)
    )


print(
    "✓ Trained DeBERTa classification head "
    "loaded successfully."
)


# ============================================================
# MOVE MODEL TO DEVICE
# ============================================================

model.to(DEVICE)

model.eval()

print("✓ DeBERTa model loaded successfully!")


# ============================================================
# LOAD FAISS INDEX
# ============================================================

print("Loading FAISS index...")

index = faiss.read_index(
    INDEX_FILE
)

print(
    "✓ FAISS index loaded:",
    index.ntotal,
    "vectors"
)


# ============================================================
# LOAD RAG TRAINING DATA
# ============================================================

print("Loading RAG dataframe...")

rag_train = pd.read_pickle(
    RAG_DATA_FILE
)

print(
    "✓ RAG dataframe:",
    rag_train.shape
)


# ============================================================
# LOAD SENTENCE TRANSFORMER
# ============================================================

print("Loading Sentence Transformer...")

embedder = SentenceTransformer(
    EMBEDDING_MODEL,
    device="cpu"
)

print("✓ Embedding model loaded!")


# ============================================================
# QUESTION NORMALIZATION
# ============================================================

QUESTION_WORDS = (
    "what",
    "which",
    "who",
    "when",
    "where",
    "why",
    "how",
    "whose",
    "whom"
)


def canonical_question(text):

    text = str(text).lower().strip()

    text = re.sub(
        r"\s+",
        " ",
        text
    )

    pattern = (
        r"\b("
        + "|".join(QUESTION_WORDS)
        + r")\b.*"
    )

    match = re.search(
        pattern,
        text
    )

    if match:
        text = match.group(0)

    text = re.sub(
        r"\s*(carefully|"
        r"based on the given context|"
        r"based on the given information|"
        r"among the listed options|"
        r"from the following choices)"
        r"\.?\s*$",
        "",
        text,
        flags=re.IGNORECASE
    )

    text = re.sub(
        r"\s+",
        " ",
        text
    ).strip()

    return text


# ============================================================
# BUILD RETRIEVAL QUERY
# ============================================================

def build_retrieval_query(
    question,
    choices
):

    return f"""
Question:
{canonical_question(question)}

A. {choices["A"]}
B. {choices["B"]}
C. {choices["C"]}
D. {choices["D"]}
E. {choices["E"]}
""".strip()


# ============================================================
# RETRIEVE SIMILAR QUESTIONS
# ============================================================

def retrieve_examples(
    question,
    choices,
    top_k=TOP_K,
    top_n=RETRIEVAL_TOP_N
):

    query = build_retrieval_query(
        question,
        choices
    )

    # --------------------------------------------------------
    # Generate query embedding
    # --------------------------------------------------------

    query_embedding = embedder.encode(
        query,
        convert_to_numpy=True,
        normalize_embeddings=True
    )

    query_embedding = query_embedding.astype(
        "float32"
    )

    # --------------------------------------------------------
    # FAISS search
    # --------------------------------------------------------

    scores, indices = index.search(
        query_embedding.reshape(1, -1),
        top_n
    )

    query_topic = canonical_question(
        question
    )

    retrieved = []

    seen = set()

    # --------------------------------------------------------
    # Process retrieved candidates
    # --------------------------------------------------------

    for score, idx in zip(
        scores[0],
        indices[0]
    ):

        if idx < 0:
            continue

        if float(score) < SIMILARITY_THRESHOLD:
            continue

        row = rag_train.iloc[int(idx)]

        # ----------------------------------------------------
        # Get canonical question
        # ----------------------------------------------------

        if "canonical_question" in row.index:

            topic = str(
                row["canonical_question"]
            )

        else:

            topic = canonical_question(
                row["prompt"]
            )

        # ----------------------------------------------------
        # Avoid exact same question
        # ----------------------------------------------------

        if topic == query_topic:
            continue

        # ----------------------------------------------------
        # Avoid duplicate questions
        # ----------------------------------------------------

        if topic in seen:
            continue

        seen.add(topic)

        # ----------------------------------------------------
        # Get correct answer
        # ----------------------------------------------------

        answer_letter = str(
            row["answer"]
        ).strip()

        if answer_letter in [
            "A",
            "B",
            "C",
            "D",
            "E"
        ]:

            answer_text = str(
                row[answer_letter]
            )

        else:

            answer_text = answer_letter

        retrieved.append(
            {
                "score": float(score),

                "question": str(
                    row["canonical_question"]
                    if "canonical_question"
                    in row.index
                    else row["prompt"]
                ),

                "answer_letter":
                    answer_letter,

                "answer_text":
                    answer_text
            }
        )

        if len(retrieved) >= top_k:
            break

    return retrieved


# ============================================================
# BUILD RAG PROMPT
# ============================================================

def build_rag_prompt(
    question,
    choices
):

    examples = retrieve_examples(
        question,
        choices
    )

    prompt_parts = []

    # --------------------------------------------------------
    # Retrieved examples
    # --------------------------------------------------------

    if examples:

        prompt_parts.append(
            "Here are some similar solved "
            "multiple-choice questions:\n"
        )

        for i, example in enumerate(
            examples,
            start=1
        ):

            prompt_parts.append(
                f"Example {i}:\n"
                f"Question: "
                f"{example['question']}\n"
                f"Correct answer: "
                f"{example['answer_letter']}. "
                f"{example['answer_text']}\n"
            )

    # --------------------------------------------------------
    # Current question
    # --------------------------------------------------------

    prompt_parts.append(
        "\nCurrent question:\n"
        f"{canonical_question(question)}\n\n"
        "Options:\n"
        f"A. {choices['A']}\n"
        f"B. {choices['B']}\n"
        f"C. {choices['C']}\n"
        f"D. {choices['D']}\n"
        f"E. {choices['E']}\n"
    )

    return "\n".join(
        prompt_parts
    ), examples


# ============================================================
# MODEL INFERENCE
# ============================================================

@spaces.GPU
def predict(
    question,
    option_a,
    option_b,
    option_c,
    option_d,
    option_e
):

    # --------------------------------------------------------
    # Validate question
    # --------------------------------------------------------

    if not question or not question.strip():

        return (
            "Please enter a question.",
            ""
        )

    choices = {
        "A": option_a,
        "B": option_b,
        "C": option_c,
        "D": option_d,
        "E": option_e
    }

    # ========================================================
    # RAG RETRIEVAL
    # ========================================================

    rag_prompt, examples = build_rag_prompt(
        question,
        choices
    )

    # ========================================================
    # CREATE FIVE QUESTION-OPTION PAIRS
    # ========================================================

    questions = [
        rag_prompt,
        rag_prompt,
        rag_prompt,
        rag_prompt,
        rag_prompt
    ]

    options = [
        choices["A"],
        choices["B"],
        choices["C"],
        choices["D"],
        choices["E"]
    ]

    # ========================================================
    # TOKENIZATION
    # ========================================================

    encoded = tokenizer(
        questions,
        options,
        padding="max_length",
        truncation=True,
        max_length=MAX_LENGTH,
        return_tensors="pt"
    )

    # ========================================================
    # CREATE MULTIPLE-CHOICE INPUT
    # ========================================================

    inputs = {
        key: value.unsqueeze(0).to(DEVICE)

        for key, value in encoded.items()
    }

    # ========================================================
    # DEBERTA INFERENCE
    # ========================================================

    with torch.no_grad():

        outputs = model(
            **inputs
        )

        probabilities = torch.softmax(
            outputs.logits,
            dim=-1
        )[0]

    # ========================================================
    # RANK ANSWERS
    # ========================================================

    labels = [
        "A",
        "B",
        "C",
        "D",
        "E"
    ]

    ranking = torch.argsort(
        probabilities,
        descending=True
    ).tolist()

    top3 = " ".join(
        labels[i]
        for i in ranking[:3]
    )

    # ========================================================
    # PREDICTION OUTPUT
    # ========================================================

    prediction_text = (
        "### 🧠 Top-3 Prediction\n\n"
        f"**{top3}**\n\n"
        "### Probabilities\n\n"
    )

    for idx in ranking:

        prediction_text += (
            f"**{labels[idx]}**: "
            f"{probabilities[idx].item():.4f}\n\n"
        )

    # ========================================================
    # RETRIEVAL OUTPUT
    # ========================================================

    retrieval_text = (
        "### 🔎 Retrieved Examples\n\n"
    )

    if not examples:

        retrieval_text += (
            "No sufficiently similar examples "
            f"were found "
            f"(threshold = "
            f"{SIMILARITY_THRESHOLD}).\n"
        )

    else:

        for i, example in enumerate(
            examples,
            start=1
        ):

            retrieval_text += (
                f"**Example {i}**  \n"
                f"Similarity: "
                f"`{example['score']:.4f}`\n\n"
            )

            retrieval_text += (
                f"{example['question']}\n\n"
            )

            retrieval_text += (
                f"Correct answer: "
                f"**{example['answer_letter']}**"
                f" — "
                f"{example['answer_text']}\n\n"
            )

            retrieval_text += (
                "---\n\n"
            )

    return (
        prediction_text,
        retrieval_text
    )


# ============================================================
# GRADIO UI
# ============================================================

with gr.Blocks(
    title="Science MCQ Solver"
) as demo:

    gr.Markdown(
        """
# 🧠 Science MCQ Solver

### DeBERTa-v3-small + RAG

Enter a multiple-choice science question and its five
answer options.

The system retrieves semantically similar training questions
using **Sentence Transformers + FAISS**, adds them as context,
and then uses the fine-tuned **DeBERTa-v3-small** model to
rank the five answer choices.
"""
    )

    # --------------------------------------------------------
    # Question
    # --------------------------------------------------------

    question = gr.Textbox(
        label="Question",
        placeholder=(
            "Enter the science question..."
        ),
        lines=4
    )

    # --------------------------------------------------------
    # Options
    # --------------------------------------------------------

    with gr.Row():

        option_a = gr.Textbox(
            label="Option A"
        )

        option_b = gr.Textbox(
            label="Option B"
        )

    with gr.Row():

        option_c = gr.Textbox(
            label="Option C"
        )

        option_d = gr.Textbox(
            label="Option D"
        )

    option_e = gr.Textbox(
        label="Option E"
    )

    # --------------------------------------------------------
    # Buttons
    # --------------------------------------------------------

    with gr.Row():

        submit_btn = gr.Button(
            "🚀 Predict Top-3",
            variant="primary"
        )

        clear_btn = gr.ClearButton(
            [
                question,
                option_a,
                option_b,
                option_c,
                option_d,
                option_e
            ],
            value="🗑️ Clear"
        )

    # --------------------------------------------------------
    # Outputs
    # --------------------------------------------------------

    gr.Markdown(
        "## 📊 Prediction"
    )

    prediction_output = gr.Markdown()

    gr.Markdown(
        "## 🔎 Retrieved Context"
    )

    retrieval_output = gr.Markdown()

    # --------------------------------------------------------
    # Submit button
    # --------------------------------------------------------

    submit_btn.click(
        fn=predict,

        inputs=[
            question,
            option_a,
            option_b,
            option_c,
            option_d,
            option_e
        ],

        outputs=[
            prediction_output,
            retrieval_output
        ]
    )

    # --------------------------------------------------------
    # Enter key submission
    # --------------------------------------------------------

    question.submit(
        fn=predict,

        inputs=[
            question,
            option_a,
            option_b,
            option_c,
            option_d,
            option_e
        ],

        outputs=[
            prediction_output,
            retrieval_output
        ]
    )


# ============================================================
# LAUNCH
# ============================================================

if __name__ == "__main__":

    demo.launch()