import gradio as gr
import torch
import numpy as np
import plotly.graph_objects as go
from cognitive_net import DynamicCognitiveNet

class ModelDemo:
    def __init__(self):
        self.net = DynamicCognitiveNet(input_size=5, output_size=1)
        self.training_history = []
        self.emotional_history = []
        
    def train_sequence(self, sequence_str, epochs):
        """Train model on input sequence and return visualizations"""
        try:
            # Parse input sequence
            sequence = [float(x.strip()) for x in sequence_str.split(',')]
            if len(sequence) < 6:
                return "Error: Please input at least 6 numbers", None, None
            
            # Prepare training data
            X = torch.tensor(sequence[:-1]).float()
            y = torch.tensor([sequence[-1]]).float()
            
            # Training loop
            losses = []
            emotions = []
            for epoch in range(epochs):
                loss = self.net.train_step(X, y)
                losses.append(loss)
                emotions.append(self.net.emotional_state.item())
            
            # Create loss plot
            loss_fig = go.Figure()
            loss_fig.add_trace(go.Scatter(y=losses, name='Loss'))
            loss_fig.update_layout(title='Training Loss', 
                                 xaxis_title='Epoch',
                                 yaxis_title='Loss')
            
            # Create emotion plot
            emotion_fig = go.Figure()
            emotion_fig.add_trace(go.Scatter(y=emotions, name='Emotional State'))
            emotion_fig.update_layout(title='Emotional State', 
                                    xaxis_title='Epoch',
                                    yaxis_title='Value')
            
            # Make prediction
            with torch.no_grad():
                pred = self.net(X)
                result = f"Prediction: {pred.item():.4f} (Target: {y.item():.4f})"
            
            return result, loss_fig, emotion_fig
            
        except Exception as e:
            return f"Error: {str(e)}", None, None
    
    def visualize_memory(self):
        """Visualize memory importance weights"""
        memories = []
        importances = []
        
        for mem in self.net.nodes['input_0'].memory.memory_queue:
            memories.append(mem['context'].numpy())
            importances.append(mem['importance'].item())
            
        if not memories:
            return "No memories stored yet"
            
        fig = go.Figure()
        fig.add_trace(go.Bar(y=importances))
        fig.update_layout(title='Memory Importance',
                         xaxis_title='Memory Index',
                         yaxis_title='Importance')
        return fig

# Initialize demo
demo = ModelDemo()

# Create Gradio interface
with gr.Blocks(title="Cognitive Network Demo") as iface:
    gr.Markdown("""
    # Cognitive Network Interactive Demo
    
    This demo shows a neural network with:
    - Dynamic memory
    - Emotional modulation
    - Adaptive structure
    
    Enter a sequence of numbers (comma-separated) to train the model to predict the next number.
    """)
    
    with gr.Row():
        with gr.Column():
            input_seq = gr.Textbox(label="Input Sequence (comma-separated)", 
                                 value="1, 2, 3, 4, 5, 6")
            epochs = gr.Slider(minimum=10, maximum=500, value=100, 
                             step=10, label="Training Epochs")
            train_btn = gr.Button("Train Model")
        
        result_text = gr.Textbox(label="Prediction Result")
    
    with gr.Row():
        loss_plot = gr.Plot(label="Training Loss")
        emotion_plot = gr.Plot(label="Emotional State")
    
    with gr.Row():
        memory_btn = gr.Button("Visualize Memory")
        memory_plot = gr.Plot(label="Memory Importance")
    
    # Connect components
    train_btn.click(
        fn=demo.train_sequence,
        inputs=[input_seq, epochs],
        outputs=[result_text, loss_plot, emotion_plot]
    )
    
    memory_btn.click(
        fn=demo.visualize_memory,
        inputs=None,
        outputs=memory_plot
    )

# Launch app
iface.launch()