# Our raw data, which consists of sentences
corpus = [
          "We always come to Paris",
          "The professor is from Australia",
          "I live in Stanford",
          "He comes from Taiwan",
          "The capital of Turkey is Ankara"
         ]

import re
s = "We! always come to Paris.12" # only keep letters
re.sub(r'[^A-Za-z ]+', '', s)

# simple lowercase all and split(by space) into words
def preprocess_sentence(sentence):
  return re.sub(r'[^A-Za-z ]+', '',sentence).lower().split()

train_sentences = [preprocess_sentence(sent) for sent in corpus]
train_sentences

# Set of locations that appear in our corpus
locations = set(["australia", "ankara", "paris", "stanford", "taiwan", "turkey"])

# Our train labels
train_labels = [[1 if word in locations else 0 for word in sent] for sent in train_sentences]
train_labels

vocabulary = set(w for s in train_sentences for w in s)
vocabulary.add("<unk>") # add the unknown token
vocabulary.add("<pad>") # add the padding for window
len(vocabulary)

# notice that the total window size is 2*window_size+1
def pad_window(sentence, window_size, pad_token = "<pad>"):
  window = [pad_token] * window_size
  return window + sentence + window

pad_window(train_sentences[1], window_size=2)

idx_to_word = sorted(list(vocabulary)) 

word_to_idx = {word: idx for idx, word in enumerate(idx_to_word)}

def convert_token_to_idx(sentence, word_to_idx):
  return [word_to_idx.get(token, word_to_idx["<unk>"]) for token in sentence]

print(f"From token list to indices: {convert_token_to_idx(train_sentences[0], word_to_idx)}")

import torch
import torch.nn as nn

embedding_dim = 5 # embedding dimension is 5
embeds = nn.Embedding(len(vocabulary), embedding_dim) # create an embedding table

list(embeds.named_parameters()) # current embeddings for each word

# Get embeddings for words (if we want to do something later...)
indices = torch.tensor([word_to_idx[v] for v in ["paris", "ankara"]], dtype=torch.long)
embeddings = embeds(indices)
embeddings

from torch.utils.data import DataLoader
from functools import partial

def custom_collate_fn(batch, window_size, word_to_idx):
  x, y = zip(*batch)

  def pad_window(sentence, window_size, pad_token="<pad>"):
    window = [pad_token]*window_size
    return window + sentence +window

  x = [pad_window(s, window_size=window_size) for s in x]

  def convert_token_to_idx(sentence, word_to_idx):
    return [word_to_idx.get(token, word_to_idx["<unk>"]) for token in sentence] 
    # use get here to have a default value for words not in dictionary
  
  x = [convert_token_to_idx(s, word_to_idx) for s in x]
  pad_token_idx = word_to_idx["<pad>"]

  # pad all sentences to equal length
  x = [torch.LongTensor(x_i) for x_i in x]
  x_padded = nn.utils.rnn.pad_sequence(x, batch_first=True, padding_value=pad_token_idx)
  
  lengths = [len(label) for label in y]
  lengths = torch.LongTensor(lengths)

  y = [torch.LongTensor(y_i) for y_i in y]
  y_padded = nn.utils.rnn.pad_sequence(y, batch_first=True, padding_value=pad_token_idx)

  return x_padded, y_padded, lengths

# parameters
data = list(zip(train_sentences, train_labels))
batch_size = 2
shuffle = True
window_size = 2
collate_fn = partial(custom_collate_fn, window_size=window_size, word_to_idx=word_to_idx)

# instantiate
loader = DataLoader(data, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn)

counter = 0
for batched_x, batched_y, batched_lengths in loader:
  print(f"Iteration {counter}")
  print("Batched Input:")
  print(batched_x)
  print("Batched Labels:")
  print(batched_y)
  print("Batched Lengths:")
  print(batched_lengths)
  print("")
  counter += 1

# create windows using unfold function
chunk = batched_x.unfold(1, window_size*2+1, 1)
print(chunk)

class WordWindowClassifier(nn.Module):
  def __init__(self, param, vocab_size, pad_idx=0):
    super(WordWindowClassifier, self).__init__()

    self.window_size = param["window_size"]
    self.embed_dim = param["embed_dim"]
    self.hidden_dim = param["hidden_dim"]
    self.freeze_embeddings = param["freeze_embeddings"]

    # embedding layer
    self.embeds = nn.Embedding(vocab_size, self.embed_dim, padding_idx=pad_idx)
    # if freeze_embeddings, set require grad to false
    if self.freeze_embeddings:
      self.embed_layer.weight.requires_grad = False 

    """ Hidden Layer
    """
    full_window_size = 2*window_size+1
    self.hidden_layer = nn.Sequential(
        nn.Linear(full_window_size * self.embed_dim, self.hidden_dim),
        nn.Tanh()
    )  

    """ Output Layer
    """
    self.output_layer = nn.Linear(self.hidden_dim, 1)

    self.prob = nn.Sigmoid()

  def forward(self, inputs):
    B, L = inputs.size()

    token_windows = inputs.unfold(1, 2*self.window_size+1, 1)
    _,adjusted_length,_ = token_windows.size()

    assert token_windows.size() == (B, adjusted_length, 2*self.window_size+1)

    # embedding layer
    embedded_windows = self.embeds(token_windows)

    # reshape to combine dim of windows and embeddings
    embedded_windows = embedded_windows.view(B, adjusted_length, -1)

    layer_1 = self.hidden_layer(embedded_windows)

    output = self.output_layer(layer_1)

    output = self.prob(output)
    output = output.view(B, -1)
    
    return output

data = list(zip(train_sentences, train_labels))
batch_size = 2
shuffle = True
window_size = 2
collate_fn = partial(custom_collate_fn, window_size=window_size, word_to_idx=word_to_idx)

loader = DataLoader(data, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn
                    )

model_param = {
    "batch_size": 4,
    "window_size": 2,
    "embed_dim": 25,
    "hidden_dim": 25,
    "freeze_embeddings": False
}

vocab_size = len(word_to_idx)
model = WordWindowClassifier(model_param, vocab_size)

# optimizer
learning_rate = 0.01
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

# loss function
def loss_function(batch_outputs, batch_labels, batch_lengths):
  bceloss = nn.BCELoss()
  loss = bceloss(batch_outputs, batch_labels.float())

  loss = loss/batch_labels.sum().float()

  return loss

def train_epoch(loss_function, optimizer, model, loader):

  total_loss = 0
  for batch_inputs, batch_labels, batch_lengths in loader:
    # clear gradients
    optimizer.zero_grad()
    # forward pass
    outputs = model.forward(batch_inputs)
    # compute loss
    loss = loss_function(outputs, batch_labels, batch_lengths)
    # gradients
    loss.backward()
    # update params
    optimizer.step()
    total_loss += loss.item()

  return total_loss

def train(loss_function, optimizer, model, loader, num_epochs=10000):
  for epoch in range(num_epochs):
    epoch_loss = train_epoch(loss_function, optimizer, model, loader)
    if epoch%100 == 0:
      print(epoch_loss)

num_epochs = 1000
train(loss_function, optimizer, model, loader, num_epochs=num_epochs)

test_corpus = ["She comes from Paris",
               "She comes from China"]
test_sentences = [preprocess_sentence(sent) for sent in test_corpus]
test_labels = [[0, 0, 0, 1],[0, 0, 0, 1]]

test_data = list(zip(test_sentences, test_labels))
batch_size = 1
shuffle = False
window_size = 2
collate_fn = partial(custom_collate_fn, window_size=2, word_to_idx=word_to_idx)
test_loader = torch.utils.data.DataLoader(test_data, 
                                           batch_size=1, 
                                           shuffle=False, 
                                           collate_fn=collate_fn)

for test_instances, labels, _ in test_loader:
  outputs = model.forward(test_instances)
  print(labels)
  print(outputs)

predict_probs = []
for test_instances, labels, _ in test_loader:
  outputs = model.forward(test_instances)
  predict_probs.append(outputs.detach().numpy())
  print(labels)
  print(outputs)

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from IPython.display import display, HTML

def colorize(words, color_array):
    cmap=matplotlib.cm.RdYlGn
    template = '<span class="barcode"; style="color: white; background-color: {}">{}</span>'
    colored_string = ''
    for word, color in zip(words, color_array):
        color = matplotlib.colors.rgb2hex(cmap(color)[:3])
        colored_string += template.format(color, '&nbsp' + word + '&nbsp')
    return colored_string

# # or simply save in an html file and open in browser
# with open('colorize.html', 'w') as f:
#     f.write(s)

for sentence, prediction in zip(test_sentences, predict_probs):
  s = colorize(sentence, prediction[0])
  display(HTML(s))