"""
Gradio app for the Brain-like Predictive Coding Code World Model.
"""

import gradio as gr
import numpy as np
import sys
import os

# Train and load model on startup
from brain_predictive_coding import PredictiveCodingNetwork, SimpleCodeTokenizer, generate_code

print("Training model...")
tok = SimpleCodeTokenizer(vocab_size=128)

# Generate and train
SEQ_LEN = 16
EMBED = 32
N_SAMPLES = 40
EPOCHS = 15

code = generate_code(n=N_SAMPLES, max_len=SEQ_LEN)
sequences = np.array([tok.embed_seq(tok.encode(c, SEQ_LEN)) for c in code])

net = PredictiveCodingNetwork(
    embed_dim=EMBED,
    l1_n=128, l2_n=96, l3_n=64,
    l1_lr=5e-5, l2_lr=5e-5, l3_lr=5e-5
)

for epoch in range(EPOCHS):
    for i in range(N_SAMPLES):
        net.context = np.zeros_like(net.context)
        net.process_seq(sequences[i], train=True)

print("Model trained!")


def predict_code(code_text, n_steps=5):
    """Predict next characters in code sequence."""
    if not code_text:
        return "Please enter some code!"
    
    net.context = np.zeros_like(net.context)
    
    tokens = tok.encode(code_text, max_len=16)
    embeddings = tok.embed_seq(tokens)
    
    preds = net.predict_next(embeddings, n_steps=n_steps)
    predicted_chars = [tok.nearest(p) for p in preds]
    
    stats = {
        "L1 mean activity": float(np.mean(net.l1.activities)),
        "L1 sparsity": f"{np.mean(net.l1.activities > 0):.1%}",
        "L2 mean activity": float(np.mean(net.l2.activities)),
        "L2 sparsity": f"{np.mean(net.l2.activities > 0):.1%}",
        "L3 mean activity": float(np.mean(net.l3.activities)),
        "L3 sparsity": f"{np.mean(net.l3.activities > 0):.1%}",
        "Context magnitude": float(np.linalg.norm(net.context)),
    }
    
    result = f"**Input:** `{code_text}`\n\n"
    result += f"**Predicted next {n_steps} characters:**\n"
    for i, ch in enumerate(predicted_chars):
        result += f"  Step {i+1}: `{ch}`\n"
    
    result += f"\n**Brain-like Network Statistics:**\n"
    for k, v in stats.items():
        result += f"  {k}: {v}\n"
    
    return result


def get_layer_activities(code_text):
    """Show layer activity statistics."""
    net.context = np.zeros_like(net.context)
    tokens = tok.encode(code_text, max_len=16)
    embeddings = tok.embed_seq(tokens)
    net.process_seq(embeddings, train=False)
    
    l1_acts = net.l1.activities
    l2_acts = net.l2.activities
    l3_acts = net.l3.activities
    
    output = "**Layer Activities (sample of active neurons):**\n\n"
    
    output += f"L1 (Sensory, {len(l1_acts)} neurons):\n"
    active_l1 = np.where(l1_acts > 0.01)[0]
    output += f"  Active: {len(active_l1)} ({len(active_l1)/len(l1_acts):.1%})\n"
    output += f"  Top 5: {l1_acts[np.argsort(l1_acts)[-5:]][::-1].round(4).tolist()}\n\n"
    
    output += f"L2 (Hidden, {len(l2_acts)} neurons):\n"
    active_l2 = np.where(l2_acts > 0.01)[0]
    output += f"  Active: {len(active_l2)} ({len(active_l2)/len(l2_acts):.1%})\n"
    output += f"  Top 5: {l2_acts[np.argsort(l2_acts)[-5:]][::-1].round(4).tolist()}\n\n"
    
    output += f"L3 (Context, {len(l3_acts)} neurons):\n"
    active_l3 = np.where(l3_acts > 0.01)[0]
    output += f"  Active: {len(active_l3)} ({len(active_l3)/len(l3_acts):.1%})\n"
    output += f"  Top 5: {l3_acts[np.argsort(l3_acts)[-5:]][::-1].round(4).tolist()}\n"
    
    return output


description = """
# 🧠 Brain-like Predictive Coding Code World Model

This model uses a **hierarchical predictive coding network** inspired by the brain's cortical hierarchy:
- **L1 (Sensory)**: Processes code token embeddings like primary visual cortex
- **L2 (Hidden)**: Learns associative patterns like inferotemporal cortex
- **L3 (Context)**: Maintains sequence context like prefrontal cortex

**Brain-like features:**
- LIF (Leaky Integrate-and-Fire) neurons
- PES (Prescribed Error Sensitivity) learning — error-driven weight updates
- Top-down predictions from higher layers
- Prediction errors drive learning (free-energy principle)
- Numba JIT acceleration for fast CPU inference
"""

with gr.Blocks(title="Brain-like PC Code Model") as demo:
    gr.Markdown(description)
    
    with gr.Tab("Code Prediction"):
        with gr.Row():
            with gr.Column():
                code_input = gr.Textbox(
                    label="Input Code",
                    placeholder="def compute(x):\n  return",
                    lines=3
                )
                n_steps = gr.Slider(1, 10, value=5, step=1, label="Prediction Steps")
                predict_btn = gr.Button("Predict Next Tokens")
            
            with gr.Column():
                output = gr.Markdown(label="Predictions")
        
        predict_btn.click(predict_code, inputs=[code_input, n_steps], outputs=output)
    
    with gr.Tab("Layer Activities"):
        with gr.Row():
            with gr.Column():
                code_input2 = gr.Textbox(
                    label="Input Code",
                    placeholder="for i in range(10):",
                    lines=2
                )
                act_btn = gr.Button("Show Activities")
            
            with gr.Column():
                act_output = gr.Markdown(label="Layer Activities")
        
        act_btn.click(get_layer_activities, inputs=code_input2, outputs=act_output)

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