import os
import re
import gradio as gr
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer
from transformers import T5ForConditionalGeneration, T5Tokenizer

# Device Configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

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

# Preprocessing helper
INSTRUCTION_PREFIXES = [
    r"^pick the best possible answer:\s*",
    r"^determine the correct option:\s*",
    r"^select the most accurate option:\s*",
    r"^identify the correct statement:\s*",
    r"^which of the following is correct\??\s*",
    r"^choose the correct answer:\s*",
]
INSTRUCTION_SUFFIXES = [
    r"\s*among the listed options\.?$",
    r"\s*carefully\.?$",
    r"\s*based on the given context\.?$",
    r"\s*from the following choices\.?$",
    r"\s*based on the context\.?$",
]

def clean_prompt(text: str) -> str:
    t = str(text).lower()
    for p in INSTRUCTION_PREFIXES:
        t = re.sub(p, "", t)
    for s in INSTRUCTION_SUFFIXES:
        t = re.sub(s, "", t)
    t = t.strip()
    return t[0].upper() + t[1:] if t else t

# Load Models
print("Loading Model Components for Hugging Face Deployment...")
flan_tok = T5Tokenizer.from_pretrained("google/flan-t5-large")
flan_model = T5ForConditionalGeneration.from_pretrained(
    "google/flan-t5-large",
    torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,
).to(device).eval()

mini_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

DEC_START = torch.tensor([[flan_model.config.decoder_start_token_id]], device=device)
ANSWER_IDS = [flan_tok.encode(c, add_special_tokens=False)[0] for c in OPTION_COLS]

@torch.no_grad()
def predict_mcq(prompt, opt_a, opt_b, opt_c, opt_d, opt_e, flan_weight=0.7):
    cleaned_q = clean_prompt(prompt)
    options = [opt_a, opt_b, opt_c, opt_d, opt_e]
    
    # --- 1. Generative Scoring (Flan-T5) ---
    opts_text = "\n".join([f"{c}) {text}" for c, text in zip(OPTION_COLS, options)])
    flan_prompt = (
        "Answer the following multiple-choice science question. "
        "Choose the single best answer and reply with ONLY the letter.\n\n"
        f"Question: {cleaned_q}\n\n{opts_text}\n\nAnswer:"
    )

    enc = flan_tok(flan_prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
    out = flan_model(
        input_ids=enc["input_ids"],
        attention_mask=enc["attention_mask"],
        decoder_input_ids=DEC_START,
    )
    log_probs = F.log_softmax(out.logits[0, 0, :], dim=-1)
    flan_scores = log_probs[ANSWER_IDS].cpu().float().numpy()
    flan_soft = np.exp(flan_scores)
    flan_probs = flan_soft / (np.sum(flan_soft) + 1e-9)

    # --- 2. Semantic Embedding Scoring (MiniLM) ---
    q_emb = mini_model.encode(cleaned_q, normalize_embeddings=True)
    o_embs = mini_model.encode(options, normalize_embeddings=True)
    minilm_scores = np.dot(o_embs, q_emb)
    
    # Min-max normalization
    min_s, max_s = minilm_scores.min(), minilm_scores.max()
    minilm_probs = (minilm_scores - min_s) / (max_s - min_s + 1e-9)

    # --- 3. Weighted Hybrid Blend ---
    final_probs = flan_weight * flan_probs + (1.0 - flan_weight) * minilm_probs
    ranked_indices = np.argsort(final_probs)[::-1]

    # Outputs
    top_3_list = [OPTION_COLS[i] for i in ranked_indices[:3]]
    top_3_str = " ".join(top_3_list)
    
    probabilities_dict = {
        OPTION_COLS[i]: round(float(final_probs[i]), 4) for i in ranked_indices
    }

    return top_3_str, probabilities_dict

# Gradio Interface Setup
demo = gr.Interface(
    fn=predict_mcq,
    inputs=[
        gr.Textbox(label="Question Prompt", lines=3, placeholder="What is the Josephson effect?"),
        gr.Textbox(label="Option A", placeholder="Exploited by superconducting devices such as SQUIDs."),
        gr.Textbox(label="Option B", placeholder="Exploited by magnetic devices such as SQUIDs."),
        gr.Textbox(label="Option C", placeholder="Used for measuring electric flux quantum."),
        gr.Textbox(label="Option D", placeholder="A classical optical diffraction phenomenon."),
        gr.Textbox(label="Option E", placeholder="Described by Maxwell thermodynamics."),
        gr.Slider(minimum=0.0, maximum=1.0, value=0.7, label="Ensemble Weight Alpha (Flan-T5 weight)"),
    ],
    outputs=[
        gr.Textbox(label="Predicted Top-3 Answer Ranking (MAP@3 Target)"),
        gr.Label(label="Option Probabilities"),
    ],
    title="Smart MCQ Solver Challenge — Hugging Face Space",
    description="An end-to-end Machine Learning pipeline deploying a hybrid zero-shot ensemble model (Flan-T5 + MiniLM) for solving STEM multiple choice questions.",
    examples=[
        [
            "What is the Josephson effect?",
            "Exploited by superconducting devices such as SQUIDs.",
            "Exploited by magnetic devices such as SQUIDs.",
            "Used for measuring electric flux quantum.",
            "A classical optical diffraction phenomenon.",
            "Described by Maxwell thermodynamics.",
            0.7
        ]
    ]
)

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