import gradio as gr
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import pickle
import os

# Load artifacts
ARTIFACTS_DIR = "."
model = tf.keras.models.load_model(os.path.join(ARTIFACTS_DIR, "eng_fra_seq2seq.h5"))
with open(os.path.join(ARTIFACTS_DIR, "src_tokenizer.pkl"), "rb") as f:
    src_tokenizer = pickle.load(f)
with open(os.path.join(ARTIFACTS_DIR, "tgt_tokenizer.pkl"), "rb") as f:
    tgt_tokenizer = pickle.load(f)
with open(os.path.join(ARTIFACTS_DIR, "meta.pkl"), "rb") as f:
    meta = pickle.load(f)

max_len_src = meta["max_len_src"]
max_len_tgt = meta["max_len_tgt"]
start_token = meta["start_token"]
end_token = meta["end_token"]

tgt_index_to_word = {i: w for w, i in tgt_tokenizer.word_index.items()}
tgt_index_to_word[0] = ""  # padding

# Extract encoder and decoder pieces from full model
# Note: this matches the training architecture above.
encoder_inputs = model.input[0]
encoder_emb_output = model.layers[2](encoder_inputs)
encoder_lstm = model.layers[3]
encoder_outputs = encoder_lstm(encoder_emb_output)
state_h_enc = encoder_outputs[1]
state_c_enc = encoder_outputs[2]
encoder_model = tf.keras.Model(encoder_inputs, [state_h_enc, state_c_enc])


decoder_inputs = model.input[1]
decoder_inputs = layers.Lambda(
    lambda x: tf.expand_dims(x, axis=1)
)(decoder_inputs)
# Find layers by types in order of creation
dec_emb_layer = model.layers[4]
decoder_lstm = model.layers[5]
decoder_dense = model.layers[6]

dec_state_input_h = tf.keras.Input(shape=(decoder_lstm.units,))
dec_state_input_c = tf.keras.Input(shape=(decoder_lstm.units,))
dec_emb2 = dec_emb_layer(decoder_inputs)
decoder_lstm_outputs,state_h, state_c = decoder_lstm(dec_emb2, initial_state=(dec_state_input_h, dec_state_input_c))
dec_outputs = decoder_lstm_outputs[0]
state_h_dec = decoder_lstm_outputs[1]
state_c_dec = decoder_lstm_outputs[2]
dec_outputs = decoder_dense(dec_outputs)
decoder_model = tf.keras.Model(
    [decoder_inputs, dec_state_input_h, dec_state_input_c],
    [dec_outputs, state_h_dec, state_c_dec],
)

def translate_eng_to_fra(text: str) -> str:
    if not text.strip():
        return ""

    # Encode source
    seq = src_tokenizer.texts_to_sequences([text])
    seq = tf.keras.preprocessing.sequence.pad_sequences(
        seq, maxlen=max_len_src, padding="post"
    )
    state_h, state_c = encoder_model.predict(seq)

    # Start decoder with <sos>
    start_idx = tgt_tokenizer.word_index.get(start_token, 1)
    end_idx = tgt_tokenizer.word_index.get(end_token, 2)

    target_seq = np.array([[start_idx]])
    decoded_sentence = []
    for _ in range(max_len_tgt):
        output_tokens, h, c = decoder_model.predict(
            [target_seq, state_h, state_c]
        )

        sampled_token_index = np.argmax(output_tokens[0, -1, :])
        sampled_word = tgt_index_to_word.get(sampled_token_index, "")

        if sampled_token_index == end_idx or sampled_word == end_token:
            break

        if sampled_word not in ["", start_token]:
            decoded_sentence.append(sampled_word)

        target_seq = np.array([[sampled_token_index]])
        state_h, state_c = h, c

    return " ".join(decoded_sentence)

with gr.Blocks() as demo:
    gr.Markdown("# English → French (from scratch)")
    with gr.Row():
        inp = gr.Textbox(label="English Input")
        out = gr.Textbox(label="French Output")
    btn = gr.Button("Translate")
    btn.click(fn=translate_eng_to_fra, inputs=inp, outputs=out)

if __name__ == "__main__":
    demo.launch()