# python
from collections import defaultdict
from pathlib import Path

import os

# pypi
from dotenv import load_dotenv
from expects import expect, contain_exactly

import nltk
import numpy
import pandas

# my other stuff
from graeae import Timer

rows, columns = data.shape
print(f"Rows: {rows:,} Columns: {columns}")

print(data.iloc[0])

training_size = 3 * 10**5
training_data = data.iloc[:training_size]
testing_data = data.iloc[training_size:]

assert len(training_data) == training_size

del(data)

duplicates = training_data[training_data.is_duplicate==1]
example = duplicates.iloc[0]
print(example.question1)
print(example.question2)
print(example.is_duplicate)

print(f"There are {len(duplicates):,} duplicates for the training data.")

def words_to_index(words):
    return [vocabulary[word] for word in words]

Q1_train = question_1_train.apply(words_to_index)
Q2_train = question_2_train.apply(words_to_index)

Q1_test = Q1_test_words.apply(words_to_index)
Q2_test = Q2_test_words.apply(words_to_index)

print('first question in the train set:\n')
print(question_1_train.iloc[0], '\n') 
print('encoded version:')
print(Q1_train.iloc[0],'\n')

print(f"{len(vocabulary):,}")

TRAINING_FRACTION = 0.8
cut_off = int(len(question_1_train) * TRAINING_FRACTION)
train_question_1, train_question_2 = Q1_train[:cut_off], Q2_train[:cut_off]
validation_question_1, validation_question_2 = Q1_train[cut_off: ], Q2_train[cut_off:]
print(f"Number of duplicate questions: {len(Q1_train):,}")
print(f"The length of the training set is:  {len(train_question_1):,}")
print(f"The length of the validation set is: {len(validation_question_1):,}")

# python
from collections import defaultdict, namedtuple
from pathlib import Path

import os

# pypi
from dotenv import load_dotenv
from pathlib import Path

import attr
import nltk
import pandas

nltk.download("punkt")

Tokens = namedtuple("Tokens", ["unknown", "padding", "padding_token"])
TOKENS = Tokens(unknown=0,
                padding=1,
                padding_token="<PAD>")

Question = namedtuple("Question", ["question_one", "question_two"])
Data = namedtuple("Data", ["train", "validate", "test", "y_test"])

@attr.s(auto_attribs=True)
class DataTokenizer:
    """Converts questions to tokens

    Args:
     data: the data-frame to tokenize
    """
    data: pandas.DataFrame
    _question_1: pandas.Series=None
    _question_2: pandas.Series=None

@attr.s(auto_attribs=True)
class DataTensorizer:
    """Convert tokenized words to numbers

    Args:
     vocabulary: word to integer mapping
     question_1: data to convert
     question_2: other data to convert
    """
    vocabulary: dict
    question_1: pandas.Series
    question_2: pandas.Series
    _tensorized_1: pandas.Series=None
    _tensorized_2: pandas.Series=None

@attr.s(auto_attribs=True)
class DataLoader:
    """Loads and transforms the data

    Args:
     env: The path to the .env file with the raw-data path
     key: key in the environment with the path to the data
     train_validation_size: number of entries for the training/validation set
     training_fraction: what fraction of the training/valdiation set for training
    """
    env: str="posts/nlp/.env"
    key: str="QUORA_TRAIN"
    train_validation_size: int=300000
    training_fraction: float=0.8
    _data_path: Path=None
    _raw_data: pandas.DataFrame=None
    _training_data: pandas.DataFrame=None
    _testing_data: pandas.DataFrame=None
    _duplicates: pandas.DataFrame=None
    _tokenized_train: DataTokenizer=None
    _tokenized_test: DataTokenizer=None
    _vocabulary: dict=None
    _tensorized_train: DataTensorizer=None
    _tensorized_test: DataTensorizer=None
    _test_labels: pandas.Series=None    
    _data: namedtuple=None

from neurotic.nlp.siamese_networks import DataLoader

loader = DataLoader()

data = loader.data

print(f"Number of duplicate questions: {len(loader.duplicates):,}")
print(f"The length of the training set is:  {len(data.train.question_one):,}")
print(f"The length of the validation set is: {len(data.validate.question_one):,}")

print('first question in the train set:\n')
print(loader.duplicates.question1.iloc[0])
print('encoded version:')
print(data.train.question_one[0],'\n')
expect(data.train.question_one[0]).to(contain_exactly(*Q1_train.iloc[0]))

assert len(loader.vocabulary) == len(vocabulary)
assert not set(vocabulary) - set(loader.vocabulary)
print(f"{len(loader.vocabulary):,}")

# python
from pathlib import Path
import os

# from pypi
from dotenv import load_dotenv

import trax.fastmath.numpy as trax_numpy

load_dotenv("posts/nlp/.env")
PREFIX = "SIAMESE_"
q1 = trax_numpy.load(Path(os.environ[PREFIX + "Q1"]).expanduser())
q2 = trax_numpy.load(Path(os.environ[PREFIX + "Q2"]).expanduser())
v1 = trax_numpy.load(Path(os.environ[PREFIX + "V1"]).expanduser())
v2 = trax_numpy.load(Path(os.environ[PREFIX + "V2"]).expanduser())
y_test = trax_numpy.load(Path(os.environ[PREFIX + "Y_TEST"]).expanduser())

print(f'q1 has shape: {q1.shape} \n\nAnd it looks like this: \n\n {q1}\n\n')

print(f'q2 has shape: {q2.shape} \n\nAnd looks like this: \n\n {q2}\n\n')

print(f'y_test has shape: {y_test.shape} \n\nAnd looks like this: \n\n {y_test}\n\n')

print(f'v1 has shape: {v1.shape} \n\nAnd looks like this: \n\n {v1}\n\n')

print(f'v2 has shape: {v2.shape} \n\nAnd looks like this: \n\n {v2}\n\n')

batch_size = 512
threshold = 0.7
batch = range(batch_size)

correct = 0

for row in batch:
    similarity = trax_numpy.dot(v1[row], v2[row])
    similar_enough = similarity > threshold
    correct += (y_test[element] == similar_enough)

accuracy = correct / batch_size

print(f"The accuracy of the model is: {accuracy:0.4f}.")

# from pypi
import numpy

v1_1 = numpy.array([1, 2, 3])
v1_2 = numpy.array([9, 8, 7])
v1_3 = numpy.array([-1, -4, -2])
v1_4 = numpy.array([1, -7, 2])
v1 = numpy.vstack([v1_1, v1_2, v1_3, v1_4])
print("v1 :")
print(v1, "\n")
v2_1 = v1_1 + numpy.random.normal(0, 2, 3)  # add some noise to create approximate duplicate
v2_2 = v1_2 + numpy.random.normal(0, 2, 3)
v2_3 = v1_3 + numpy.random.normal(0, 2, 3)
v2_4 = v1_4 + numpy.random.normal(0, 2, 3)
v2 = numpy.vstack([v2_1, v2_2, v2_3, v2_4])
print("v2 :")
print(v2, "\n")

assert len(v1) == len(v2)

alpha = 0.25

# pypi
from jax.interpreters.xla import _DeviceArray as DeviceArray
from trax import layers

import numpy
import trax
import trax.fastmath.numpy as fast_numpy

def normalize(x: numpy.ndarray) -> DeviceArray:
    """L2 Normalization

    Args:
     x: the data to normalize

    Returns:
     normalized version of x
    """
    return x / fast_numpy.sqrt(fast_numpy.sum(x * x, axis=-1, keepdims=True))

tensor = numpy.random.random((2,5))
print(f'The tensor is of type: {type(tensor)}\n\nAnd looks like this:\n\n {tensor}')

norm_tensor = normalize(tensor)
print(f'The normalized tensor is of type: {type(norm_tensor)}\n\nAnd looks like this:\n\n {norm_tensor}')

vocab_size = 500
model_dimension = 128

# Define the LSTM model
LSTM = layers.Serial(
        layers.Embedding(vocab_size=vocab_size, d_feature=model_dimension),
        layers.LSTM(model_dimension),
        layers.Mean(axis=1),
        layers.Fn('Normalize', lambda x: normalize(x))
    )

# Use the Parallel combinator to create a Siamese model out of the LSTM 
Siamese = layers.Parallel(LSTM, LSTM)

def show_layers(model, layer_prefix):
    print(f"Total layers: {len(model.sublayers)}\n")
    for i in range(len(model.sublayers)):
        print('========')
        print(f'{layer_prefix}_{i}: {model.sublayers[i]}\n')

print('Siamese model:\n')
show_layers(Siamese, 'Parallel.sublayers')

print('Detail of LSTM models:\n')
show_layers(LSTM, 'Serial.sublayers')

data = NERData().data
model = NER(vocabulary_size=len(data.vocabulary),
            tag_count=len(data.tags)).model
model.init_from_file(Path("~/models/ner/model.pkl.gz", weights_only=True).expanduser())
print(model)

# python
from collections import namedtuple
from pathlib import Path

# pypi
import holoviews
import hvplot.pandas
import jax
import numpy
import pandas
import trax

# this project
from neurotic.nlp.named_entity_recognition import (DataGenerator,
                                                   NER,
                                                   NERData,
                                                   TOKEN)
# another project
from graeae import EmbedHoloviews

x, y = next(test_generator)
print(f"x's shape: {x.shape} y's shape: {y.shape}")

predictions = model(x)
print(type(predictions))
print(f"predictions has shape: {predictions.shape}")

def evaluate_prediction(pred: jax.interpreters.xla._DeviceArray,
                        labels: numpy.ndarray,
                        pad: int=SETTINGS.padding_id) -> float:
    """Calculates the accuracy of a prediction

    Args:
      pred: prediction array with shape 
           (num examples, max sentence length in batch, num of classes)
      labels: array of size (batch_size, seq_len)
      pad: integer representing pad character

    Returns:
      accuracy: fraction of correct predictions
    """
    outputs = numpy.argmax(pred, axis=-1)
    mask = labels != pad
    return numpy.sum((outputs==labels)[mask])/numpy.sum(mask)

accuracy = evaluate_prediction(model(x), y)
print("accuracy: ", accuracy)

repetitions = range(
    int(len(data.data_sets.x_test)/SETTINGS.batch_size))
nexts = (next(test_generator) for repetition in repetitions)
accuracies = [evaluate_prediction(model(x), y) for x, y in nexts]

data = pandas.DataFrame.from_dict(dict(Accuracy=accuracies))
plot = data.Accuracy.hvplot(kind="hist", color=PLOT.tan).opts(
    title="Accuracy Distribution",
    height=PLOT.height,
    width=PLOT.width,
    fontscale=PLOT.fontscale)

output = Embed(plot=plot, file_name="accuracy_distribution")()

print(output)

# from python
from collections import namedtuple
from functools import partial
from tempfile import TemporaryFile

import random
import sys

# from pypi
from holoviews import opts
from trax import layers
from trax.supervised import training

import holoviews
import hvplot.pandas
import pandas
import trax

# this project
from neurotic.nlp.named_entity_recognition import (DataGenerator,
                                                   NER,
                                                   NERData,
                                                   TOKEN)
# another project
from graeae import EmbedHoloviews, Timer

train_generator = trax.data.inputs.add_loss_weights(
    training_generator,
    id_to_mask=ner.data.vocabulary[TOKEN.pad])

evaluate_generator = trax.data.inputs.add_loss_weights(
    validation_generator,
    id_to_mask=ner.data.vocabulary[TOKEN.pad])

def train_model(NER: trax.layers.Serial,
                train_generator: type,
                eval_generator: type,
                train_steps: int=1,
                steps_per_checkpoint: int=100,
                learning_rate: float=SETTINGS.learning_rate,
                verbose: bool=False,
                output_dir="~/models/ner/") -> training.Loop:
    """Train the Named Entity Recognition Model
    Args: 
      NER: the model you are building
      train_generator: The data generator for training examples
      eval_generator: The data generator for validation examples,
      train_steps: number of training steps
      output_dir: folder to save your model

    Returns:
      training_loop: a trax supervised training Loop
    """
    train_task = training.TrainTask(
        labeled_data=train_generator,
        loss_layer = layers.WeightedCategoryCrossEntropy(),
        optimizer = trax.optimizers.Adam(learning_rate),
        n_steps_per_checkpoint=steps_per_checkpoint,
    )

    eval_task = training.EvalTask(
      labeled_data = eval_generator,
      metrics = [layers.WeightedCategoryCrossEntropy(),
                 layers.Accuracy()],
      n_eval_batches = SETTINGS.batch_size
    )

    training_loop = training.Loop(
        NER,
        train_task,
        eval_tasks=[eval_task],
        output_dir=output_dir)

    if verbose:
        print(f"Running {train_steps} steps")
    training_loop.run(n_steps = train_steps)
    return training_loop

training_steps = 1500
real_stdout = sys.stdout

TIMER.emit = False
TIMER.start()
with TemporaryFile("w") as temp_file:
    sys.stdout = temp_file
    training_loop = train_model(trainee.model, train_generator,
                                evaluate_generator,
                                steps_per_checkpoint=10,
                                train_steps=training_steps,
                                verbose=False)
TIMER.stop()
sys.stdout = real_stdout

print(f"{TIMER.ended - TIMER.started}")

# pypi
from trax import layers

# this project
from neurotic.nlp.named_entity_recognition import DataGenerator, NERData, TOKEN

ner = NERData()

vocab = vocabulary = ner.data.vocabulary
tag_map = tags = ner.data.tags

model = NER()
# display your model
print(model)

from neurotic.nlp.named_entity_recognition import NER

builder = NER(122, 666)

print(builder.model)

# from python
import random

# from pypi
import numpy

# this project
from neurotic.nlp.named_entity_recognition import NERData, TOKEN

ner = NERData()

# to make the functions pass we need to use their names (initially)
vocab = vocabulary = ner.data.vocabulary
tag_map = tags = ner.data.tags

print(vocabulary["the"])
print(vocabulary[TOKEN.pad])
print(vocabulary["The"])

for tag, index in tags.items():
    print(f" - {tag}: {index}")

print(f"The number of tags is {len(tag_map)}")
print(f"The vocabulary size is {len(vocab):,}")
print(f"The training size is {len(ner.data.data_sets.x_train):,}")
print(f"The validation size is {len(ner.data.data_sets.x_validate):,}")
print("The first training sentence is ")
print(f"'{' '.join(ner.data.raw_data_sets.x_train[0])}'")
print("Its corresponding label is")
print(f" '{' '.join(ner.data.raw_data_sets.y_train[0])}'")

print("The first training encoded sentence is ")
print(f"{ner.data.data_sets.x_train[0]}")
print("Its corresponding encoded label is")
print(f"{ner.data.data_sets.y_train[0]}")

def data_generator(batch_size: int, x: list, y: list, pad: int,
                   shuffle: bool=False, verbose: bool=False):
    """Generate batches of data for training

    Args: 
      batch_size - size of each batch generated
      x - sentences where words are represented as integers
      y - tags associated with the sentences
      pad - number to use as the padding character
      shuffle - Whether to shuffle the data
      verbose - Whether to print information to stdout

    Yields:
     a tuple containing 2 elements:
       X - np.ndarray of dim (batch_size, max_len) of padded sentences
       Y - np.ndarray of dim (batch_size, max_len) of tags associated with the sentences in X
    """    
    # count the number of lines in data_lines
    num_lines = len(x)

    # create an array with the indexes of data_lines that can be shuffled
    lines_index = list(range(num_lines))

    # shuffle the indexes if shuffle is set to True
    if shuffle:
        random.shuffle(lines_index)

    index = 0 # tracks current location in x, y
    while True:
        buffer_x = [0] * batch_size
        buffer_y = [0] * batch_size
        max_len = 0
        for i in range(batch_size):
             # if the index is greater than or equal to the number of lines in x
            if index >= num_lines:
                # then reset the index to 0
                index = 0
                # re-shuffle the indexes if shuffle is set to True
                if shuffle:
                    random.shuffle(lines_index)

            # The current position is obtained using `lines_index[index]`
            # Store the x value at the current position into the buffer_x
            buffer_x[i] = x[lines_index[index]]

            # Store the y value at the current position into the buffer_y
            buffer_y[i] = y[lines_index[index]]

            lenx = len(buffer_x[i])    #length of current x[]
            if lenx > max_len:
                max_len = lenx                   #max_len tracks longest x[]

            # increment index by one
            index += 1


        # create X,Y, NumPy arrays of size (batch_size, max_len) 'full' of pad value
        X = numpy.full((batch_size, max_len), pad)
        Y = numpy.full((batch_size, max_len), pad)

        # copy values from lists to NumPy arrays. Use the buffered values
        for i in range(batch_size):
            # get the example (sentence as a tensor)
            # in `buffer_x` at the `i` index
            x_i = buffer_x[i]

            # similarly, get the example's labels
            # in `buffer_y` at the `i` index
            y_i = buffer_y[i]

            # Walk through each word in x_i
            for j in range(len(x_i)):
                # store the word in x_i at position j into X
                X[i, j] = x_i[j]

                # store the label in y_i at position j into Y
                Y[i, j] = y_i[j]

        if verbose: print("index=", index)
        yield((X,Y))

batch_size = 5
mini_sentences = ner.data.data_sets.x_train[0: 8]
mini_labels = ner.data.data_sets.y_train[0: 8]
dg = data_generator(batch_size, mini_sentences, mini_labels, vocab["<PAD>"], shuffle=False, verbose=True)
X1, Y1 = next(dg)
X2, Y2 = next(dg)
print(Y1.shape, X1.shape, Y2.shape, X2.shape)
print(X1[0][:], "\n", Y1[0][:])

# from python
from typing import List, Tuple
import random

# from pypi
import attr
import numpy

Vectors = List[List[int]]
Batch = Tuple[numpy.ndarray]

@attr.s(auto_attribs=True)
class DataGenerator:
    """A generator of data to train the NER Model

    Args:
     batch_size: how many lines to generate at once
     x: the encoded sentences
     y: the encoded labels 
     padding: encoding to use for padding lines
     shuffle: whether to shuffle the data
     verbose: whether to print messages to stdout
    """
    batch_size: int
    x: Vectors
    y: Vectors
    padding: int
    shuffle: bool=False
    verbose: bool=False
    _batch: iter=None

from neurotic.nlp.named_entity_recognition import DataGenerator

generator = DataGenerator(x=ner.data.data_sets.x_train[0:8],
                          y=ner.data.data_sets.y_train[0: 8],
                          batch_size=5,
                          padding=vocabulary[TOKEN.pad])

X1, Y1 = next(generator)
X2, Y2 = next(generator)
print(Y1.shape, X1.shape, Y2.shape, X2.shape)
print(X1[0][:], "\n", Y1[0][:])