"""
app.py  —  Gradio demo for bert-cpu-benchmark
"""

import gradio as gr
from benchmark import load_model, profile_flops, benchmark_latency

# ── Defaults ──────────────────────────────────────────────────────────────────

DEFAULT_MODEL     = "katrjohn/TinyGreekNewsBERT"
DEFAULT_TOKENIZER = "nlpaueb/bert-base-greek-uncased-v1"
DEFAULT_TEXT      = "The government announced new support measures for workers today."

# ── Core function ─────────────────────────────────────────────────────────────

def run_benchmark(model_id, tokenizer_id, trust_remote_code, runs, seq_len, sample_text):
    model_id     = model_id.strip()
    tokenizer_id = tokenizer_id.strip() or model_id
    sample_text  = sample_text.strip() or DEFAULT_TEXT

    if not model_id:
        return "⚠️ Please enter a model ID.", "", "", "", "", "", ""

    try:
        model, tokenizer = load_model(
            model_id=model_id,
            tokenizer_id=tokenizer_id,
            trust_remote_code=trust_remote_code,
        )
    except Exception as e:
        return f"⚠️ Failed to load model: {e}", "", "", "", "", "", ""

    try:
        flops_data = profile_flops(model, tokenizer, seq_len=seq_len)
    except Exception as e:
        return f"⚠️ FLOPs profiling failed: {e}", "", "", "", "", "", ""

    try:
        latency_data = benchmark_latency(
            model, tokenizer,
            text=sample_text,
            warm=20,
            runs=runs,
        )
    except Exception as e:
        return f"⚠️ Latency benchmark failed: {e}", "", "", "", "", "", ""

    model_size_mb = (
        sum(param.numel() * param.element_size() for param in model.parameters()) +
        sum(buffer.numel() * buffer.element_size() for buffer in model.buffers())
    ) / (1024 ** 2)
    params  = f"{flops_data['params'] / 1e6:.1f} M"
    size_mb = f"{model_size_mb:.2f} MB"
    macs    = f"{flops_data['macs']   / 1e9:.2f} GMac"
    flops   = f"{flops_data['flops']  / 1e9:.2f} GFLOPs"
    mean_ms = f"{latency_data['mean_ms']:.2f} ms"
    p95_ms  = f"{latency_data['p95_ms']:.2f} ms"

    return "", params, size_mb, macs, flops, mean_ms, p95_ms


# ── CSS ───────────────────────────────────────────────────────────────────────

CSS = """
.result-box textarea {
    font-family: 'Courier New', monospace !important;
    font-size: 1.1rem !important;
    text-align: center !important;
    font-weight: 600 !important;
}
"""

# ── UI ────────────────────────────────────────────────────────────────────────

with gr.Blocks(theme=gr.themes.Soft(), title="BERT CPU Benchmark", css=CSS) as demo:

    gr.Markdown("""
    # BERT CPU Benchmark
    **CPU inference profiler for any BERT-family encoder model on Hugging Face**

    Measures parameters, MACs, FLOPs, mean latency, and p95 latency — all on CPU, no GPU required.
    > Compatible with any `AutoModel`-loadable encoder model: BERT, RoBERTa, DeBERTa, ELECTRA, DistilBERT, and custom distilled models.
    """)

    with gr.Row():
        with gr.Column(scale=2):
            model_input = gr.Textbox(
                label="Model ID",
                value=DEFAULT_MODEL,
                placeholder="e.g. bert-base-uncased",
            )
            tokenizer_input = gr.Textbox(
                label="Tokenizer ID",
                value=DEFAULT_TOKENIZER,
                placeholder="Leave blank to use the same as Model ID",
            )
            trust_remote = gr.Checkbox(
                label="trust_remote_code  (required for custom architectures)",
                value=True,
            )
            sample_text = gr.Textbox(
                label="Sample text for latency benchmark",
                value=DEFAULT_TEXT,
                lines=2,
            )
            with gr.Row():
                runs_input   = gr.Slider(minimum=10, maximum=500, value=100, step=10,
                                         label="Latency runs")
                seqlen_input = gr.Slider(minimum=64, maximum=512, value=512, step=64,
                                         label="Sequence length")
            run_btn = gr.Button("⚡ Run Benchmark", variant="primary", size="lg")
            warning_box = gr.Textbox(label="", show_label=False, interactive=False,
                                     container=False, visible=True)

        with gr.Column(scale=3):
            gr.Markdown("### 📐 Model Complexity")
            with gr.Row():
                out_params = gr.Textbox(label="Parameters", interactive=False,
                                        elem_classes="result-box")
                out_size   = gr.Textbox(label="Model size", interactive=False,
                                        elem_classes="result-box")
                out_macs   = gr.Textbox(label="MACs", interactive=False,
                                        elem_classes="result-box")
                out_flops  = gr.Textbox(label="FLOPs  (2 × MACs)", interactive=False,
                                        elem_classes="result-box")

            gr.Markdown("### ⏱️ CPU Latency")
            with gr.Row():
                out_mean = gr.Textbox(label="Mean latency", interactive=False,
                                      elem_classes="result-box")
                out_p95  = gr.Textbox(label="p95 latency", interactive=False,
                                      elem_classes="result-box")

    gr.Markdown("""
    ---
    > **FLOPs** are hardware-agnostic — they measure the model's computational cost, not the machine's speed.
    > **Latency** is measured with `torch.inference_mode()` after 20 warm-up passes to avoid cold-start bias.
    """)

    gr.Examples(
        label="Try an example",
        examples=[
            ["katrjohn/TinyGreekNewsBERT",       "nlpaueb/bert-base-greek-uncased-v1", True,  100, 512, "Η κυβέρνηση ανακοίνωσε νέα μέτρα στήριξης."],
            ["bert-base-uncased",                 "",                                   False, 100, 512, "The government announced new support measures."],
            ["FacebookAI/xlm-roberta-base",       "",                                   False, 100, 512, "The government announced new support measures."],
            ["microsoft/deberta-v3-base",         "",                                   False, 100, 512, "The government announced new support measures."],
            ["distilbert-base-uncased",           "",                                   False, 100, 512, "The government announced new support measures."],
        ],
        inputs=[model_input, tokenizer_input, trust_remote, runs_input, seqlen_input, sample_text],
    )

    run_btn.click(
        fn=run_benchmark,
        inputs=[model_input, tokenizer_input, trust_remote, runs_input, seqlen_input, sample_text],
        outputs=[warning_box, out_params, out_size, out_macs, out_flops, out_mean, out_p95],
    )

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