# app.py
import gradio as gr
import tensorflow as tf
import pickle
import numpy as np
import os

# --- 1. CONFIGURATION & MODEL LOADING ---
# This section loads your trained AI models and the tokenizers needed to understand text.
MAX_SEQ_LENGTH = 30 # Must match the value used during training!

print("Loading models and tokenizers...")
try:
    # Load the "Go Larger" model and its vocabulary
    successor_model = tf.keras.models.load_model('successor_model.h5')
    with open('successor_model_tokenizers.pkl', 'rb') as f:
        successor_tokenizers = pickle.load(f)
    
    # Load the "Go Smaller" model and its vocabulary
    predecessor_model = tf.keras.models.load_model('predecessor_model.h5')
    with open('predecessor_model_tokenizers.pkl', 'rb') as f:
        predecessor_tokenizers = pickle.load(f)

    print("Models and tokenizers loaded successfully.")
except Exception as e:
    # This helps debug issues on Hugging Face Spaces if a file is missing
    print(f"FATAL ERROR loading files: {e}")
    successor_model, predecessor_model = None, None

# --- 2. THE CORE PREDICTION LOGIC ---
# This function is the "brain" of the application.
def predict_next_state(direction, current_unit, current_analogy, current_commentary):
    # Safety check in case models failed to load
    if not all([successor_model, predecessor_model]):
        return "Error: Models are not loaded.", "Please check the server logs on Hugging Face.", "---"

    # A. Select the correct AI model and tokenizers based on user's click
    model = successor_model if direction == "larger" else predecessor_model
    tokenizers = successor_tokenizers if direction == "larger" else predecessor_tokenizers
    
    # B. Prepare the input data for the model
    # The input text must be converted to numbers exactly as it was during training.
    input_data = {
        'current_unit_name': [current_unit],
        'current_analogy': [current_analogy],
        'current_commentary': [current_commentary]
    }
    
    processed_input = {}
    for col, text_list in input_data.items():
        sequences = tokenizers[col].texts_to_sequences(text_list)
        padded_sequences = tf.keras.preprocessing.sequence.pad_sequences(sequences, maxlen=MAX_SEQ_LENGTH, padding='post')
        processed_input[col] = padded_sequences
        
    # C. Get the AI's prediction
    predictions = model.predict(processed_input)
    
    # D. Decode the prediction from numbers back to human-readable text
    target_texts = {}
    output_cols = ['target_unit_name', 'target_analogy', 'target_commentary']
    
    for i, col in enumerate(output_cols):
        # The model outputs probabilities; we take the most likely token (word) at each step.
        pred_indices = np.argmax(predictions[i], axis=-1)
        # Use the tokenizer to convert the sequence of indices back into a sentence.
        predicted_sequence = tokenizers[col].sequences_to_texts(pred_indices)[0]
        # Clean up padding and unknown words
        target_texts[col] = predicted_sequence.replace('<oov>', '').replace(' end', '').strip()
    
    # E. Handle the "Infinity" Sentinel
    # Check if the AI returned our special signal.
    if "end of knowledge" in target_texts['target_unit_name'].lower():
        # If so, switch to the simple rule-based procedural engine.
        prefix = "Giga-" if direction == "larger" else "pico-"
        new_unit = f"{prefix}{current_unit}"
        new_analogy = "A procedurally generated unit beyond the AI's known universe."
        new_commentary = "This represents a step into true infinity, where rules replace learned knowledge."
        return new_unit, new_analogy, new_commentary
    else:
        # Otherwise, return the AI's generated response.
        return target_texts['target_unit_name'], target_texts['target_analogy'], target_texts['target_commentary']

# Wrapper functions for the buttons
def go_larger(unit, analogy, commentary):
    return predict_next_state("larger", unit, analogy, commentary)

def go_smaller(unit, analogy, commentary):
    return predict_next_state("smaller", unit, analogy, commentary)

# --- 3. THE GRADIO USER INTERFACE ---
# This section defines the layout and interactivity of the web page.
initial_unit = "Byte"
initial_analogy = "A single character of text, like 'R'"
initial_commentary = "From binary choices, a building block is formed, ready to hold a single, recognizable symbol."

# Use gr.Blocks for a custom layout
with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as demo:
    gr.Markdown("# 🤖 Digital Scale Explorer AI")
    gr.Markdown("An AI trained from scratch to explore the infinite ladder of data sizes. Click the buttons to traverse the universe of data!")
    
    with gr.Row():
        # Define the output text boxes
        unit_name_out = gr.Textbox(value=initial_unit, label="Unit Name", interactive=False, elem_id="unit_name_style")
        analogy_out = gr.Textbox(value=initial_analogy, label="Analogy", lines=4, interactive=False, elem_id="analogy_style")
        commentary_out = gr.Textbox(value=initial_commentary, label="AI Commentary", lines=3, interactive=False, elem_id="commentary_style")

    with gr.Row():
        # Define the buttons
        smaller_btn = gr.Button("Go Smaller ⬇️", variant="secondary", size="lg")
        larger_btn = gr.Button("Go Larger ⬆️", variant="primary", size="lg")
    
    # Connect the "Go Larger" button to its function
    larger_btn.click(
        fn=go_larger,
        inputs=[unit_name_out, analogy_out, commentary_out],
        outputs=[unit_name_out, analogy_out, commentary_out]
    )
    # Connect the "Go Smaller" button to its function
    smaller_btn.click(
        fn=go_smaller,
        inputs=[unit_name_out, analogy_out, commentary_out],
        outputs=[unit_name_out, analogy_out, commentary_out]
    )

# Launch the app when the script is run
if __name__ == "__main__":
    demo.launch()