import gradio as gr
import tempfile
import shutil
import os
import subprocess
import re
import tensorflow as tf
import numpy as np
from PIL import Image

DEFAULT_CONFIG = "u55_eval_with_TA_config_400_and_200_MHz.ini"

def extract_summary_from_log(log_text):
    summary_keys = [
        "Accelerator configuration",
        "Accelerator clock",
        "Total SRAM used",
        "Total On-chip Flash used",
        "CPU operators",
        "NPU operators",
        "Batch Inference time"
    ]
    summary = []
    for key in summary_keys:
        match = re.search(rf"{re.escape(key)}\s+(.+)", log_text)
        if match:
            value = match.group(1).strip()
            if key == "Batch Inference time":
                value = value.split(",")[0].strip()
                key = "Inference time"
            summary.append((key, value))
    return summary

def run_vela(model_file):
    accel = "ethos-u55-128"
    optimise = "Size"
    mem_mode = "Sram_Only"
    sys_config = "Ethos_U55_400MHz_SRAM_3.2_GBs_Flash_0.05_GBs"
    tmpdir = tempfile.mkdtemp()
    try:
        # Use the original uploaded model filename
        original_model_name = os.path.basename(model_file)
        model_path = os.path.join(tmpdir, original_model_name)
        shutil.copy(model_file, model_path)

        config_path = os.path.join(tmpdir, DEFAULT_CONFIG)
        shutil.copy(DEFAULT_CONFIG, config_path)

        output_dir = os.path.join(tmpdir, "vela_out")
        os.makedirs(output_dir, exist_ok=True)

        cmd = [
            "vela",
            f"--accelerator-config={accel}",
            f"--optimise={optimise}",
            f"--config={config_path}",
            f"--memory-mode={mem_mode}",
            f"--system-config={sys_config}",
            model_path,
            "--verbose-cycle-estimate",
            "--verbose-performance",
            f"--output-dir={output_dir}"
        ]

        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        vela_stdout = result.stdout

        # Check for unsupported model patterns in logs
        unsupported_patterns = [
            "Warning: Unsupported TensorFlow Lite semantics",
            "Network Tops/s                                    nan Tops/s",
            "Neural network macs                                 0 MACs/batch"
        ]
        if any(pat in vela_stdout for pat in unsupported_patterns):
            summary_html = (
                '<div class="card">'
                '<div class="card-header error-header"><span style="color: white; font-weight: 600;">Unsupported Model</span></div>'
                '<div class="card-content">'
                'This model has unsupported layers and needs investigation based on layers.<br><br>'
                'Please use Vela compiler on your Host Machine for further analysis.'
                '</div></div>'
            )
            # Try to provide per-layer.csv if available for download
            per_layer_csv = None
            for log_fname in os.listdir(output_dir):
                if log_fname.endswith("per-layer.csv"):
                    per_layer_csv = os.path.join("/tmp", log_fname)
                    shutil.copy(os.path.join(output_dir, log_fname), per_layer_csv)
                    break
            return summary_html, None, per_layer_csv

        model_filename = os.path.basename(model_file)
        if model_filename:
            vela_stdout = vela_stdout.replace(
                "Network summary for",
                f"Network summary for {model_filename} ("
            )

        summary_items = extract_summary_from_log(vela_stdout)
        # Convert summary_items to dict for easy access
        summary_dict = dict(summary_items) if summary_items else {}

        # Build 4 cards for results
        def clean_ops(val):
            # Remove '=' and leading/trailing spaces
            return val.lstrip("= ").strip() if isinstance(val, str) else val

        summary_html = (
            '<div class="card">'
            '<div class="card-header"><span style="color: white; font-weight: 600;">Estimated Results on SR110</span></div>'
            '<div class="card-content">'
            '<div class="stats-grid">'
            # Card 1: Accelerator
            '<div class="stat-item">'
            '<div class="stat-label">Accelerator</div>'
            f'<div class="stat-value"><strong>Configuration:</strong> {summary_dict.get("Accelerator configuration","-")}<br><strong>Clock:</strong> {summary_dict.get("Accelerator clock","-")}</div>'
            '</div>'
            # Card 2: Memory Usage
            '<div class="stat-item">'
            '<div class="stat-label">Memory Usage</div>'
            f'<div class="stat-value"><strong>Total SRAM:</strong> {summary_dict.get("Total SRAM used","-")}<br><strong>Total Flash:</strong> {summary_dict.get("Total On-chip Flash used","-")}</div>'
            '</div>'
            # Card 3: Operator Distribution
            '<div class="stat-item">'
            '<div class="stat-label">Operator Distribution</div>'
            f'<div class="stat-value"><strong>CPU Operators:</strong> {clean_ops(summary_dict.get("CPU operators","-"))}<br><strong>NPU Operators:</strong> {clean_ops(summary_dict.get("NPU operators","-"))}</div>'
            '</div>'
            # Card 4: Performance
            '<div class="stat-item">'
            '<div class="stat-label">Performance</div>'
            f'<div class="stat-value"><strong>Inference time:</strong> {summary_dict.get("Inference time","-")}</div>'
            '</div>'
            '</div></div></div>'
        ) if summary_items else '<div class="card"><div class="card-header error-header"><span style="color: white; font-weight: 600;">Error</span></div><div class="card-content">Summary info not found in log.</div></div>'

        for fname in os.listdir(output_dir):
            if fname.endswith("vela.tflite"):
                final_path = os.path.join("/tmp", fname)
                shutil.copy(os.path.join(output_dir, fname), final_path)
                # Find per-layer.csv file for logs
                per_layer_csv = None
                for log_fname in os.listdir(output_dir):
                    if log_fname.endswith("per-layer.csv"):
                        per_layer_csv = os.path.join("/tmp", log_fname)
                        shutil.copy(os.path.join(output_dir, log_fname), per_layer_csv)
                        break
                return summary_html, final_path, per_layer_csv

        # If no tflite, still try to return per-layer.csv if present
        per_layer_csv = None
        for log_fname in os.listdir(output_dir):
            if log_fname.endswith("per-layer.csv"):
                per_layer_csv = os.path.join("/tmp", log_fname)
                shutil.copy(os.path.join(output_dir, log_fname), per_layer_csv)
                break
        return summary_html, None, per_layer_csv
    finally:
        shutil.rmtree(tmpdir)

# Create Gradio interface
with gr.Blocks(
    theme=gr.themes.Default(primary_hue="blue", neutral_hue="gray"),
    title="SR Vela Compiler",
    css="""
    body {
        background: #fafafa !important;
    }
    .gradio-container {
        max-width: none !important;
        margin: 0 !important;
        background-color: #fafafa !important;
        font-family: 'Inter', 'Segoe UI', -apple-system, sans-serif !important;
        width: 100vw !important;
    }
    .main-header {
        text-align: center;
        margin: 0 !important;
        color: #3b82f6 !important;
        font-weight: 600;
        font-size: 2.5rem;
        letter-spacing: -0.025em;
    }
    .main-header h1 {
        color: #3b82f6 !important;
        margin: 0 !important;
    }
    /* Force all h1 to be blue */
    h1 {
        color: #3b82f6 !important;
    }
    /* Force all markdown h1 to be blue */
    .gr-markdown h1 {
        color: #3b82f6 !important;
    }
    .subtitle {
        text-align: center;
        color: #4b5563 !important;
        font-size: 1.1rem;
        margin-top: 0.5rem !important;
        margin-bottom: 2rem !important;
    }
    .card {
        background: #fafafa !important;
        border-radius: 12px !important;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06) !important;
        border: 1px solid #e5e7eb !important;
        margin-bottom: 1.5rem !important;
        transition: all 0.2s ease-in-out !important;
        overflow: hidden !important;
    }
    /* Fix Downloads card styling */
    .card .gr-form {
        background: transparent !important;
        border: none !important;
        box-shadow: none !important;
        padding: 0 !important;
    }
    /* Ensure gr.Group with card class looks like a proper card */
    .gr-group.card {
        background: #fafafa !important;
        border-radius: 12px !important;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06) !important;
        border: 1px solid #e5e7eb !important;
        margin-bottom: 1.5rem !important;
        overflow: hidden !important;
        padding: 0 !important;
    }
    /* Remove inner container styling */
    .gr-group.card > div:first-child {
        background: transparent !important;
        border: none !important;
        padding: 0 !important;
    }
    /* Style download buttons like MobileNet's example buttons */
    .card-content button {
        background: #f1f5f9 !important;
        border: 1px solid #cbd5e1 !important;
        color: #4b5563 !important;
        border-radius: 6px !important;
        transition: all 0.2s ease !important;
        margin: 0.5rem 0 !important;
        width: 100% !important;
    }
    .card-content button:hover {
        background: #1975cf !important;
        border-color: #1975cf !important;
        color: white !important;
    }
    .card > * {
        padding: 0 !important;
        margin: 0 !important;
    }
    .card:hover {
        box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1), 0 4px 6px -2px rgba(0, 0, 0, 0.05) !important;
        transform: translateY(-1px) !important;
    }
    .card-header {
        background: linear-gradient(135deg, #1975cf 0%, #1557b0 100%) !important;
        color: white !important;
        padding: 1rem 1.5rem !important;
        border-radius: 12px 12px 0 0 !important;
        font-weight: 600 !important;
        font-size: 1.1rem !important;
    }
    .card-header * {
        color: white !important;
    }
    .card-content {
        padding: 1.5rem !important;
        color: #4b5563 !important;
        line-height: 1.6 !important;
        background: #fafafa !important;
    }
    .stats-grid {
        display: grid !important;
        grid-template-columns: 1fr 1fr !important;
        gap: 1.5rem !important;
        margin-top: 1.5rem !important;
    }
    .stat-item {
        background: #f8fafc !important;
        padding: 1rem !important;
        border-radius: 8px !important;
        border-left: 4px solid #1975cf !important;
    }
    .stat-label {
        font-weight: 600 !important;
        color: #4b5563 !important;
        font-size: 0.9rem !important;
        margin-bottom: 0.5rem !important;
    }
    .stat-value {
        color: #4b5563 !important;
        font-size: 0.85rem !important;
    }
    .stat-value strong {
        color: #4b5563 !important;
        font-weight: 600 !important;
    }
    .btn-primary {
        background: #1975cf !important;
        border-color: #1975cf !important;
        color: white !important;
    }
    .btn-primary:hover {
        background: #1557b0 !important;
        border-color: #1557b0 !important;
    }
    .markdown {
        color: #374151 !important;
    }
    .results-text {
        color: #4b5563 !important;
        font-weight: 500 !important;
        padding: 0 !important;
        margin: 0 !important;
    }
    .results-text p {
        color: #4b5563 !important;
        margin: 0.5rem 0 !important;
    }
    .results-text * {
        color: #4b5563 !important;
    }
    div[data-testid="markdown"] p {
        color: #4b5563 !important;
    }
    .prose {
        color: #4b5563 !important;
    }
    .prose * {
        color: #4b5563 !important;
    }
    .card-header,
    .card-header * {
        color: white !important;
    }
    .error-header {
        background: linear-gradient(135deg, #dc2626 0%, #b91c1c 100%) !important;
    }
    .custom-footer {
        max-width: 800px !important;
        margin: 2rem auto !important;
        background: white !important;
        border-radius: 12px !important;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06) !important;
        border: 1px solid #e5e7eb !important;
        padding: 1.5rem !important;
        text-align: center !important;
    }
    .custom-footer a {
        color: #1975cf !important;
        text-decoration: none !important;
        font-weight: 600 !important;
    }
    .custom-footer a:hover {
        text-decoration: underline !important;
    }
    """
) as demo:
    
    gr.HTML("""
    <div class="main-header">
        <h1 style="color: #3b82f6 !important;">SR Vela Compiler</h1>
    </div>
    <div class="subtitle">
        Upload your quantized TFLite model to get SR NPU performance estimates.<br>
        These estimates are being done in a Vela Toolchain Environment specific to SR110 configurations.
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            
            model_file = gr.File(
                label="",
                file_types=[".tflite"],
                type="filepath",
                height=280
            )
            
            compile_btn = gr.Button(
                "Compile Model for SR110",
                variant="primary",
                size="lg",
                elem_classes=["btn-primary"]
            )
            
            with gr.Group(elem_classes=["card"]):
                gr.HTML('<div class="card-header"><span style="color: white; font-weight: 600;">Downloads</span></div>')
                
                with gr.Column(elem_classes=["card-content"]):
                    download_btn = gr.DownloadButton(
                        label="Download Compiled Model",
                        visible=False,
                        size="sm"
                    )
                    download_logs_btn = gr.DownloadButton(
                        label="Download Logs Per Layer",
                        visible=False,
                        size="sm"
                    )
        
        with gr.Column(scale=1):
            summary_table = gr.HTML(elem_classes=["results-summary-box"])
    
    # Footer
    gr.HTML("""
    <div class="custom-footer">
        <div style="margin-bottom: 0.5rem;">
            For a detailed walkthrough, please see our
            <a href="http://localhost:3000/sr/evaluate-sr" target="_blank">Evaluate Model Guide</a>.
        </div>
        <div>
            To get started quickly, visit our
            <a href="http://localhost:3000/sr/quick-start" target="_blank">SR Quick Start page</a>.
        </div>
    </div>
    """)

    def wrapped_run_vela(model_file):
        if model_file is None:
            # Show error in a card-like box
            return (
                '<div class="card">'
                '<div class="card-header error-header"><span style="color: white; font-weight: 600;">No Model</span></div>'
                '<div class="card-content">'
                'No model file uploaded.'
                '</div></div>',
                gr.update(visible=False, value=None),
                gr.update(visible=False, value=None)
            )

        summary_html, compiled_model_dl, per_layer_csv = run_vela(model_file)
        
        return (
            summary_html,
            gr.update(visible=compiled_model_dl is not None, value=compiled_model_dl),
            gr.update(visible=per_layer_csv is not None, value=per_layer_csv)
        )

    compile_btn.click(
        wrapped_run_vela,
        inputs=[model_file],
        outputs=[summary_table, download_btn, download_logs_btn]
    )

# Launch the demo
if __name__ == "__main__":
    demo.launch(show_api=False)