import gradio as gr
import os

# Paths for images
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")

# Netron HTML templates
def get_netron_html(model_url):
    return f"""
        <iframe 
            src="{model_url}" 
            width="100%" 
            height="800" 
            frameborder="0" 
            style="background-color: black; border: 2px solid white;">
        </iframe>
    """

# URLs for Netron visualizations
yolov5_url = "https://netron.app/?url=https://huggingface.co/FFusion/FFusionXL-BASE/blob/main/vae_encoder/model.onnx"
yolov8_url = "https://netron.app/?url=https://huggingface.co/spaces/BhumikaMak/NeuralVista/resolve/main/weight_files/yolov8s.pt"

custom_css = """
body {
    background-color: black;
    background-size: 1800px 1800px;
    height: 100%;
    margin: 0;
    overflow-y: auto;
}
#neural-vista-title {
    color: #800000 !important;
    font-size: 32px;
    font-weight: bold;
    text-align: center;
}
#neural-vista-text {
    color: white  !important;
    font-size: 18px;
    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 <span style="color: yellow; font-family: 'Papyrus', cursive; font-size: 18px;">☼</span> 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")
            gr.HTML(get_netron_html(yolov5_url))

        with gr.Column():
            gr.Markdown("Yolov8")
            gr.Image(yolov8_result, label="Detections & Interpretability Map")
            gr.Image(yolov8_dff, label="Feature Factorization & discovered concept")
            gr.HTML(get_netron_html(yolov8_url))

demo.launch()