# Importăm bibliotecile necesare
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow import keras
from sklearn.model_selection import train_test_split
import re
import random
import chardet

# Pregătim datele pentru modelul Transformer
# Ne asigurăm că tokenizer-ul nu filtrează caractere speciale sau simboluri importante.
tokenizer = keras.preprocessing.text.Tokenizer(filters='', oov_token='<OOV>')
tokenizer.fit_on_texts(voynich_sentences + english_sentences)

# Convertim textele în secvențe de tokeni
voynich_sequences = tokenizer.texts_to_sequences(voynich_sentences)
english_sequences = tokenizer.texts_to_sequences(english_sentences)

# Calculăm lungimea maximă a secvenței pentru padding
max_length = max(max(len(seq) for seq in voynich_sequences),
                 max(len(seq) for seq in english_sequences))

# Aplicăm padding pentru a obține secvențe de lungime fixă
voynich_padded = keras.preprocessing.sequence.pad_sequences(voynich_sequences, maxlen=max_length, padding='post')
english_padded = keras.preprocessing.sequence.pad_sequences(english_sequences, maxlen=max_length, padding='post')

# Verificăm dimensiunea vocabularului și lungimea maximă
vocab_size = len(tokenizer.word_index) + 1  # +1 pentru tokenul <OOV>
print(f"Dimensiunea vocabularului: {vocab_size}")
print(f"Lungimea maximă a secvenței: {max_length}")

# Împărțim datele în seturi de antrenament și validare
X_train, X_val, y_train, y_val = train_test_split(voynich_padded, english_padded, test_size=0.2, random_state=42)

# Reshape pentru a fi compatibil cu sparse_categorical_crossentropy
y_train = np.expand_dims(y_train, axis=-1)
y_val = np.expand_dims(y_val, axis=-1)

# Verificăm formele seturilor de date
print(f"Forma X_train: {X_train.shape}")
print(f"Forma y_train: {y_train.shape}")
print(f"Forma X_val: {X_val.shape}")
print(f"Forma y_val: {y_val.shape}")

# Funcție pentru blocul Transformer Encoder
def transformer_encoder(inputs, head_size, num_heads, ff_dim, dropout=0):
    # Atenție multi-cap
    x = keras.layers.MultiHeadAttention(
        key_dim=head_size, num_heads=num_heads, dropout=dropout
    )(inputs, inputs)
    x = keras.layers.Dropout(dropout)(x)
    x = keras.layers.LayerNormalization(epsilon=1e-6)(x)
    res = x + inputs  # Reziduul adăugat
    
    # Rețea feed-forward
    x = keras.layers.Conv1D(filters=ff_dim, kernel_size=1, activation="relu")(res)
    x = keras.layers.Dropout(dropout)(x)
    x = keras.layers.Conv1D(filters=inputs.shape[-1], kernel_size=1)(x)
    x = keras.layers.LayerNormalization(epsilon=1e-6)(x)
    return x + res  # Reziduul adăugat din nou

# Definim modelul Transformer
def transformer_model(vocab_size, max_length):
    inputs = keras.layers.Input(shape=(max_length,))
    
    # Adăugăm un strat de embedding
    embedding_layer = keras.layers.Embedding(vocab_size, 128)(inputs)
    x = embedding_layer
    
    # Adăugăm două blocuri Transformer
    for _ in range(2):
        x = transformer_encoder(x, head_size=64, num_heads=2, ff_dim=32, dropout=0.1)
    
    # Strat TimeDistributed pentru a prezice pe fiecare pas de timp
    x = keras.layers.TimeDistributed(keras.layers.Dense(vocab_size, activation="softmax"))(x)
    
    # Definim modelul final
    return keras.Model(inputs=inputs, outputs=x)

# Construim și compilăm modelul
model = transformer_model(vocab_size, max_length)
model.compile(optimizer=keras.optimizers.Adam(learning_rate=0.001),
              loss="sparse_categorical_crossentropy", 
              metrics=["accuracy"])

# Afișăm rezumatul modelului
print("Rezumatul modelului:")
model.summary()

# Antrenăm modelul
history = model.fit(X_train, y_train, batch_size=32, epochs=50, validation_data=(X_val, y_val))

# Funcție pentru a traduce text Voynich
def translate_voynich(text):
    sequence = tokenizer.texts_to_sequences([text])
    padded = keras.preprocessing.sequence.pad_sequences(sequence, maxlen=max_length, padding='post')
    predicted = model.predict(padded)
    predicted_sequence = np.argmax(predicted, axis=-1)[0]
    translated = ' '.join([tokenizer.index_word.get(i, '') for i in predicted_sequence if i != 0])
    return translated

# Exemplu de utilizare
sample_voynich = voynich_sentences[0]
print("Text Voynich:", sample_voynich)
print("Traducere:", translate_voynich(sample_voynich))

# Salvăm modelul pentru utilizare ulterioară
model.save('voynich_translator.h5')

# Generăm traducerea pentru întregul manuscris
full_translation = [translate_voynich(sentence) for sentence in voynich_sentences]

# Salvăm traducerea într-un fișier
with open('voynich_translation.txt', 'w', encoding='utf-8') as f:
    for original, translation in zip(voynich_sentences, full_translation):
        f.write(f"Original: {original}\n")
        f.write(f"Traducere: {translation}\n\n")

print("Traducerea completă a fost salvată în fișierul 'voynich_translation.txt'")