from tensorflow.python.keras.models import Sequential

import utils as utils
import tensorflow as tf
import numpy as np
from tensorflow import keras
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.text import Tokenizer
from typing import Tuple, List
from utils import ModelSaver, DATAFILES
from assemble_data import TSV_PATH

# -------- Data Parameters --------
TRAINING_SIZE = 160_000
"""Number of examples from the dataset to use in training"""
MODEL_PATH = 'savefiles/Classifier'
"""Path to save the model to and load it from"""
TOKENIZER_PATH = 'savefiles/classifier_tokenizer.json'
"""Path to save the tokenizer to and load it from"""
# ----------------


# ======== MODEL CLASSES ========
# ================


# -------- Tokenizing/Embedding parameters --------
VOCABULARY_SIZE = 10_000
"""Maximum number of words in tokenizer index"""
EMBEDDING_DIMENSIONS = 16
"""Dimensions in model's embedding layer"""
INPUT_LENGTH = 600
"""Maximum amount of tokens in model input. This will be the size of the 
input layer and shorter values will be padded to fit it."""
OOV_TOKEN = "<OOV>"
"""Token to print for out-of-vocabulary words when tokenizing."""
# ----------------

# -------- Training Parameters --------
MAX_EPOCHS = 5
"""Maximum amount of epochs to run the model training for."""
ACCURACY_THRESHOLD = 0.97
"""If the model passes this threshold in its accuracy on the validation dataset, it is saved and the training ends."""
CALLBACKS = [ModelSaver(ACCURACY_THRESHOLD, MODEL_PATH)]
"""Callback functions to run at the end of each input (objects must inherit from"""
# ----------------
# ================

# ======== MODEL DATA FUNCTIONS ========
def unpack_tsv() -> Tuple[List[str], List[List[int]]]:
    """
    Unpacks the dataset TSV generated by assemble_data.py into the list of paragraphs and the list of their corresponding
    label vectors

    :param filepath: the path the TSV file was generated to
    :return: A tuple of the list of training paragraphs and the list of their corresponding label vectors
    """
    raw_data = []
    labels = []

    with open(TSV_PATH) as tsvfile, open('datafiles/neither.txt') as neitherfile:
        for i, item in enumerate(tsvfile.readlines()):
            cells = item.split('\t')
            raw_data.append(cells[0])
            label = cells[1].split(',')
            label = [int(i) for i in label]
            labels.append(label)

            if i % 2 == 0:
                raw_data.append(neitherfile.readline())
                labels.append([0, 0])

        return raw_data, labels


def get_data() -> Tuple[Tuple[np.ndarray, np.array], Tuple[np.ndarray, np.array]]:
    """
    :return: A tuple of the training and testing datasets.
    """
    raw_data, labels = unpack_tsv()

    tokenizer = utils.generate_tokenizer(VOCABULARY_SIZE, raw_data[:TRAINING_SIZE], TOKENIZER_PATH)

    processed_data = utils.preprocess_text(raw_data, tokenizer, INPUT_LENGTH)
    processed_labels = utils.preprocess_labels(labels)

    training_data, training_labels = processed_data[:TRAINING_SIZE], processed_labels[:TRAINING_SIZE]
    testing_data, testing_labels = processed_data[TRAINING_SIZE:], processed_labels[TRAINING_SIZE:]
    return (training_data, training_labels), (testing_data, testing_labels)


def generate_tokenizer(texts_to_fit: List[str]) -> Tokenizer:
    """
    Generates a new tokenizer to use based on given data to fit it to, and saves it to the `TOKENIZER_PATH` set above.
    This should only run once each time the model is regenerated and retrained, when gathering the data to train it on.
    Otherwise, the tokenizer should be achieved via `get_tokenizer()`

    :param texts_to_fit: Texts to generate a vocabulary from
    :return: a new tokenizer based on the `VOCABULARY_SIZE` constant above, fit to the given texts.
    """
    return utils.generate_tokenizer(VOCABULARY_SIZE, texts_to_fit, TOKENIZER_PATH)


def get_tokenizer() -> Tokenizer:
    """
    :return: the tokenizer saved at `TOKENIZER_PATH`, set above.
    """
    return utils.load_tokenizer(TOKENIZER_PATH)


def predict(texts: List[str]) -> List[List[float]]:
    """
    Predicts the classes of textual inputs: ingredient, instruction, and junk.
    :param texts: A list of textual inputs to classify
    :return:
    """
    model = get_model()
    processed_data = utils.preprocess_text(texts, get_tokenizer(), INPUT_LENGTH)
    return model.predict(processed_data)
# ================


# ======== MODEL GENERATION FUNCTIONS ========
def generate_model() -> Sequential:
    """
    Generates a new classifier model, that maps `INPUT_LENGTH` input to 2 sigmoid-activated neurons indicating
    instruction and ingredient probabilities, and is compiled with MSE loss, adam optimizer and accuracy metric.
    see this function definition for more details.

    This should only run once each time the model is regenerated and retrained. Otherwise, the tokenizer
    should be achieved via `get_model()`

    :return: A new model by the classifier's architecture
    """
    model = tf.keras.Sequential([

        # Input layer. Size is INPUT_LENGTH; preprocessing converts each sentence into an INPUT_LENGTH list of
        # real tokens, padded if needed.
        # Uses word embedding, which converts sequences into tensors where each word is a vector,
        # indicating its position on a number (EMBEDDING_DIMENNSIONS) of semantic axes (gender, tense, verb/noun etc.)
        tf.keras.layers.Embedding(VOCABULARY_SIZE, EMBEDDING_DIMENSIONS, input_length=INPUT_LENGTH),
        # Noise/Dimension reduction layer into the average embedding in the sentence.
        tf.keras.layers.GlobalAveragePooling1D(),
        # Basic neural network black box, with standard selu activation
        tf.keras.layers.Dense(24, activation='selu'),
        tf.keras.layers.Dense(32, activation='selu'),
        tf.keras.layers.Dense(32, activation='selu'),
        tf.keras.layers.Dense(24, activation='selu'),
        # Sigmoid ensures the model outputs probabilities between 0 and 1. Output size is two,
        # for classifications as ingredient and as instruciton.
        tf.keras.layers.Dense(2, activation='sigmoid')
    ])

    # Notice that even though this is a classifier, the loss function is MSE, rather than the more common cross entropy.
    # The reason for this is that the model is also trained on "junk" data with expected value [0, 0]. This is because
    # cross entropy is used
    model.compile(loss=keras.losses.mean_squared_error, optimizer='adam', metrics=['accuracy'])

    return model


def get_model() -> Sequential:
    """
    :return: The model saved in the set `MODEL_PATH`
    """
    return load_model(MODEL_PATH)


def train_model():
    """
    Trains the model based on the TSV file generated by assemble_data and the neither.txt file generated by scrape_data.
    """
    model = generate_model()
    training_dataset, testing_dataset = get_data()
    model.fit(*training_dataset, validation_data=testing_dataset, epochs=MAX_EPOCHS, callbacks=CALLBACKS)
# ================


# ======== MODEL TESTING FUNCTIONS ========
def test_on_json():
    """
    Evaluates the model on the test data from JSON dataset from kraggle
    """
    model = get_model()
    _, (test_data, test_labels) = get_data()
    model.evaluate(test_data, np.array(test_labels))


def test_on_scraped():
    """
    Evaluates the model on the site data generated by scrape_data, and show the average guess for each actual class.
    """
    # DATAFILES = [ingredients,instructions, junk]
    # Labels matching data files.
    labels = [[1, 0], [0, 1], [0, 0]]
    # For each file
    for filename, label in zip(DATAFILES, labels):
        # Get the file data
        with open(filename, 'r') as datafile:
            test_data = datafile.readlines()
        # Generate valid datasets from it
        test_labels = np.array([label for _ in test_data])
        test_data = utils.preprocess_text(test_data, get_tokenizer(), INPUT_LENGTH)
        # Show model evaluations and average prediction
        print(filename)
        model.evaluate(test_data, test_labels)
        predictions = model.predict(test_data)
        print('AVERAGE: ' + str(np.mean(predictions, axis=0)))
# ================


# ======== RUNNING SCRIPT ========
if __name__ == '__main__':
    train_model()

    model = load_model(MODEL_PATH)
    test_on_json()
    test_on_scraped()
# ================
