import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
import spaces

from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer


# =========================================================
# Device
# =========================================================

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

print("Device:", device)


# =========================================================
# Hugging Face repository
# =========================================================

MODEL_REPO = "rohitk123/strongtextcnn-mcq-solver"


# =========================================================
# StrongTextCNN
# =========================================================

class StrongTextCNN(nn.Module):

    def __init__(
        self,
        vocab_size,
        embed_dim=300,
        num_filters=128,
        dropout=0.4
    ):
        super().__init__()

        self.embedding = nn.Embedding(
            vocab_size,
            embed_dim,
            padding_idx=0
        )

        self.convs = nn.ModuleList([
            nn.Conv1d(
                embed_dim,
                num_filters,
                k
            )
            for k in [2, 3, 4, 5]
        ])

        self.bn = nn.BatchNorm1d(
            num_filters * 4
        )

        self.fc1 = nn.Linear(
            num_filters * 4,
            256
        )

        self.fc2 = nn.Linear(
            256,
            64
        )

        self.out = nn.Linear(
            64,
            1
        )

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):

        x = self.embedding(x)

        x = x.permute(0, 2, 1)

        conv_outputs = []

        for conv in self.convs:

            c = F.relu(conv(x))

            p = F.max_pool1d(
                c,
                kernel_size=c.shape[2]
            ).squeeze(2)

            conv_outputs.append(p)

        x = torch.cat(
            conv_outputs,
            dim=1
        )

        x = self.bn(x)

        x = self.dropout(x)

        x = F.relu(self.fc1(x))

        x = self.dropout(x)

        x = F.relu(self.fc2(x))

        x = self.dropout(x)

        x = self.out(x)

        return x.squeeze(1)


# =========================================================
# Download model checkpoint
# =========================================================

print("Downloading model...")

model_file = hf_hub_download(
    repo_id=MODEL_REPO,
    filename="best_model.pt"
)

print("Model downloaded:", model_file)


# =========================================================
# Download tokenizer
# =========================================================

print("Downloading tokenizer...")

tokenizer_file = hf_hub_download(
    repo_id=MODEL_REPO,
    filename="tokenizer/tokenizer.json"
)

print("Tokenizer downloaded:", tokenizer_file)


# =========================================================
# Load tokenizer
# =========================================================

tokenizer = Tokenizer.from_file(
    tokenizer_file
)

print(
    "Tokenizer loaded successfully!"
)

print(
    "Vocabulary size:",
    tokenizer.get_vocab_size()
)


# =========================================================
# Load checkpoint
# =========================================================

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

print("Checkpoint loaded!")

print(
    "Best F1:",
    checkpoint["best_f1"]
)


# =========================================================
# Create model
# =========================================================

model = StrongTextCNN(
    vocab_size=checkpoint["vocab_size"],
    embed_dim=checkpoint["embed_dim"],
    num_filters=checkpoint["num_filters"],
    dropout=checkpoint["dropout"]
)

model.load_state_dict(
    checkpoint["model_state_dict"]
)

model.to(device)

model.eval()

print(
    "StrongTextCNN loaded successfully!"
)


# =========================================================
# Prediction
# =========================================================

@spaces.GPU(duration=30)
def predict(
    question,
    A,
    B,
    C,
    D,
    E
):

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

    probabilities = []

    for choice in choices:

        # SAME preprocessing as training
        text = (
            "Pick the best possible answer: "
            + question
            + " "
            + choice
        )

        # Tokenize
        encoding = tokenizer.encode(
            text
        )

        # Get token IDs
        ids = encoding.ids

        # Truncate
        ids = ids[:256]

        # Padding
        if len(ids) < 256:

            ids = ids + [
                0
            ] * (256 - len(ids))

        # Convert to tensor
        input_ids = torch.tensor(
            [ids],
            dtype=torch.long,
            device=device
        )

        # Prediction
        with torch.no_grad():

            logit = model(
                input_ids
            )

            probability = torch.sigmoid(
                logit
            ).item()

        probabilities.append(
            probability
        )


    # =====================================================
    # Prediction
    # =====================================================

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

    prediction = max(
        range(5),
        key=lambda i: probabilities[i]
    )

    answer = letters[prediction]


    # =====================================================
    # Result
    # =====================================================

    result = (
        f"Predicted Answer: {answer}\n\n"
        "Scores\n\n"
    )

    for letter, probability in zip(
        letters,
        probabilities
    ):

        result += (
            f"{letter}: "
            f"{probability:.4f}\n"
        )

    return result


# =========================================================
# Gradio
# =========================================================

demo = gr.Interface(

    fn=predict,

    inputs=[

        gr.Textbox(
            lines=4,
            label="Question"
        ),

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

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

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

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

        gr.Textbox(
            label="Option E"
        )
    ],

    outputs=gr.Textbox(
        label="Prediction"
    ),

    title="StrongTextCNN MCQ Solver",

    description=(
        "MCQ solver using StrongTextCNN."
    ),

    examples=[
        [
            "Which planet is known as the Red Planet?",
            "Earth",
            "Mars",
            "Venus",
            "Jupiter",
            "Saturn"
        ]
    ]
)


# =========================================================
# Launch
# =========================================================

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