import gradio as gr
import torch
import numpy as np
import plotly.graph_objects as go
from huggingface_hub import hf_hub_download
import sys
import os
from pathlib import Path

# 1. Download Model Files ------------------------------------------------
def setup_model():
    REPO_ID = "VLabTech/cognitive_net"
    FILE_LIST = [
        "cognitive_net/__init__.py",
        "cognitive_net/memory.py",
        "cognitive_net/node.py", 
        "cognitive_net/network.py"
    ]
    
    # Create package directory
    model_dir = Path("cognitive_net")
    model_dir.mkdir(exist_ok=True)
    
    # Download files
    for file in FILE_LIST:
        try:
            downloaded_file = hf_hub_download(
                repo_id=REPO_ID,
                filename=file,
                local_dir=model_dir.parent,
                force_filename=file
            )
        except Exception as e:
            print(f"Error downloading {file}: {str(e)}")
    
    # Add to Python path
    if str(model_dir.absolute()) not in sys.path:
        sys.path.insert(0, str(model_dir.absolute()))

# 2. Initialize Model ----------------------------------------------------
class CognitiveDemo:
    def __init__(self):
        setup_model()
        
        try:
            from cognitive_net import DynamicCognitiveNet
            self.net = DynamicCognitiveNet(input_size=5, output_size=1)
            self.net.optimizer = torch.optim.AdamW(self.net.parameters(), lr=0.001)
        except ImportError as e:
            raise RuntimeError(f"Gagal memuat model: {str(e)}")
        
        self.training_history = []

    def _adapt_model(self, X: torch.Tensor, y: torch.Tensor):
        """Penyesuaian dimensi tensor untuk arsitektur kognitif"""
        X = X.view(-1, 1)  # Bentuk (seq_len, 1)
        y = y.view(1)      # Bentuk (1,)
        return X, y

    def train(self, sequence: str, epochs: int):
        try:
            # Parse dan validasi input
            nums = [float(n.strip()) for n in sequence.split(',')]
            if len(nums) < 6:
                raise ValueError("Input minimal 6 angka")
            
            X = torch.tensor(nums[:-1])
            y = torch.tensor([nums[-1]])
            
            # Adaptasi dimensi
            X, y = self._adapt_model(X, y)
            
            # Training loop
            losses = []
            for _ in range(epochs):
                loss = self.net.train_step(X, y)
                losses.append(loss)
                
            return {
                "prediction": self.net(X).detach().numpy()[0],
                "loss_plot": self._create_plot(losses, "Loss Training"),
                "emotion": self.net.emotional_state.item()
            }
            
        except Exception as e:
            return {"error": str(e)}

    def _create_plot(self, data, title):
        fig = go.Figure()
        fig.add_trace(go.Scatter(y=data, mode='lines'))
        fig.update_layout(title=title)
        return fig

# 3. Gradio Interface ----------------------------------------------------
demo = CognitiveDemo()

with gr.Blocks(title="Cognitive Network Demo") as app:
    gr.Markdown("""## Demo Jaringan Kognitif VLabTech""")
    
    with gr.Row():
        input_seq = gr.Textbox(label="Deret Input (contoh: 0.1, 0.5, 1.0,...)", 
                             value="0.1, 0.3, 0.5, 0.7, 0.9, 1.1")
        epochs = gr.Slider(10, 500, value=100, label="Jumlah Epoch")
    
    with gr.Row():
        train_btn = gr.Button("🚀 Latih Model")
        pred_output = gr.Label(label="Hasil Prediksi")
        emotion_output = gr.Number(label="Status Emosional")
    
    loss_plot = gr.Plot(label="Progress Training")
    
    train_btn.click(
        fn=lambda s, e: demo.train(s, e),
        inputs=[input_seq, epochs],
        outputs=[pred_output, loss_plot, emotion_output]
    )

# 4. Run App -------------------------------------------------------------
if __name__ == "__main__":
    app.launch(debug=True)