import gradio as gr
import os

yolov5_result = os.path.join(os.getcwd(), "data/xai/yolov5.png")
yolov8_result = os.path.join(os.getcwd(), "data/xai/yolov8.png")
yolov5_dff = os.path.join(os.getcwd(), "data/xai/yolov5_dff.png")
yolov8_dff = os.path.join(os.getcwd(), "data/xai/yolov8_dff.png")

def view_model(selected_models):
    """Generate Netron visualization for the selected models."""
    netron_html = ""
    for model in selected_models:
        if model=="yolov8s":
            netron_html = f"""
            <iframe 
                src="https://netron.app/?url=https://huggingface.co/spaces/BhumikaMak/NeuralVista/resolve/main/weight_files/yolov8s.pt" 
                width="100%" 
                height="800" 
                frameborder="0">
            </iframe>
            """
        if model == "yolov5":
            netron_html = f"""
            <iframe 
                src="https://netron.app/?url=https://huggingface.co/FFusion/FFusionXL-BASE/blob/main/vae_encoder/model.onnx" 
                width="100%" 
                height="800" 
                frameborder="0">
            </iframe>
            """
    return netron_html if netron_html else "<p>No valid models selected for visualization.</p>"

custom_css = """
body {
    background-color: black; /* Navy blue background */
    background-size: 1800px 1800px; /* Grid cell size */
    height: 100%; /* Ensure body height is 100% of the viewport */
    margin: 0; /* Remove default margin */
    overflow-y: auto; /* Allow vertical scrolling if needed */
}
#neural-vista-title {
    color: #800000 !important;  /* Purple color for the title */
    font-size: 32px;           /* Adjust font size as needed */
    font-weight: bold;
    text-align: center;
}
#neural-vista-text {
    color: white  !important;  /* Purple color for the title */
    font-size: 18px;           /* Adjust font size as needed */
    font-weight: bold;
    text-align: center;
    
}

#highlighted-text {
    font-weight: bold;
    color: #1976d2;
}
"""

with gr.Blocks(css=custom_css) as demo:
    gr.HTML("""
      <div style="border: 2px solid #a05252; padding: 20px; border-radius: 8px;">
        <span style="color: #E6E6FA; font-family: 'Papyrus', cursive; font-weight: bold; font-size: 32px;">NeuralVista</span><br><br>
        <span style="color: #E6E6FA; font-family: 'Papyrus', cursive; font-size: 18px;">A harmonious framework of tools ☼ designed to illuminate the inner workings of AI.</span>
      </div>
    """)
    with gr.Row():
        with gr.Column():
            gr.Markdown("Yolov5")
            gr.Image(yolov5_result, label="Detections & Interpretability Map")
            gr.Image(yolov5_dff, label="Feature Factorization & discovered concept")
        with gr.Column():
            gr.Markdown("Yolov8")
            gr.Image(yolov8_result, label="Detections & Interpretability Map")
            gr.Image(yolov8_dff, label="Feature Factorization & discovered concept")

demo.launch()