# app.py
import os
import torch
import gradio as gr
from huggingface_hub import hf_hub_download
from cognitive_net import DynamicCognitiveNet, CognitiveMemory, CognitiveNode

class ModelInterface:
    def __init__(self):
        self.model = self.load_model()
        
    def load_model(self):
        """Load model dari HuggingFace Hub"""
        try:
            # Inisialisasi model
            model = DynamicCognitiveNet(input_size=5, output_size=2)
            
            # Download checkpoint
            checkpoint_path = hf_hub_download(
                repo_id="VLabTech/cognitive_net",
                filename="model.pt"
            )
            
            # Load weights
            model.load_state_dict(torch.load(checkpoint_path))
            model.eval()
            return model
        except Exception as e:
            print(f"Error loading model: {e}")
            return None

    def predict(self, input_text):
        """Proses input dan generate prediksi"""
        try:
            # Parse input
            values = [float(x.strip()) for x in input_text.split(",")]
            if len(values) != 5:
                return f"Error: Masukkan tepat 5 nilai (dipisahkan koma). Anda memasukkan {len(values)} nilai."
            
            # Convert ke tensor
            input_tensor = torch.tensor(values, dtype=torch.float32)
            
            # Generate prediksi
            with torch.no_grad():
                output = self.model(input_tensor)
            
            # Format output
            result = f"Hasil Prediksi:\n"
            result += f"Output 1: {output[0]:.4f}\n"
            result += f"Output 2: {output[1]:.4f}"
            
            return result
            
        except ValueError:
            return "Error: Pastikan semua input adalah angka valid"
        except Exception as e:
            return f"Error: {str(e)}"

# Inisialisasi interface
model_interface = ModelInterface()

# Setup Gradio Interface
demo = gr.Interface(
    fn=model_interface.predict,
    inputs=gr.Textbox(
        label="Input Values",
        placeholder="Masukkan 5 nilai numerik (pisahkan dengan koma). Contoh: 1.0, 2.0, 3.0, 4.0, 5.0"
    ),
    outputs=gr.Textbox(label="Hasil Prediksi"),
    title="Cognitive Network Demo",
    description="""
    ## Cognitive Network Inference Demo
    
    Model ini menerima 5 input numerik dan menghasilkan 2 output numerik.
    
    ### Cara Penggunaan:
    1. Masukkan 5 angka yang dipisahkan dengan koma
    2. Klik Submit untuk melihat hasil prediksi
    
    ### Format Input:
    - Masukkan tepat 5 nilai numerik
    - Pisahkan nilai dengan koma
    - Gunakan titik untuk desimal
    """,
    examples=[
        ["1.0, 2.0, 3.0, 4.0, 5.0"],
        ["0.5, -1.0, 2.5, 1.5, -0.5"],
        ["0.1, 0.2, 0.3, 0.4, 0.5"]
    ]
)

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