import os
import sys
import json

# ==============================================================================
# AUTO-CONFIG: Find and link hidden XLA compiler libraries inside the venv
# ==============================================================================
def configure_xla_paths():
    venv_base = sys.prefix
    for root, _, files in os.walk(venv_base):
        if "libdevice.10.bc" in files:
            cuda_dir = root.split("/nvvm")[0]
            os.environ["XLA_FLAGS"] = f"--xla_gpu_cuda_data_dir={cuda_dir}"
            print(f"[XLA Config] Successfully linked compiler data to: {cuda_dir}")
            return True
    print("[XLA Config] Warning: libdevice.10.bc not found. Training may crash.")
    return False

configure_xla_paths()

import tensorflow as tf
#tf.debugging.set_log_device_placement(True)
import numpy as np

# ==========================================
# 1. ARCHITECTURE CONFIGURATION
# ==========================================
EMBEDDING_DIM = 64
RNN_UNITS = 256

# Pipeline & Training Hyperparameters
TEXT_FILE = "combined_training_data.txt"
VOCAB_FILE = "vocab.txt"

# DYNAMICALLY MANGLED DIRECTORY:
# This will automatically evaluate to something like: ./checkpoints_emb256_rnn512
CHECKPOINT_DIR = f"./checkpoints_emb{EMBEDDING_DIM}_rnn{RNN_UNITS}"
CONFIG_FILE = os.path.join(CHECKPOINT_DIR, "config.json")

SEQ_LENGTH = 128
BATCH_SIZE = 128
EPOCHS = 100
BUFFER_SIZE = 10000
SEED_TEXT = "it was a dark and stormy night"

print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
os.makedirs(CHECKPOINT_DIR, exist_ok=True)

# ==========================================
# 2. CONFIG CHECKPOINT GUARD (THE MAGIC SOUP)
# ==========================================
# If an existing config file is found, OVERWRITE your local variables
# to force the model to build at the size matching the saved weights.
if os.path.exists(CONFIG_FILE):
    print(f"Found existing model configuration file at '{CONFIG_FILE}'.")
    with open(CONFIG_FILE, 'r') as f:
        saved_config = json.load(f)

    EMBEDDING_DIM = saved_config["embedding_dim"]
    RNN_UNITS = saved_config["rnn_units"]
    print(f"-> Overrode network sizes to match saved profile: Embedding={EMBEDDING_DIM}, RNN Units={RNN_UNITS}")
else:
    # Save the current configuration since it's a brand new run
    current_config = {
        "embedding_dim": EMBEDDING_DIM,
        "rnn_units": RNN_UNITS
    }
    with open(CONFIG_FILE, 'w') as f:
        json.dump(current_config, f, indent=4)
    print(f"Created a new model configuration footprint file at '{CONFIG_FILE}'.")

# ==========================================
# 3. DATA PREPROCESSING
# ==========================================
if os.path.exists(VOCAB_FILE):
    print(f"Loading existing vocabulary file from '{VOCAB_FILE}'...")
    with open(VOCAB_FILE, 'r', encoding='utf-8') as f:
        vocab = json.load(f)
else:
    if not os.path.exists(TEXT_FILE):
        print(f"File '{TEXT_FILE}' not found locally. Please run the process.py script!")
        quit(1)

    print(f"Loading raw text from {TEXT_FILE} to build vocabulary...")
    with open(TEXT_FILE, 'r', encoding='utf-8') as f:
        text = f.read()
    temp_vectorizer = tf.keras.layers.TextVectorization(split="character", standardize="lower")
    temp_vectorizer.adapt(tf.data.Dataset.from_tensor_slices([text]))
    vocab = temp_vectorizer.get_vocabulary()
    with open(VOCAB_FILE, 'w', encoding='utf-8') as f:
        json.dump(vocab, f, ensure_ascii=False)

vocab_size = len(vocab)
vectorize_layer = tf.keras.layers.TextVectorization(
    split="character", standardize="lower", vocabulary=vocab, output_mode="int"
)

dataset = None
if os.path.exists(TEXT_FILE):
    with open(TEXT_FILE, 'r', encoding='utf-8') as f:
        text = f.read()
    all_ids = vectorize_layer(tf.constant([text]))[0]
    ids_dataset = tf.data.Dataset.from_tensor_slices(all_ids)
    sequences = ids_dataset.batch(SEQ_LENGTH + 1, drop_remainder=True)

    def split_input_target(sequence):
        return sequence[:-1], sequence[1:]

    dataset = sequences.map(split_input_target).shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True).prefetch(tf.data.AUTOTUNE)

# ==========================================
# 4. MODEL ARCHITECTURE
# ==========================================

class CharacterTextModel(tf.keras.Model):
    def __init__(self, vocab_size, embedding_dim, rnn_units, num_layers=2, dropout_rate=0.4):
        super().__init__()
        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)

        # Create paired lists of GRUs and Dropouts based on your desired depth
        self.gru_layers = []
        self.dropout_layers = []

        # Loop runs (num_layers - 1) times, leaving the final layer separate
        for _ in range(num_layers - 1):
            self.gru_layers.append(
                tf.keras.layers.GRU(rnn_units, return_sequences=True)
            )
            self.dropout_layers.append(
                tf.keras.layers.Dropout(dropout_rate)
            )

        # Switched to GRU. Note: GRU only has 1 state output (h), not 2 (h, c)
        self.gru_final = tf.keras.layers.GRU(rnn_units, return_sequences=True, return_state=True)
        self.dropout_final = tf.keras.layers.Dropout(dropout_rate)
        self.dense = tf.keras.layers.Dense(vocab_size)

    def call(self, inputs, states=None, return_state=False, training=False):
        x = self.embedding(inputs, training=training)

        # Loop through both lists simultaneously using zip()
        for gru, dropout in zip(self.gru_layers, self.dropout_layers):
            x = gru(x, training=training)
            x = dropout(x, training=training)

        # Final GRU Layer returns the output 'x' and a single hidden state 'h'
        x, h = self.gru_final(x, initial_state=states, training=training)
        x = self.dropout_final(x, training=training)
        x = self.dense(x, training=training)

        if return_state:
            return x, h  # Returns just 'h' instead of '(h, c)'
        return x

# Automatically scales to whatever dimensions were chosen or loaded!
model = CharacterTextModel(vocab_size=vocab_size, embedding_dim=EMBEDDING_DIM, rnn_units=RNN_UNITS)
model.build(input_shape=(BATCH_SIZE, SEQ_LENGTH))

import glob

# Find all files matching the pattern
checkpoint_files = glob.glob(os.path.join(CHECKPOINT_DIR, "ckpt_*.weights.h5"))

if checkpoint_files:
    # Sort files naturally or by modification time to get the latest one
    latest_checkpoint = max(checkpoint_files, key=os.path.getmtime)
    print(f"Restoring model layers from: {latest_checkpoint}")
    model.load_weights(latest_checkpoint)
else:
    print("Starting a clean initialization.")


loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer='adam', loss=loss)

# ==========================================
# 5. SAMPLING & GENERATION LOGIC
# ==========================================
def produce_sample(model, seed, num_generate=300, temperature=0.7):
    input_chars = vectorize_layer(tf.constant([seed]))
    input_ids = input_chars[0][:len(seed)].numpy().tolist()
    generated_ids, states = [], None

    for _ in range(num_generate):
        current_tokens = input_ids[-SEQ_LENGTH:]
        predictions, states = model(tf.expand_dims(current_tokens, 0), states=states, return_state=True, training=False)
        predictions = predictions[0, -1, :] / temperature
        predicted_id = tf.random.categorical(tf.expand_dims(predictions, 0), num_samples=1)[0, 0].numpy()
        generated_ids.append(predicted_id)
        input_ids.append(predicted_id)

    return "".join([vocab[idx] for idx in generated_ids])

class GenerationCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        print(f"\n--- End of Epoch {epoch+1} ---")
        if epoch % 10 == 0:
            for temp in [0.1, 0.4, 0.7]:
                print(f"Seed: \"{SEED_TEXT}\" -> {produce_sample(self.model, seed=SEED_TEXT, num_generate=200, temperature=temp)}\n")

checkpoint_prefix = os.path.join(CHECKPOINT_DIR, "ckpt_{epoch}.weights.h5")
checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_prefix, save_weights_only=True)

# Stop training when validation loss stops improving for 3 epochs
early_stopping = tf.keras.callbacks.EarlyStopping(
    monitor='val_loss',
    patience=3,
    restore_best_weights=True # Automatically rolls back to epoch 15 weights!
)

# ==========================================
# 6. START RUN
# ==========================================
if dataset is not None:
    model.fit(dataset, epochs=EPOCHS, callbacks=[GenerationCallback(), checkpoint_callback])
else:
    print(f"\n--- Inference Mode ({EMBEDDING_DIM}dim, {RNN_UNITS}units) ---")
    print(f"Result: {produce_sample(model, seed=SEED_TEXT, num_generate=300, temperature=0.6)}")