import tensorflow as tf
keras = tf.keras
from tensorflow import Tensor
from keras.api.layers import TextVectorization, Embedding, LSTM, Dense, Input
from keras.api.models import Sequential
import numpy as np
import numpy.typing as npt
import pandas as pd
from pandas import DataFrame



#
#
#  Description: The purpose of this model is to qualify sentences into different 'functions'
#
#    - "Utility" - A sentence whose purpose is to perform some non-informational duty in a conversation.
#                  An example of this could be a greeting or farewell.
#
#    - "Transfer" - Transfers a piece of information to another participant. For instance, describing
#                   the details of an object.
#
#    - "Query" - Queries a piece of information from another participant.
#
#    - "Imperative" - Commands a participant to do something.
#



#
#
# MODEL CONSTANTS
#
#

# this is the maximum allowed size of the vocabulary
max_features: int = 20000

# the dimension of the output from the embedding layer
embedding_dim: int = 128

# The number of epochs to train for
epochs: int = 50

# Maximum size of the vocab for this layer
max_tokens: int = 5000

# (Only valid in INT mode) If set, the output will have its time dimension padded or truncated to exactly output_sequence_length values
output_sequence_length: int = 4

# The number of classes we're training for
num_classes: int = 4




#
#
# LOAD DATA
#
#


# read training sentences
train_csv_path: str = './data/semantic/sent_func/train.csv'
train_csv_raw: DataFrame = pd.read_csv(open(train_csv_path), dtype={
    "utility": "int64",
    "transfer": "int64",
    "query": "int64",
    "imperative": "int64",
    "sentence": "string",
})
train_data_raw: DataFrame = train_csv_raw[["sentence"]]
train_data_split: DataFrame = train_data_raw
vocab: list[str] = train_data_raw.to_numpy()
train_labels: DataFrame = train_csv_raw[["utility", "transfer", "query", "imperative"]]

# read testing sentences
test_csv_path: str = './data/semantic/sent_func/test.csv'
test_csv_raw: DataFrame = pd.read_csv(open(test_csv_path), dtype={
    "utility": "int64",
    "transfer": "int64",
    "query": "int64",
    "imperative": "int64",
    "sentence": "string",
})
test_data_raw: DataFrame = test_csv_raw[["sentence"]]
test_data_split: DataFrame = test_data_raw
test_labels: DataFrame = test_csv_raw[["utility", "transfer", "query", "imperative"]]



#
#
# CREATE VECTORIZER
#
#


# init vectorizer
textVec: TextVectorization = TextVectorization(
    max_tokens=max_tokens,
    output_mode='int',
    output_sequence_length=output_sequence_length,
    pad_to_max_tokens=True
)

# Add the vocab to the tokenizer
textVec.adapt(vocab) 
input_data: list[str] = train_data_split
train_data: Tensor = textVec.call(input_data)



#
#
# CREATE MODEL
#
#


# construct model
model: Sequential = Sequential([
    keras.Input(shape=(1,), dtype=tf.string),
    textVec,
    Embedding(max_features + 1, embedding_dim),
    LSTM(128),
    Dense(num_classes, activation='sigmoid')
])

#compile the model
# model.build(keras.Input(shape=(None,), dtype="int64"))
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])



#
#
# TRAIN MODEL
#
#

# Final formatting of data
npTrainData = train_data_split.to_numpy(dtype=object).flatten()
npTrainLabel: npt.NDArray[np.complex64] = train_labels.to_numpy()
if(npTrainData.shape[0] != npTrainLabel.shape[0]):
    print("Data label size mismatch!")
    print(npTrainData.shape)
    print(npTrainLabel.shape)

# try fitting data
print("Training..")
model.fit(npTrainData,npTrainLabel,epochs=epochs)



#
#
# EVALUATE MODEL
#
#


# evaluate here
npTestData: npt.NDArray = test_data_split.to_numpy(dtype=object).flatten()
npTestLabel: npt.NDArray[np.complex64] = test_labels.to_numpy()
if(npTestData.shape[0] != npTestLabel.shape[0]):
    print("Data label size mismatch!")
    print(npTestData.shape)
    print(npTestLabel.shape)

print("Evaluating..")
model.evaluate(npTestData,npTestLabel)



#
#
# PREDICT
#
#

# predictTargetRaw: Tensor = tf.constant(['Hello'])
# npPredict: npt.NDArray = np.array(predictTargetRaw, dtype=object)
# print("Prediction test..")
# result: list[int] = model.predict(predictTargetRaw)
# print("predict result:")
# print(predictTargetRaw)
# print(result)
# print(train_data)
# print(train_labels)



#
#
# SAVE (DEVELOPMENT) MODEL
#
#

# save the model so keras can reload
# savePath: str = './data/semantic/model.keras'
# model.save(savePath)


#
#
# SAVE (PRODUCTION) MODEL
#
#

# export the model so java can leverage it
print("Saving..")
exportPath: str = './data/model/sent_func'
model.export(exportPath)

