!pip install qlearnkit['pennylane']
!pip install matplotlib

import qlearnkit as ql
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from matplotlib import pyplot as plt


tag_to_ix = {"DET": 0, "NN": 1, "V": 2}  # Assign each tag with a unique index
ix_to_tag = {i: k for k, i in tag_to_ix.items()}


def prepare_sequence(seq, to_ix):
    idxs = [to_ix[w] for w in seq]
    return torch.tensor(idxs, dtype=torch.long)


training_data = [
    # Tags are: DET - determiner; NN - noun; V - verb
    # For example, the word "The" is a determiner
    ("The dog ate the apple".split(), ["DET", "NN", "V", "DET", "NN"]),
    ("Everybody read that book".split(), ["NN", "V", "DET", "NN"])
]
word_to_ix = {}

# For each words-list (sentence) and tags-list in each tuple of training_data
for sent, tags in training_data:
    for word in sent:
        if word not in word_to_ix:  # word has not been assigned an index yet
            word_to_ix[word] = len(word_to_ix)  # Assign each word with a unique index

print(f"Vocabulary: {word_to_ix}")
print(f"Entities: {ix_to_tag}")

class LSTMTagger(nn.Module):

    def __init__(self,
                 model,
                 embedding_dim,
                 hidden_dim,
                 vocab_size,
                 tagset_size):
        super(LSTMTagger, self).__init__()
        self.hidden_dim = hidden_dim

        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)

        self.model = model


        # The linear layer that maps from hidden state space to tag space
        self.hidden2tag = nn.Linear(hidden_dim, tagset_size)

    def forward(self, sentence):
        embeds = self.word_embeddings(sentence)
        lstm_out, _ = self.model(embeds.view(len(sentence), 1, -1))
        tag_logits = self.hidden2tag(lstm_out.view(len(sentence), -1))
        tag_scores = F.log_softmax(tag_logits, dim=1)
        return tag_scores

embedding_dim = 8
hidden_dim = 6
n_layers = 1
n_qubits = 4
n_epochs = 300
device = 'default.qubit'

def trainer(lstm_model, embedding_dim, hidden_dim,
            n_epochs, model_label):
    # the LSTM Tagger Model
    model = LSTMTagger(lstm_model,
                   embedding_dim,
                   hidden_dim,
                   vocab_size=len(word_to_ix),
                   tagset_size=len(tag_to_ix))
    # loss function and Stochastic Gradient Descend
    # as optimizers
    loss_function = nn.NLLLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.1)

    history = {
        'loss': [],
        'acc': []
    }
    for epoch in range(n_epochs):
        losses = []
        preds = []
        targets = []
        for sentence, tags in training_data:
            # Step 1. Remember that Pytorch accumulates gradients.
            # We need to clear them out before each instance
            model.zero_grad()

            # Step 2. Get our inputs ready for the network, that is, turn them into
            # Tensors of word indices.
            sentence_in = prepare_sequence(sentence, word_to_ix)
            labels = prepare_sequence(tags, tag_to_ix)

            # Step 3. Run our forward pass.
            tag_scores = model(sentence_in)

            # Step 4. Compute the loss, gradients, and update the parameters by
            #  calling optimizer.step()
            loss = loss_function(tag_scores, labels)
            loss.backward()
            optimizer.step()
            losses.append(float(loss))

            probs = torch.softmax(tag_scores, dim=-1)
            preds.append(probs.argmax(dim=-1))
            targets.append(labels)
        avg_loss = np.mean(losses)
        history['loss'].append(avg_loss)

        # print("preds", preds)
        preds = torch.cat(preds)
        targets = torch.cat(targets)
        corrects = (preds == targets)
        accuracy = corrects.sum().float() / float(targets.size(0))
        history['acc'].append(accuracy)

        print(f"Epoch {epoch + 1} / {n_epochs}: Loss = {avg_loss:.3f} Acc = {accuracy:.2f}")

    with torch.no_grad():
        input_sentence = training_data[0][0]
        labels = training_data[0][1]
        inputs = prepare_sequence(input_sentence, word_to_ix)
        tag_scores = model(inputs)

        tag_ids = torch.argmax(tag_scores, dim=1).numpy()
        tag_labels = [ix_to_tag[k] for k in tag_ids]
        print(f"Sentence:  {input_sentence}")
        print(f"Labels:    {labels}")
        print(f"Predicted: {tag_labels}")

    fig, ax1 = plt.subplots()

    ax1.set_xlabel("Epoch")
    ax1.set_ylabel("Loss")
    ax1.plot(history['loss'], label=f"{model_label} Loss")

    ax2 = ax1.twinx()
    ax2.set_ylabel("Accuracy")
    ax2.plot(history['acc'], label=f"{model_label} LSTM Accuracy", color='tab:red')

    plt.title("Part-of-Speech Tagger Training")
    plt.ylim(0., 1.5)
    plt.legend(loc="upper right")

    plt.show()

import qlearnkit.nn as qnn
qlstm = qnn.QLongShortTermMemory(
    embedding_dim,
    hidden_dim,
    n_layers,
    n_qubits=n_qubits,
    device=device
)

# random values to feed into the QNode
inputs = torch.rand(embedding_dim, n_qubits)
weights = torch.rand(n_layers, n_qubits, 3)
import pennylane as qml
dev = qml.device(device, wires=n_qubits)
circ = qml.QNode(qlstm._construct_vqc, dev)

print(qml.draw(circ, expansion_strategy='device', show_all_wires=True)(inputs, weights))

trainer(qlstm, embedding_dim, hidden_dim, n_epochs, model_label='Hybrid Quantum')

clstm = nn.LSTM(embedding_dim, hidden_dim, num_layers=n_layers)
trainer(clstm, embedding_dim, hidden_dim, n_epochs,model_label='Classical')