import os
import sys
import torch
import gradio as gr
from huggingface_hub import hf_hub_download, snapshot_download
from pathlib import Path
import importlib.util

class CognitiveNetworkDemo:
    def __init__(self):
        self.model = None
        self.repo_path = None
        self.setup_environment()
        
    def setup_environment(self):
        """Setup the environment and download the model"""
        # Download repository content
        self.repo_path = Path(snapshot_download(
            repo_id="VLabTech/cognitive_net",
            repo_type="model",
            local_dir="./model_repo"
        ))
    
    def import_module_from_file(self, module_name, file_path):
        """Import a module from file path"""
        spec = importlib.util.spec_from_file_location(module_name, file_path)
        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        spec.loader.exec_module(module)
        return module
            
    def load_model(self):
        """Load the model if not already loaded"""
        if self.model is None:
            try:
                # Find and import required modules
                network_path = next(self.repo_path.rglob("network.py"))
                node_path = next(self.repo_path.rglob("node.py"))
                memory_path = next(self.repo_path.rglob("memory.py"))
                
                # Import modules using absolute paths
                node_module = self.import_module_from_file("node", node_path)
                memory_module = self.import_module_from_file("memory", memory_path)
                network_module = self.import_module_from_file("network", network_path)
                
                # Create model instance
                self.model = network_module.DynamicCognitiveNet(input_size=5, output_size=2)
                
            except StopIteration:
                raise ImportError("Tidak dapat menemukan file modul yang diperlukan")
            except Exception as e:
                print("Debug - repo path:", self.repo_path)
                print("Debug - sys.path:", sys.path)
                raise ImportError(f"Gagal mengimpor model: {str(e)}")
        return self.model

    def predict(self, input_text):
        """Make predictions using the model"""
        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."
            
            # Load model and generate prediction
            model = self.load_model()
            input_tensor = torch.tensor(values, dtype=torch.float32)
            output = model(input_tensor)
            
            # Format output
            result = "Hasil Prediksi:\n"
            result += f"Output 1: {output[0]:.4f}\n"
            result += f"Output 2: {output[1]:.4f}"
            
            return result
            
        except ValueError as e:
            return f"Error dalam format input: {str(e)}"
        except Exception as e:
            return f"Error: {str(e)}\n\nDebug info:\nRepo path: {self.repo_path}"

def main():
    # Initialize the demo
    demo_app = CognitiveNetworkDemo()
    
    # Setup Gradio Interface
    demo = gr.Interface(
        fn=demo_app.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 menggunakan
        arsitektur Cognitive Network yang terinspirasi dari cara kerja otak biologis.
        Model diambil dari VLabTech/cognitive_net.
        """,
        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"]
        ]
    )
    
    demo.launch()

if __name__ == "__main__":
    main()