import tensorflow as tf
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.layers import Embedding, LSTM, Dense, Bidirectional, Dropout
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping
import numpy as np
import re

# --------------------
# 1. Preprocessing for Phoneme Sequences
# --------------------
def preprocess_line(line):
    line = re.sub(r"<\/?s>", "", line)      # remove <s> </s>
    line = re.sub(r"\(\d+\)", "", line)     # remove (IDs)
    line = line.replace("_", "")            # remove underscores
    line = line.replace("7", "h")           # replace 7 with h
    return line.strip()

# --------------------
# 2. Load Phones Inventory (optional, for validation/debug)
# --------------------
with open("Phones.txt", encoding="utf-8") as f:
    phones = [p.strip() for p in f if p.strip()]
print(f"Loaded {len(phones)} phones from inventory")

# --------------------
# 3. Load & Clean Corpus
# --------------------
with open("Transcription-ROMAN.txt", encoding="utf-8") as f:
    raw_lines = [preprocess_line(line) for line in f if line.strip()]

# Keep only phoneme tokens (space separated)
corpus = []
for line in raw_lines:
    tokens = line.split()
    if tokens:
        corpus.append(" ".join(tokens))

print(f"Corpus size: {len(corpus)} lines")
print("Example preprocessed line:", corpus[0][:200])

# --------------------
# 4. Tokenizer (fit on corpus instead of only phones)
# --------------------
tokenizer = Tokenizer(num_words=10000, oov_token="<OOV>")

tokenizer.fit_on_texts(corpus)

total_vocab = len(tokenizer.word_index) + 1
print("Total vocab size:", total_vocab)

# --------------------
# 5. Create Training Sequences
# --------------------
input_sequences = []
for line in corpus:
    token_list = tokenizer.texts_to_sequences([line])[0]
    for i in range(1, len(token_list)):
        n_gram_seq = token_list[:i+1]
        input_sequences.append(n_gram_seq)

max_sequence_len = max(len(x) for x in input_sequences)
print("Max sequence length:", max_sequence_len)

input_sequences = np.array(
    pad_sequences(input_sequences, maxlen=max_sequence_len, padding="pre")
)

xs, labels = input_sequences[:, :-1], input_sequences[:, -1]  # no one-hot, sparse labels

# --------------------
# 6. Model Definition
# --------------------
model = Sequential([
    Embedding(total_vocab, 128),
    Bidirectional(LSTM(256, return_sequences=True)),
    Dropout(0.3),
    LSTM(128),
    Dense(256, activation="relu"),
    Dropout(0.3),
    Dense(total_vocab, activation="softmax")
])

adam = Adam(learning_rate=0.0005)
model.compile(loss="sparse_categorical_crossentropy",
              optimizer=adam,
              metrics=["accuracy"])

# Build before summary
model.build(input_shape=(None, max_sequence_len-1))
print(model.summary())

# --------------------
# 7. Train with EarlyStopping
# --------------------
early_stop = EarlyStopping(monitor="val_loss", patience=7, restore_best_weights=True)

history = model.fit(
    xs, labels,
    epochs=50,
    batch_size=128,
    validation_split=0.1,
    shuffle=True,
    callbacks=[early_stop],
    verbose=1
)

# --------------------
# 8. Sampling & Text Generation
# --------------------
def sample_with_temp(preds, temperature=1.0):
    preds = np.asarray(preds).astype("float64")
    preds = np.log(preds + 1e-8) / temperature
    exp_preds = np.exp(preds)
    preds = exp_preds / np.sum(exp_preds)
    return np.random.choice(len(preds), p=preds)

def generate_sequence(seed_text, next_tokens=30, temperature=0.8):
    """Generate a phoneme sequence from a seed"""
    for _ in range(next_tokens):
        token_list = tokenizer.texts_to_sequences([seed_text])[0]
        token_list = pad_sequences([token_list], maxlen=max_sequence_len-1, padding="pre")
        preds = model.predict(token_list, verbose=0)[0]
        predicted = sample_with_temp(preds, temperature)
        output_token = tokenizer.index_word.get(predicted, "")
        seed_text += " " + output_token
    return seed_text

def reconstruct_words(sequence: str) -> str:
    """Join phoneme tokens back into Roman Urdu words"""
    tokens = sequence.split()
    return " ".join(tokens).replace("  ", " ").strip()

# --------------------
# 9. Example Usage
# --------------------
seed = "zindagi ek"
raw_output = generate_sequence(seed, next_tokens=7, temperature=0.7)
print("Generated phoneme sequence:", raw_output)

roman_urdu = reconstruct_words(raw_output)
print("Reconstructed Roman Urdu:", roman_urdu)