"""Gradio app for the interactive machine-learning demo suite."""

from __future__ import annotations

import argparse
import os

import gradio as gr

from vibe_ml_lab.backprop_demo import run_backprop_demo_ui
from vibe_ml_lab.cnn_demo import predict_cnn_digit_ui, train_cnn_demo_ui
from vibe_ml_lab.hog_bow_svm import CLASS_NAMES, predict_hog_bow_ui, train_hog_bow_demo_ui
from vibe_ml_lab.resnet_compare import MODEL_SPECS, compare_resnets_ui

CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Manrope:wght@400;600;700;800&family=IBM+Plex+Mono:wght@400;500&display=swap');

body, .gradio-container {
  background:
    radial-gradient(circle at top right, rgba(239, 125, 50, 0.18), transparent 28%),
    radial-gradient(circle at top left, rgba(47, 107, 108, 0.18), transparent 26%),
    linear-gradient(180deg, #f9f4ec 0%, #f1ece4 100%);
  font-family: 'Manrope', 'Aptos', 'Segoe UI', sans-serif !important;
}

.gradio-container {
  max-width: 1280px !important;
}

.hero-shell {
  border-radius: 28px;
  padding: 30px 34px;
  background: linear-gradient(135deg, rgba(255,255,255,0.88), rgba(255,243,225,0.92));
  border: 1px solid rgba(191, 149, 111, 0.35);
  box-shadow: 0 18px 40px rgba(100, 73, 36, 0.12);
  margin-bottom: 14px;
}

.hero-shell h1 {
  margin: 8px 0 12px 0;
  font-size: 2.4rem;
  line-height: 1.04;
  color: #302117;
  letter-spacing: -0.04em;
}

.hero-shell p {
  color: #4c3b31;
  font-size: 1.05rem;
  line-height: 1.7;
}

.eyebrow {
  font-family: 'IBM Plex Mono', monospace;
  text-transform: uppercase;
  letter-spacing: 0.18em;
  color: #2f6b6c !important;
  font-size: 0.74rem !important;
}

.hero-grid {
  display: grid;
  grid-template-columns: repeat(4, minmax(0, 1fr));
  gap: 14px;
  margin-top: 22px;
}

.hero-card {
  background: rgba(255,255,255,0.74);
  border: 1px solid rgba(199, 160, 126, 0.34);
  border-radius: 20px;
  padding: 16px 18px;
}

.hero-card strong {
  display: block;
  color: #2d2017;
  font-size: 1.25rem;
  margin-bottom: 6px;
}

.hero-card span {
  color: #725846;
  font-size: 0.95rem;
}

.section-note {
  background: rgba(255,255,255,0.72);
  border: 1px solid rgba(203, 164, 130, 0.24);
  border-radius: 20px;
  padding: 14px 18px;
}

button.primary {
  background: linear-gradient(135deg, #ef7d32, #c95f1a) !important;
  border: none !important;
  color: white !important;
}

.tabs button {
  font-weight: 700 !important;
}

@media (max-width: 860px) {
  .hero-grid {
    grid-template-columns: repeat(2, minmax(0, 1fr));
  }
}
"""

APP_THEME = gr.themes.Soft(
    primary_hue="orange",
    secondary_hue="teal",
    neutral_hue="stone",
)


def build_interface() -> gr.Blocks:
    with gr.Blocks(title="Visual ML Playground") as demo:
        gr.HTML(
            """
            <div class="hero-shell">
              <div class="eyebrow">Vibe Coding Design</div>
              <h1>Visual ML Playground</h1>
              <p>
                一个面向展示与教学的机器学习交互应用，把传统视觉、反向传播、卷积神经网络、
                以及预训练 ResNet 模型对比放到同一个 Web 页面里。它既能解释算法，也能给别人直接点开体验。
              </p>
              <div class="hero-grid">
                <div class="hero-card"><strong>HOG + BoW + SVM</strong><span>传统图像分类流程可视化</span></div>
                <div class="hero-card"><strong>Backprop Demo</strong><span>从梯度到决策边界的动态演示</span></div>
                <div class="hero-card"><strong>LeNet-style CNN</strong><span>从零训练手写数字卷积网络</span></div>
                <div class="hero-card"><strong>ResNet Compare</strong><span>真实预训练模型测速与预测对比</span></div>
              </div>
            </div>
            """
        )

        with gr.Tabs():
            with gr.Tab("HOG + BoW + SVM"):
                gr.Markdown(
                    "使用合成几何图形数据集演示传统图像分类管线：先提取局部 HOG 描述子，"
                    "再通过 KMeans 构建视觉词袋，最后交给线性 SVM 完成分类。",
                    elem_classes=["section-note"],
                )
                hog_state = gr.State()

                with gr.Row():
                    with gr.Column(scale=1):
                        hog_train_per_class = gr.Slider(20, 60, value=36, step=4, label="每类训练样本数")
                        hog_test_per_class = gr.Slider(6, 24, value=12, step=2, label="每类测试样本数")
                        hog_vocab = gr.Slider(8, 32, value=18, step=2, label="视觉词袋大小")
                        hog_train_btn = gr.Button("训练并评估 HOG 管线", variant="primary")
                    with gr.Column(scale=2):
                        hog_summary = gr.Markdown()
                        with gr.Row():
                            hog_confusion = gr.Image(label="混淆矩阵", type="numpy")
                            hog_gallery = gr.Image(label="样本画廊", type="numpy")

                gr.Markdown("### 交互测试")
                with gr.Row():
                    with gr.Column(scale=1):
                        hog_input = gr.Image(label="上传待分类图像（可留空）", type="numpy", sources=["upload", "clipboard"])
                        hog_random_class = gr.Dropdown(
                            ["随机"] + CLASS_NAMES,
                            value="随机",
                            label="随机样本类别",
                        )
                        hog_predict_btn = gr.Button("预测上传图像 / 生成随机样本", variant="primary")
                    with gr.Column(scale=2):
                        with gr.Row():
                            hog_preview = gr.Image(label="输入预览", type="numpy")
                            hog_similarity = gr.Image(label="类别相似度", type="numpy")
                        hog_prediction = gr.Markdown()

                hog_train_btn.click(
                    train_hog_bow_demo_ui,
                    inputs=[hog_train_per_class, hog_test_per_class, hog_vocab],
                    outputs=[hog_state, hog_summary, hog_confusion, hog_gallery],
                )
                hog_predict_btn.click(
                    predict_hog_bow_ui,
                    inputs=[hog_state, hog_input, hog_random_class],
                    outputs=[hog_preview, hog_similarity, hog_prediction],
                )

            with gr.Tab("反向传播演示"):
                gr.Markdown(
                    "在一个带噪声的 XOR 二分类任务上，从零训练一个小型多层感知机，"
                    "展示反向传播如何推动损失下降、准确率提升，以及决策边界逐步成形。",
                    elem_classes=["section-note"],
                )
                with gr.Row():
                    with gr.Column(scale=1):
                        bp_hidden = gr.Slider(3, 16, value=8, step=1, label="隐藏层神经元数量")
                        bp_lr = gr.Slider(0.05, 0.8, value=0.35, step=0.05, label="学习率")
                        bp_epochs = gr.Slider(80, 800, value=320, step=40, label="训练轮数")
                        bp_noise = gr.Slider(0.05, 0.35, value=0.18, step=0.01, label="样本噪声强度")
                        bp_run_btn = gr.Button("运行反向传播演示", variant="primary")
                    with gr.Column(scale=2):
                        bp_summary = gr.Markdown()
                        with gr.Row():
                            bp_metrics = gr.Image(label="Loss / Accuracy 曲线", type="numpy")
                            bp_boundary = gr.Image(label="决策边界", type="numpy")
                        bp_table = gr.Dataframe(
                            headers=["参数", "首轮梯度范数", "末轮梯度范数"],
                            datatype=["str", "number", "number"],
                            interactive=False,
                            wrap=True,
                            row_count=(4, "fixed"),
                            column_count=(3, "fixed"),
                            label="梯度流观测",
                        )

                bp_run_btn.click(
                    run_backprop_demo_ui,
                    inputs=[bp_hidden, bp_lr, bp_epochs, bp_noise],
                    outputs=[bp_summary, bp_metrics, bp_boundary, bp_table],
                )

            with gr.Tab("CNN 训练与测试"):
                gr.Markdown(
                    "这一页训练一个纯 NumPy 实现的类 LeNet-5 卷积神经网络。"
                    "数据集由程序动态生成，适合用来展示卷积、池化、全连接层如何协同完成手写数字分类。",
                    elem_classes=["section-note"],
                )
                cnn_state = gr.State()

                with gr.Row():
                    with gr.Column(scale=1):
                        cnn_train_per_class = gr.Slider(25, 80, value=45, step=5, label="每类训练样本数")
                        cnn_test_per_class = gr.Slider(8, 30, value=15, step=1, label="每类测试样本数")
                        cnn_filters = gr.Slider(4, 10, value=6, step=1, label="卷积核数量")
                        cnn_hidden = gr.Slider(16, 64, value=32, step=8, label="全连接隐藏层大小")
                        cnn_epochs = gr.Slider(6, 24, value=14, step=2, label="训练轮数")
                        cnn_lr = gr.Slider(0.04, 0.2, value=0.12, step=0.01, label="学习率")
                        cnn_train_btn = gr.Button("训练 CNN 并评估", variant="primary")
                    with gr.Column(scale=2):
                        cnn_summary = gr.Markdown()
                        with gr.Row():
                            cnn_curves = gr.Image(label="训练曲线", type="numpy")
                            cnn_confusion = gr.Image(label="测试混淆矩阵", type="numpy")
                        with gr.Row():
                            cnn_filters_img = gr.Image(label="卷积核可视化", type="numpy")
                            cnn_gallery = gr.Image(label="测试样本预测", type="numpy")

                gr.Markdown("### 上传单数字图片测试")
                with gr.Row():
                    with gr.Column(scale=1):
                        cnn_input = gr.Image(label="上传一个单数字图片", type="numpy", sources=["upload", "clipboard"])
                        cnn_predict_btn = gr.Button("用当前 CNN 预测", variant="primary")
                    with gr.Column(scale=2):
                        cnn_preview = gr.Image(label="原图 / 预处理结果", type="numpy")
                        cnn_prediction = gr.Markdown()

                cnn_train_btn.click(
                    train_cnn_demo_ui,
                    inputs=[
                        cnn_train_per_class,
                        cnn_test_per_class,
                        cnn_filters,
                        cnn_hidden,
                        cnn_epochs,
                        cnn_lr,
                    ],
                    outputs=[cnn_state, cnn_summary, cnn_curves, cnn_confusion, cnn_filters_img, cnn_gallery],
                )
                cnn_predict_btn.click(
                    predict_cnn_digit_ui,
                    inputs=[cnn_state, cnn_input],
                    outputs=[cnn_preview, cnn_prediction],
                )

            with gr.Tab("预训练 ResNet 对比"):
                gr.Markdown(
                    "这里使用真实的 ImageNet 预训练 ONNX 权重，对不同深度的 ResNet 在同一张图片上进行推理与测速。"
                    "第一次运行会自动下载模型文件。",
                    elem_classes=["section-note"],
                )
                with gr.Row():
                    with gr.Column(scale=1):
                        resnet_input = gr.Image(
                            label="上传常见物体照片",
                            type="numpy",
                            sources=["upload", "clipboard"],
                        )
                        resnet_models = gr.CheckboxGroup(
                            choices=list(MODEL_SPECS.keys()),
                            value=list(MODEL_SPECS.keys()),
                            label="参与比较的模型",
                        )
                        resnet_compare_btn = gr.Button("开始比较预训练 ResNet", variant="primary")
                    with gr.Column(scale=2):
                        resnet_summary = gr.Markdown()
                        resnet_table = gr.Dataframe(
                            headers=["模型", "深度", "模型大小(MB)", "中位延迟(ms)", "Top-1", "Top-3"],
                            datatype=["str", "number", "number", "number", "str", "str"],
                            interactive=False,
                            wrap=True,
                            label="模型比较结果",
                        )
                        resnet_chart = gr.Image(label="延迟 / 模型体积对比", type="numpy")

                resnet_compare_btn.click(
                    compare_resnets_ui,
                    inputs=[resnet_input, resnet_models],
                    outputs=[resnet_summary, resnet_table, resnet_chart],
                )

    demo.queue(default_concurrency_limit=2)
    return demo


def main() -> None:
    parser = argparse.ArgumentParser(description="Launch the Visual ML Playground app.")
    parser.add_argument("--share", action="store_true", help="Create a Gradio public sharing link.")
    parser.add_argument(
        "--port",
        type=int,
        default=int(os.environ.get("PORT", "7860")),
        help="Server port.",
    )
    parser.add_argument("--no-browser", action="store_true", help="Do not open the browser automatically.")
    args = parser.parse_args()

    demo = build_interface()
    demo.launch(
        server_name=os.environ.get("HOST", "0.0.0.0"),
        server_port=args.port,
        share=args.share,
        inbrowser=not args.no_browser,
        theme=APP_THEME,
        css=CUSTOM_CSS,
    )


if __name__ == "__main__":
    main()