# SentimentCNN class based on Sentiment Analysis tutorial by Ben Trevett
# https://github.com/bentrevett/pytorch-sentiment-analysis

import torch
import torch.nn as nn
import torchtext
 
class SentimentCNN(nn.Module):
    def __init__(self, state_dict=None, vocab=None, tokenizer='basic_english'):
        super().__init__()

        # tokenizer setup
        self.tokenizer = torchtext.data.utils.get_tokenizer(tokenizer)
        self.state_dict_name = state_dict

        if vocab:
            self.load_vocab(vocab)

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    def _setup_model(self):
        # cnn parameters
        n_filters=100
        filter_sizes=[3,5,7]
        dropout_rate=0.25
        self.min_length = max(filter_sizes)

        # language space parameters
        embedding_dim=300
        output_dim=2

        # model setup
        self.embedding = nn.Embedding(
               len(self.vocab),
               embedding_dim,
               padding_idx=self.pad_index)
        self.convs = nn.ModuleList([nn.Conv1d(embedding_dim, 
                                              n_filters, 
                                              filter_size) 
                                    for filter_size in filter_sizes])
        self.fc = nn.Linear(len(filter_sizes) * n_filters, output_dim)
        self.dropout = nn.Dropout(dropout_rate)

        if self.state_dict_name:
            self.load_state_dict(torch.load(self.state_dict_name))

    def load_vocab(self, vocab):
        # vocabulary parameters
        self.vocab = torch.load(vocab)
        self.pad_index = self.vocab['<pad>']
        self._setup_model()
        
    def forward(self, ids):
        embedded = self.dropout(self.embedding(ids))
        embedded = embedded.permute(0,2,1)
        conved = [torch.relu(conv(embedded)) for conv in self.convs]
        pooled = [conv.max(dim=-1).values for conv in conved]
        cat = self.dropout(torch.cat(pooled, dim=-1))
        prediction = self.fc(cat)
        return prediction

    def predict_sentiment(self, text):
        tokens = self.tokenizer(text)
        ids = [self.vocab[t] for t in tokens]
        if len(ids) < self.min_length:
            ids += [self.pad_index] * (self.min_length - len(ids))
        tensor = torch.LongTensor(ids).unsqueeze(dim=0).to(self.device)
        prediction = self(tensor).squeeze(dim=0)
        probability = torch.softmax(prediction, dim=-1)
        predicted_class = prediction.argmax(dim=-1).item()
        predicted_probability = probability[predicted_class].item()

        return predicted_class, predicted_probability