import pickle
import tensorflow as tf
import pandas as pd
import re


# Define the function to preprocess the input text
def preprocess_text(text):
    # Convert the text to lowercase
    text = text.lower()
    # Replace any non-alphanumeric characters with a space
    text = re.sub(r'[^a-zA-Z0-9\'+\-*/\s]', ' ', text)
    # Replace the word 'plus' with the '+' symbol
    text = text.replace('plus', '+')
    # Replace the word 'minus' with the '-' symbol
    text = text.replace('minus', '-')
    # Replace the word 'times' with the '*' symbol
    text = text.replace('times', '*')
    # Replace the word 'divided by' with the '/' symbol
    text = text.replace('divided by', '/')
    # Replace the characters * / - + with a space and the character itself
    text = text.replace('*', ' * ').replace('/', ' / ').replace('-', ' - ').replace('+', ' + ').replace('.', ' . ')
    # Replace all numbers with a special token
    text = re.sub(r'\d+', '<num>', text)
    # Replace all numbers with a special token
    text = text.replace(r'oov', '<oov>')
    # Remove any extra whitespace
    text = re.sub(r'\s+', ' ', text)
    return text.strip()


epochs = int(input("epochs: "))
batch_size = int(input("batch size: "))

# Load the training data from a CSV file
data = pd.read_csv("data.csv")
input_data = data['input'].values
labels = data['class'].values

# Preprocess the input data
input_data = [preprocess_text(text) for text in input_data]

# Define the vocabulary and encode the input data
vocab = tf.keras.preprocessing.text.Tokenizer(filters='')
vocab.fit_on_texts(input_data)
encoded_input = vocab.texts_to_sequences(input_data)

# Pad the encoded input data to ensure all inputs are of the same length
max_length = max([len(seq) for seq in encoded_input])
padded_input = tf.keras.preprocessing.sequence.pad_sequences(encoded_input, maxlen=max_length, padding='post')

# Define the neural network model
model = tf.keras.Sequential([
    tf.keras.layers.Embedding(input_dim=len(vocab.word_index) + 1, output_dim=64, input_length=max_length),
    tf.keras.layers.Conv1D(filters=64, kernel_size=3, activation='relu', padding='same'),
    tf.keras.layers.GlobalMaxPooling1D(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(128, activation=tf.keras.layers.ELU(alpha=1.0)),
    tf.keras.layers.Dense(3, activation='softmax')
])

# Compile the model
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Train the model
model.fit(padded_input, labels, epochs=epochs, batch_size=batch_size)

# Save the trained model
model.save('main.h5')

# Save the tokenizer using pickle
with open('tokenizer.pkl', 'wb') as file:
    pickle.dump(vocab, file)

token_to_word = {token: word for word, token in vocab.word_index.items()}
print(token_to_word)

model.summary()

# Enter into an interactive loop to test the model