import os
import json
import glob
import argparse
import sys
import re

# ==============================================================================
# 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
import numpy as np

# ==========================================
# 1. MODEL ARCHITECTURE
# ==========================================
class CharacterTextModel(tf.keras.Model):
    def __init__(self, vocab_size, embedding_dim, rnn_units, num_layers=1, dropout_rate=0.3):
        super().__init__()
        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)

        self.lstm_layers = [
            tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
            for _ in range(num_layers)
        ]
        self.dropout_layers = [tf.keras.layers.Dropout(dropout_rate) for _ in range(num_layers)]
        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)

        if states is None:
            states = [None] * len(self.lstm_layers)

        new_states = []
        for i, (lstm, dropout) in enumerate(zip(self.lstm_layers, self.dropout_layers)):
            x, h, c = lstm(x, initial_state=states[i], training=training)
            x = dropout(x, training=training)
            new_states.append([h, c])

        # FIXED: Moved out of the loop so multi-layer architectures function properly
        x = self.dense(x, training=training)

        if return_state:
            return x, new_states
        return x



# ==========================================
# 2. HELPER UTILITIES
# ==========================================
def load_or_create_vocab(text_file, vocab_file):
    if os.path.exists(vocab_file):
        print(f"[Vocab] Loading existing vocabulary from '{vocab_file}'...")
        with open(vocab_file, 'r', encoding='utf-8') as f:
            return json.load(f)

    if not os.path.exists(text_file):
        print(f"[Error] Text file '{text_file}' not found. Cannot build vocabulary.")
        quit(1)

    print(f"[Vocab] Building vocabulary from {text_file}...")
    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)
    return vocab

def produce_sample(model, vectorize_layer, vocab, 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

    predictions, states = model(tf.expand_dims(input_ids, 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)

    for _ in range(num_generate - 1):
        predictions, states = model(tf.expand_dims([predicted_id], 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)

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

class GenerationCallback(tf.keras.callbacks.Callback):
    def __init__(self, vectorize_layer, vocab, seed_text):
        super().__init__()
        self.vectorize_layer = vectorize_layer
        self.vocab = vocab
        self.seed_text = seed_text

    def on_epoch_end(self, epoch, logs=None):
        print(f"\n--- End of Epoch {epoch+1} ---")
        if epoch % 10 == 0:
            for temp in [0.3, 0.6, 0.9, 1.2]:
                sample = produce_sample(self.model, self.vectorize_layer, self.vocab, seed=self.seed_text, num_generate=350, temperature=temp)
                print(f"Seed (Temp {temp}): \"{self.seed_text}\" -> {sample}\n")







# ==========================================
# 3. MAIN EXECUTION PIPELINE
# ==========================================
def main():
    parser = argparse.ArgumentParser(description="Dynamically build, train, or infer from custom character LSTM models.")
    parser.add_argument("--modelfile", type=str, default=None, help="Path to a specific '.weights.h5' checkpoint file.")
    parser.add_argument("--prompt", type=str, default=None, help="Prompt string to run inference generation.")
    parser.add_argument("--train", action="store_true", help="Flag to trigger or resume training on a model.")
    args = parser.parse_args()

    # Default settings for clean, scratch runs
    DEFAULT_EMBEDDING_DIM = 512
    DEFAULT_RNN_UNITS = 512
    DEFAULT_NUM_LAYERS = 1
    DEFAULT_DROPOUT_RATE = 0.3

    TEXT_FILE = "emily_post.txt"
    VOCAB_FILE = "vocab.txt"
    SEQ_LENGTH = 512
    BATCH_SIZE = 128
    EPOCHS_PER_RUN = 200  # Easily run in chunks of 20 epochs
    BUFFER_SIZE = 10000
    SEED_TEXT = "In the case of unlawful enterance"

    vocab = load_or_create_vocab(TEXT_FILE, VOCAB_FILE)
    vocab_size = len(vocab)

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

    # Architectural parameters (will overwrite if loading an old model profile)
    emb_dim, rnn_units, num_layers, dropout_rate = DEFAULT_EMBEDDING_DIM, DEFAULT_RNN_UNITS, DEFAULT_NUM_LAYERS, DEFAULT_DROPOUT_RATE
    checkpoint_dir = f"./checkpoints_emb{emb_dim}_rnn{rnn_units}_layers{num_layers}"

    # Determine execution behavior based on user arguments
    is_resuming = args.modelfile is not None
    initial_epoch = 0

    if is_resuming:
        # Check for a matching configuration JSON file alongside the weights file
        base_path, _ = os.path.splitext(args.modelfile)
        # We strip off the ".weights" part if present to look for the matching json
        if base_path.endswith('.weights'):
            base_path = base_path[:-8]

        config_path = f"{base_path}_config.json"

        if os.path.exists(config_path):
            print(f"[Config] Found matching footprint configuration at: {config_path}")
            with open(config_path, 'r') as f:
                saved_config = json.load(f)
            emb_dim = saved_config["embedding_dim"]
            rnn_units = saved_config["rnn_units"]
            num_layers = saved_config.get("num_layers", DEFAULT_NUM_LAYERS)
            dropout_rate = saved_config.get("dropout_rate", DEFAULT_DROPOUT_RATE)
            saved_lr = saved_config.get("learning_rate", None)
            checkpoint_dir = saved_config.get("checkpoint_dir", os.path.dirname(args.modelfile))
            print(f"-> Architecture Adjusted to match saved model: Embedding={emb_dim}, RNN={rnn_units}, Layers={num_layers}")
        else:
            print(f"[Warning] No unique configuration file found at '{config_path}'. Falling back to script defaults.")
        filename = os.path.basename(args.modelfile)
        epoch_match = re.search(r'ckpt_(\d+)', filename)
        if epoch_match:
            initial_epoch = int(epoch_match.group(1))
            print(f"[Tracking] Detected training resume point. Resuming starting at Epoch: {initial_epoch}")

    else:
        print("[Mode] No model file passed. Initializing a brand-new training run profile.")
        os.makedirs(checkpoint_dir, exist_ok=True)

    # Initialize model dynamically using the determined metrics
    model = CharacterTextModel(vocab_size=vocab_size, embedding_dim=emb_dim, rnn_units=rnn_units, num_layers=num_layers, dropout_rate=dropout_rate)
    model.build(input_shape=(BATCH_SIZE, SEQ_LENGTH))

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

    if is_resuming:
        print(f"[Weights] Restoring model layer states from: {args.modelfile}")
        model.load_weights(args.modelfile)
        if saved_lr is not None:
            model.optimizer.learning_rate.assign(saved_lr)
            print(f"-> Optimizer Learning Rate restored to: {saved_lr}")

    # Process data pipeline if training is requested OR if no parameters were specified at all
    should_train = args.train or (args.modelfile is None and args.prompt is None)

    if should_train:
        if not os.path.exists(TEXT_FILE):
            print(f"[Error] Raw text source '{TEXT_FILE}' required for training. Exiting.")
            quit(1)

        print(f"[Data] Preparing dataset batches from {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)

        # Callbacks
        # Note: Weights and JSON foot-prints share matching index sequences
        checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")

        class SaveConfigCallback(tf.keras.callbacks.Callback):
            def on_epoch_end(self, epoch, logs=None):
                # Save out metadata alongside weights files every epoch
                epoch_config_path = f"{checkpoint_prefix.format(epoch=epoch+1)}_config.json"
                current_footprint = {
                    "embedding_dim": emb_dim,
                    "rnn_units": rnn_units,
                    "num_layers": num_layers,
                    "dropout_rate": dropout_rate,
                    "checkpoint_dir": checkpoint_dir,
                    "learning_rate": float(self.model.optimizer.learning_rate.numpy())
                }
                with open(epoch_config_path, 'w') as f:
                    json.dump(current_footprint, f, indent=4)

        checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
            filepath=checkpoint_prefix + ".weights.h5",
            save_weights_only=True
        )

        early_stopping = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=4, restore_best_weights=True)
        lr_scheduler = tf.keras.callbacks.ReduceLROnPlateau(monitor='loss', factor=0.5, patience=2, min_lr=1e-6, verbose=1)
        gen_callback = GenerationCallback(vectorize_layer, vocab, SEED_TEXT)

        total_epochs = initial_epoch + EPOCHS_PER_RUN
        print(f"[Train] Commencing training loop inside checkpoint tree: {checkpoint_dir}")
        print(f"[Train] Running Epochs {initial_epoch + 1} through {total_epochs}...")

        model.fit(
            dataset,
            epochs=total_epochs,
            initial_epoch = initial_epoch,
            callbacks=[gen_callback, checkpoint_callback, SaveConfigCallback(), early_stopping, lr_scheduler]
        )

    # Process custom Inference Generations
    if args.prompt is not None:
        print(f"\n--- Custom Inference Mode (Seed: \"{args.prompt}\") ---")
        result = produce_sample(model, vectorize_layer, vocab, seed=args.prompt, num_generate=300, temperature=0.6)
        print(f"Result:\n{result}\n")
    elif not should_train and args.modelfile is not None:
        # User specified a file but gave no training or prompt orders; run a quick test generation
        print(f"\n--- Default Quick Test (Seed: \"{SEED_TEXT}\") ---")
        result = produce_sample(model, vectorize_layer, vocab, seed=SEED_TEXT, num_generate=300, temperature=0.6)
        print(f"Result:\n{result}\n")

if __name__ == "__main__":
    main()
