import tensorflow as tf
import io
import os
import numpy as np

# 1. Download dataset (eng–fra)
path_to_zip = tf.keras.utils.get_file(
    fname="fra-eng.zip",
    origin="http://storage.googleapis.com/download.tensorflow.org/data/fra-eng.zip",
    extract=True,
)
path_to_file = os.path.join(os.path.dirname(path_to_zip), "fra.txt")

# 2. Load sentence pairs
with io.open(path_to_file, encoding="UTF-8") as f:
    lines = f.read().strip().split("\n")

pairs = [l.split("\t")[:2] for l in lines]  # [eng, fra]
eng_texts = [p[0] for p in pairs]
fra_texts = [p[1] for p in pairs]

# Add special tokens to target (French)
start_token = "<sos>"
end_token = "<eos>"
fra_texts_in = [start_token + " " + t for t in fra_texts]
fra_texts_out = [t + " " + end_token for t in fra_texts]

# 3. Build tokenizers
num_words_src = 10000
num_words_tgt = 10000

src_tokenizer = tf.keras.preprocessing.text.Tokenizer(
    num_words=num_words_src, filters="", lower=True, oov_token="<unk>"
)
tgt_tokenizer = tf.keras.preprocessing.text.Tokenizer(
    num_words=num_words_tgt, filters="", lower=True, oov_token="<unk>"
)

src_tokenizer.fit_on_texts(eng_texts)
tgt_tokenizer.fit_on_texts(fra_texts_in + fra_texts_out)

src_sequences = src_tokenizer.texts_to_sequences(eng_texts)
tgt_sequences_in = tgt_tokenizer.texts_to_sequences(fra_texts_in)
tgt_sequences_out = tgt_tokenizer.texts_to_sequences(fra_texts_out)

# 4. Pad sequences
max_len_src = max(len(s) for s in src_sequences)
max_len_tgt = max(len(s) for s in tgt_sequences_in)

src_input = tf.keras.preprocessing.sequence.pad_sequences(
    src_sequences, maxlen=max_len_src, padding="post"
)
tgt_input = tf.keras.preprocessing.sequence.pad_sequences(
    tgt_sequences_in, maxlen=max_len_tgt, padding="post"
)
tgt_output = tf.keras.preprocessing.sequence.pad_sequences(
    tgt_sequences_out, maxlen=max_len_tgt, padding="post"
)

# 5. Define encoder–decoder model (from scratch)
embed_dim = 256
enc_units = 256
dec_units = 256
src_vocab_size = min(num_words_src, len(src_tokenizer.word_index) + 1)
tgt_vocab_size = min(num_words_tgt, len(tgt_tokenizer.word_index) + 1)

# Encoder
encoder_inputs = tf.keras.Input(shape=(None,), name="encoder_inputs")
enc_emb = tf.keras.layers.Embedding(src_vocab_size, embed_dim)(encoder_inputs)
encoder_lstm = tf.keras.layers.LSTM(
    enc_units, return_state=True, return_sequences=False
)
_, state_h, state_c = encoder_lstm(enc_emb)
encoder_states = [state_h, state_c]

# Decoder
decoder_inputs = tf.keras.Input(shape=(None,), name="decoder_inputs")
dec_emb_layer = tf.keras.layers.Embedding(tgt_vocab_size, embed_dim)
dec_emb = dec_emb_layer(decoder_inputs)
decoder_lstm = tf.keras.layers.LSTM(
    dec_units, return_sequences=True, return_state=True
)
decoder_outputs, _, _ = decoder_lstm(dec_emb, initial_state=encoder_states)
decoder_dense = tf.keras.layers.Dense(tgt_vocab_size, activation="softmax")
decoder_outputs = decoder_dense(decoder_outputs)

model = tf.keras.Model([encoder_inputs, decoder_inputs], decoder_outputs)
model.compile(
    optimizer="adam",
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    metrics=["accuracy"],
)

# 6. Train model
batch_size = 64
epochs = 10

model.fit(
    [src_input, tgt_input],
    np.expand_dims(tgt_output, -1),
    batch_size=batch_size,
    epochs=epochs,
    validation_split=0.1,
)

# 7. Save everything
os.makedirs("artifacts", exist_ok=True)
model.save("artifacts/eng_fra_seq2seq.h5")

# Save tokenizers and metadata
import pickle

with open("artifacts/src_tokenizer.pkl", "wb") as f:
    pickle.dump(src_tokenizer, f)
with open("artifacts/tgt_tokenizer.pkl", "wb") as f:
    pickle.dump(tgt_tokenizer, f)

meta = {
    "max_len_src": max_len_src,
    "max_len_tgt": max_len_tgt,
    "start_token": start_token,
    "end_token": end_token,
}
with open("artifacts/meta.pkl", "wb") as f:
    pickle.dump(meta, f)