import torch
import torch.nn as nn
import lightning as L
import torchmetrics as tm
from tokenizers import Tokenizer
import gradio as gr
from huggingface_hub import hf_hub_download

COARSE_LABELS = [
    "ABBR (0): Abbreviation",
    "ENTY (1): Entity",
    "DESC (2): Description and abstract concept",
    "HUM (3): Human being",
    "LOC (4): Location",
    "NUM (5): Numeric value",
]

FINE_LABELS = [
    "ABBR (0): Abbreviation",
    "ABBR (1): Expression abbreviated",
    "ENTY (2): Animal",
    "ENTY (3): Organ of body",
    "ENTY (4): Color",
    "ENTY (5): Invention, book and other creative piece",
    "ENTY (6): Currency name",
    "ENTY (7): Disease and medicine",
    "ENTY (8): Event",
    "ENTY (9): Food",
    "ENTY (10): Musical instrument",
    "ENTY (11): Language",
    "ENTY (12): Letter like a-z",
    "ENTY (13): Other entity",
    "ENTY (14): Plant",
    "ENTY (15): Product",
    "ENTY (16): Religion",
    "ENTY (17): Sport",
    "ENTY (18): Element and substance",
    "ENTY (19): Symbols and sign",
    "ENTY (20): Techniques and method",
    "ENTY (21): Equivalent term",
    "ENTY (22): Vehicle",
    "ENTY (23): Word with a special property",
    "DESC (24): Definition of something",
    "DESC (25): Description of something",
    "DESC (26): Manner of an action",
    "DESC (27): Reason",
    "HUM (28): Group or organization of persons",
    "HUM (29): Individual",
    "HUM (30): Title of a person",
    "HUM (31): Description of a person",
    "LOC (32): City",
    "LOC (33): Country",
    "LOC (34): Mountain",
    "LOC (35): Other location",
    "LOC (36): State",
    "NUM (37): Postcode or other code",
    "NUM (38): Number of something",
    "NUM (39): Date",
    "NUM (40): Distance, linear measure",
    "NUM (41): Price",
    "NUM (42): Order, rank",
    "NUM (43): Other number",
    "NUM (44): Lasting time of something",
    "NUM (45): Percent, fraction",
    "NUM (46): Speed",
    "NUM (47): Temperature",
    "NUM (48): Size, area and volume",
    "NUM (49): Weight",
]


class Classifier:
    def __init__(self, tokenizer_ckpt_path, model_ckpt_path):
        self.tokenizer = Tokenizer.from_file(tokenizer_ckpt_path)
        self.model = LSTMWithAttentionClassifier.load_from_checkpoint(
            model_ckpt_path,
            map_location="cpu",
        )

    def predict(self, text):
        encoding = self.tokenizer.encode(text)
        ids = torch.tensor([encoding.ids])
        logits, _ = self.model(ids)
        probs = torch.softmax(logits, dim=1).squeeze().tolist()
        return {
            category: prob
            for category, prob in zip(
                FINE_LABELS if self.model.fine else COARSE_LABELS, probs
            )
        }


class Attention(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.WQuery = nn.Linear(hidden_dim, hidden_dim)
        self.WKey = nn.Linear(hidden_dim, hidden_dim)
        self.WValue = nn.Linear(hidden_dim, 1)

    def forward(self, x):
        query = torch.tanh(self.WQuery(x))
        key = torch.tanh(self.WKey(x))

        attention_weights = torch.softmax(self.WValue(query + key), dim=1)

        return (attention_weights * x).sum(dim=1), attention_weights


class LSTMWithAttentionClassifier(L.LightningModule):
    def __init__(
        self,
        vocab_size,
        embedding_dim,
        hidden_dim,
        num_classes,
        lr=1e-3,
        weight_decay=1e-2,
        num_layers=1,
        bidirectional=False,
        dropout=0.0,
        padding_idx=3,
        fine=False,
        **kwargs,
    ):
        super().__init__()
        self.save_hyperparameters()
        self.lr = lr
        self.weight_decay = weight_decay
        self.fine = fine

        self.embedding = nn.Embedding(
            vocab_size,
            embedding_dim,
            padding_idx=padding_idx,
        )
        self.lstm = nn.LSTM(
            embedding_dim,
            hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=bidirectional,
            dropout=dropout,
        )
        self.attention = Attention(
            hidden_dim * (1 + bidirectional),
        )
        self.fc = nn.Linear(
            hidden_dim * (1 + bidirectional),
            num_classes,
        )

        self.criteria = nn.CrossEntropyLoss()
        self.accuracy = tm.Accuracy(
            task="multiclass",
            num_classes=num_classes,
        )

    def forward(self, input_ids):
        x = self.embedding(input_ids)
        x, _ = self.lstm(x)
        x, attention_weights = self.attention(x)
        x = self.fc(x)
        return x, attention_weights

    def training_step(self, batch, batch_idx):
        input_ids = batch["input_ids"]
        coarse = batch["coarse"]
        fine = batch["fine"]
        logits, _ = self(input_ids)
        loss = self.criteria(logits, fine if self.fine else coarse)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        input_ids = batch["input_ids"]
        coarse = batch["coarse"]
        fine = batch["fine"]
        logits, _ = self(input_ids)
        loss = self.criteria(logits, fine if self.fine else coarse)
        self.log("val_loss", loss)
        pred = logits.argmax(dim=1)
        self.accuracy(pred, fine if self.fine else coarse)
        self.log("val_acc", self.accuracy, prog_bar=True)

    def configure_optimizers(self):
        return torch.optim.AdamW(
            self.parameters(),
            lr=self.lr,
            weight_decay=self.weight_decay,
        )


tokenizer_ckpt_path = hf_hub_download(
    repo_id="SatwikKambham/trec-classifier",
    filename="tokenizer.json",
)
model_ckpt_path = hf_hub_download(
    repo_id="SatwikKambham/trec-classifier",
    filename="lstm_attention.ckpt",
)
classifier = Classifier(tokenizer_ckpt_path, model_ckpt_path)
interface = gr.Interface(
    fn=classifier.predict,
    inputs=gr.components.Textbox(
        label="Question",
        placeholder="Enter a question here...",
    ),
    outputs=gr.components.Label(
        label="Predicted class",
        num_top_classes=3,
    ),
    examples=[
        [
            "What does LOL mean?",
        ],
        [
            "What is the meaning of life?",
        ],
        [
            "How long does it take for light from the sun to reach the earth?",
        ],
        [
            "When is friendship day?",
        ],
    ],
)
interface.launch()