# app.py (Final Version with gr.State for Robust State Management)
import gradio as gr
import tensorflow as tf
import pickle
import numpy as np

# --- 1. CONFIGURATION & MODEL LOADING ---
MAX_SEQ_LENGTH = 30
print("--- App Starting Up ---")
print("Loading models and tokenizers...")
try:
    successor_model = tf.keras.models.load_model('successor_model.h5')
    with open('successor_model_tokenizers.pkl', 'rb') as f:
        successor_tokenizers = pickle.load(f)
    
    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:
    print(f"FATAL ERROR loading files: {e}")
    successor_model, predecessor_model = None, None

# --- 2. THE CORE PREDICTION LOGIC ---
# This function is the same, but it will now receive its input from the reliable gr.State
def predict_next_state(model, tokenizers, current_state_dict):
    if not model or not tokenizers or not current_state_dict:
        return {"error": "Model or state is not loaded"}, "Error", "Error", "Error"

    # Prepare input data from the state dictionary
    input_data = {
        'current_unit_name': [current_state_dict['unit_name']],
        'current_analogy': [current_state_dict['analogy']],
        'current_commentary': [current_state_dict['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
        
    predictions = model.predict(processed_input)
    
    # Decode prediction
    target_texts = {}
    output_cols = ['target_unit_name', 'target_analogy', 'target_commentary']
    for i, col in enumerate(output_cols):
        pred_indices = np.argmax(predictions[i], axis=-1)
        predicted_sequence = tokenizers[col].sequences_to_texts(pred_indices)[0]
        clean_text = ' '.join([word for word in predicted_sequence.split() if word not in ['<oov>', 'end']])
        target_texts[col] = clean_text.strip()
    
    print(f"Decoded Unit Name: {target_texts['target_unit_name']}")

    # Create the new state dictionary
    new_state = {
        'unit_name': target_texts['target_unit_name'],
        'analogy': target_texts['target_analogy'],
        'commentary': target_texts['target_commentary']
    }

    # Handle "Infinity" Sentinel
    if "end of knowledge" in new_state['unit_name'].lower():
        direction = "larger" if model == successor_model else "smaller"
        prefix = "Giga-" if direction == "larger" else "pico-"
        new_state['unit_name'] = f"{prefix}{current_state_dict['unit_name']}"
        new_state['analogy'] = "A procedurally generated unit beyond the AI's known universe."
        new_state['commentary'] = "This represents a step into true infinity, where rules replace learned knowledge."
    
    # Return the new state object and the values for the textboxes
    return new_state, new_state['unit_name'], new_state['analogy'], new_state['commentary']

# --- WRAPPER FUNCTIONS ---
# They now take the state dictionary as input and return the new state dictionary
def go_larger(current_state):
    print("\n>>> 'Go Larger' button clicked. Using SUCCESSOR model.")
    return predict_next_state(successor_model, successor_tokenizers, current_state)

def go_smaller(current_state):
    print("\n>>> 'Go Smaller' button clicked. Using PREDECESSOR model.")
    return predict_next_state(predecessor_model, predecessor_tokenizers, current_state)

# --- 3. THE GRADIO USER INTERFACE (RE-ARCHITECTED) ---
initial_state = {
    "unit_name": "Byte",
    "analogy": "a single character of text, like 'R'",
    "commentary": "From binary choices, a building block is formed, ready to hold a single, recognizable symbol."
}

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!")
    
    # *** THIS IS THE KEY CHANGE ***
    # Create an invisible component to reliably hold our state
    app_state = gr.State(value=initial_state)

    with gr.Row():
        unit_name_out = gr.Textbox(value=initial_state['unit_name'], label="Unit Name", interactive=False)
        analogy_out = gr.Textbox(value=initial_state['analogy'], label="Analogy", lines=4, interactive=False)
        commentary_out = gr.Textbox(value=initial_state['commentary'], label="AI Commentary", lines=3, interactive=False)

    with gr.Row():
        smaller_btn = gr.Button("Go Smaller ⬇️", variant="secondary", size="lg")
        larger_btn = gr.Button("Go Larger ⬆️", variant="primary", size="lg")
    
    # --- The button clicks now use the app_state as their primary input and output ---
    larger_btn.click(
        fn=go_larger,
        inputs=[app_state],  # The INPUT is the reliable state object
        # The OUTPUT is the new state object AND the values for the textboxes
        outputs=[app_state, unit_name_out, analogy_out, commentary_out] 
    )
    smaller_btn.click(
        fn=go_smaller,
        inputs=[app_state], # The INPUT is the reliable state object
        outputs=[app_state, unit_name_out, analogy_out, commentary_out]
    )

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