import gradio as gr
import tensorflow as tf
import numpy as np
from model import TeraV3

# --- CONFIGURATION ---
VOCAB_SIZE = 100 
WEIGHTS_PATH = 'tera_v3_weights.weights.h5'

# Initialize the model
print("Loading Tera v3 Sovereign Model...")
model = TeraV3(vocab_size=VOCAB_SIZE, dim=256, depth=6)

# Load the weights
try:
    model.load_weights(WEIGHTS_PATH)
    print("✅ Weights loaded successfully!")
except Exception as e:
    print(f"❌ Error loading weights: {e}")

# Simple character map to match your training
chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 "
char_to_int = {c: i for i, c in enumerate(chars)}
int_to_char = {i: c for i, c in enumerate(chars)}

def generate_response(message, history):
    # Tokenize input
    tokens = [char_to_int.get(c, 0) for c in message]
    input_tensor = tf.constant([tokens])
    
    # Generate output (predict next 50 characters)
    response = ""
    current_input = input_tensor
    
    for _ in range(50):
        predictions = model.predict(current_input, verbose=0)
        predicted_id = tf.argmax(predictions[0, -1]).numpy()
        next_char = int_to_char.get(predicted_id, "?")
        response += next_char
        
        # Append predicted token to input for next step
        next_token = tf.constant([[predicted_id]])
        current_input = tf.concat([current_input, next_token], axis=1)
        
        if next_char == ".": break
        
    return response

# Create the Gradio Interface
demo = gr.ChatInterface(
    fn=generate_response,
    title="Tera v3: Sovereign AI",
    description="Experience the power of Time Mix and Squared ReLU. A model built to compete with the frontier.",
    examples=["Hello Tera v3!", "Explain Quantum AI", "What is the future of AI?"]
)

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