"""
Insurance Pricing Model Training Frontend
A Gradio-based web interface for configuring and running insurance pricing models.
"""

import inspect
import json
import os

from ins_pricing.frontend.app_controller import PricingApp, _check_frontend_deps


def create_ui():
    """Create the Gradio interface."""
    _check_frontend_deps()
    import gradio as gr

    app = PricingApp()

    def _dump_json_template(payload):
        return json.dumps(payload, indent=2, ensure_ascii=False)

    xgb_search_space_template = _dump_json_template(
        app.config_builder._default_xgb_search_space()
    )
    resn_search_space_template = _dump_json_template(
        app.config_builder._default_resn_search_space()
    )
    ft_search_space_template = _dump_json_template(
        app.config_builder._default_ft_search_space()
    )
    ft_unsupervised_search_space_template = _dump_json_template(
        app.config_builder._default_ft_unsupervised_search_space()
    )
    workflow_template = _dump_json_template(
        {
            "workflow": {
                "mode": "plot_direct",
                "cfg_path": "config_plot.json",
                "xgb_cfg_path": "config_xgb_direct.json",
                "resn_cfg_path": "config_resn_direct.json",
            }
        }
    )
    xgb_step2_overrides_template = _dump_json_template(
        {
            "output_dir": "./ResultsXGBFromFTUnsupervised",
            "optuna_storage": "./ResultsXGBFromFTUnsupervised/optuna/bayesopt.sqlite3",
            "optuna_study_prefix": "pricing_ft_unsup_xgb",
            "loss_name": "mse",
            "build_oht": False,
            "final_refit": False,
            "runner": {
                "model_keys": ["xgb"],
                "nproc_per_node": 1,
                "plot_curves": False,
            },
            "plot_curves": False,
            "plot": {"enable": False},
        }
    )
    resn_step2_overrides_template = _dump_json_template(
        {
            "use_resn_ddp": True,
            "output_dir": "./ResultsResNFromFTUnsupervised",
            "optuna_storage": "./ResultsResNFromFTUnsupervised/optuna/bayesopt.sqlite3",
            "optuna_study_prefix": "pricing_ft_unsup_resn_ddp",
            "loss_name": "mse",
            "build_oht": True,
            "runner": {
                "model_keys": ["resn"],
                "nproc_per_node": 2,
                "plot_curves": False,
            },
            "plot_curves": False,
            "plot": {"enable": False},
        }
    )
    dir_status_init, dir_choices_init, dir_value_init = app.list_directory_candidates(
        str(app.working_dir)
    )

    def _initial_oneway_factor_state():
        return app.suggest_oneway_factors("config_plot.json")

    pre_factor_status_init, pre_factor_choices_init, pre_factor_value_init = _initial_oneway_factor_state()
    direct_factor_status_init, direct_factor_choices_init, direct_factor_value_init = _initial_oneway_factor_state()
    embed_factor_status_init, embed_factor_choices_init, embed_factor_value_init = _initial_oneway_factor_state()

    def _set_working_dir_ui(path_text: str):
        status, resolved = app.set_working_dir(path_text)
        _, choices, selected = app.list_directory_candidates(resolved)
        return (
            status,
            resolved,
            resolved,
            gr.update(choices=choices, value=selected),
        )

    def _refresh_working_dir_choices_ui(root_dir: str):
        status, choices, selected = app.list_directory_candidates(root_dir)
        return status, gr.update(choices=choices, value=selected)

    def _suggest_compare_defaults(model_key: str):
        key = str(model_key or "").strip().lower()
        if key == "resn":
            return (
                "config_resn_direct.json",
                "config_resn_from_ft_unsupervised.json",
                "ResN_raw",
                "ResN_ft_embed",
            )
        return (
            "config_xgb_direct.json",
            "config_xgb_from_ft_unsupervised.json",
            "XGB_raw",
            "XGB_ft_embed",
        )

    def _load_oneway_factors_ui(cfg_path: str):
        status, choices, selected = app.suggest_oneway_factors(cfg_path)
        return status, gr.update(choices=choices, value=selected)

    layout_css = """
    .gradio-container {
        max-width: 1480px !important;
        margin: 0 auto !important;
    }
    .gradio-container .tabitem {
        padding-top: 8px;
    }
    .gradio-container .gr-form {
        gap: 6px !important;
    }
    .gradio-container .gr-row {
        align-items: stretch;
    }
    .gradio-container .gr-column {
        min-width: 0;
    }
    .gradio-container textarea {
        line-height: 1.35;
    }
    .gradio-container .gr-textbox,
    .gradio-container .gr-number,
    .gradio-container .gr-dropdown,
    .gradio-container .gr-slider,
    .gradio-container .gr-checkbox {
        margin-bottom: 2px !important;
    }
    .section-header {
        margin: 4px 0 2px 0 !important;
        font-size: 0.95em;
        color: #555;
    }
    """

    with gr.Blocks(
        title="Insurance Pricing Model Training",
        theme=gr.themes.Soft(),
        css=layout_css,
    ) as demo:
        gr.Markdown("# Insurance Pricing Model Training Interface")

        # ── Working Directory ──
        with gr.Accordion("Working Directory", open=False):
            with gr.Row():
                working_dir_input = gr.Textbox(
                    label="Working Directory",
                    value=str(app.working_dir),
                    placeholder="Type a path, or select from the folder list below",
                    scale=3,
                )
                set_working_dir_btn = gr.Button(
                    "Set Working Directory", variant="secondary", scale=1)
            with gr.Row():
                working_dir_browse_root = gr.Textbox(
                    label="Browse Root",
                    value=str(app.working_dir),
                    placeholder="List folders under this path (depth=2)",
                    scale=3,
                )
                refresh_working_dir_btn = gr.Button(
                    "Refresh Folder List", variant="secondary", scale=1
                )
            with gr.Row():
                working_dir_picker = gr.Dropdown(
                    label="Select Existing Folder",
                    choices=dir_choices_init,
                    value=dir_value_init,
                    scale=3,
                )
                use_selected_working_dir_btn = gr.Button(
                    "Use Selected Folder", variant="secondary", scale=1
                )
            working_dir_status = gr.Textbox(
                label="Working Directory Status", value=dir_status_init, interactive=False)

        # ════════════════════════════════════════════════════════════════
        #  Tab: Configuration
        # ════════════════════════════════════════════════════════════════
        with gr.Tab("Configuration"):
            # ── Load Config ──
            with gr.Row():
                with gr.Column(scale=1):
                    json_file = gr.File(
                        label="Upload JSON Config",
                        file_types=[".json"],
                        type="filepath",
                    )
                    load_btn = gr.Button("Load Config", variant="primary")
                    load_status = gr.Textbox(label="Load Status", interactive=False)
                with gr.Column(scale=3):
                    config_display = gr.JSON(label="Current Configuration", value={})

            gr.Markdown("---")

            # ── Core: Data Settings ──
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("#### Data Settings")
                    data_dir = gr.Textbox(label="Data Directory", value="./Data")
                    model_list = gr.Textbox(label="Model List (comma-separated)", value="od")
                    model_categories = gr.Textbox(label="Model Categories (comma-separated)", value="bc")
                    target = gr.Textbox(label="Target Column", value="response")
                    weight = gr.Textbox(label="Weight Column", value="weights")
                    feature_list = gr.Textbox(
                        label="Feature List (comma-separated)",
                        placeholder="feature_1, feature_2, feature_3",
                        lines=4,
                    )
                    categorical_features = gr.Textbox(
                        label="Categorical Features (comma-separated)",
                        placeholder="feature_2, feature_3",
                        lines=3,
                    )

                with gr.Column(scale=1):
                    gr.Markdown("#### Task & Training")
                    with gr.Row():
                        task_type = gr.Dropdown(
                            label="Task Type",
                            choices=["regression", "binary", "multiclass"],
                            value="regression",
                        )
                        distribution = gr.Dropdown(
                            label="Distribution",
                            choices=["", "tweedie", "poisson", "gamma", "gaussian", "mse", "mae"],
                            value="",
                            allow_custom_value=True,
                        )
                    binary_resp_nme = gr.Textbox(
                        label="Binary Response Column (optional)",
                        value="",
                        placeholder="e.g. claim_flag",
                    )
                    with gr.Row():
                        split_strategy = gr.Dropdown(
                            label="Split Strategy",
                            choices=["random", "stratified", "time", "group"],
                            value="random",
                        )
                        rand_seed = gr.Number(label="Random Seed", value=13, precision=0)
                    with gr.Row():
                        epochs = gr.Number(label="Epochs", value=50, precision=0)
                        max_evals = gr.Number(label="Max Evaluations", value=50, precision=0)
                    with gr.Row():
                        prop_test = gr.Slider(
                            label="Test Proportion", minimum=0.05, maximum=0.5, value=0.25, step=0.05)
                        holdout_ratio = gr.Slider(
                            label="Holdout Ratio", minimum=0.05, maximum=0.5, value=0.25, step=0.05)
                    val_ratio = gr.Slider(
                        label="Validation Ratio", minimum=0.05, maximum=0.5, value=0.25, step=0.05)
                    output_dir = gr.Textbox(label="Output Directory", value="./Results")
                    with gr.Row():
                        use_gpu = gr.Checkbox(label="Use GPU", value=True)
                        plot_curves = gr.Checkbox(label="Plot Curves", value=False)
                    model_keys = gr.Textbox(
                        label="Model Keys (comma-separated)",
                        value="xgb, resn",
                        placeholder="xgb, resn, ft, gnn",
                    )

            # ── Split & Pre-split Data ──
            with gr.Accordion("Split & Pre-split Data", open=False):
                with gr.Row():
                    split_group_col = gr.Textbox(label="Split Group Column", value="")
                    split_time_col = gr.Textbox(label="Split Time Column", value="")
                    split_time_ascending = gr.Checkbox(label="Time Ascending", value=True)
                with gr.Row():
                    train_data_path = gr.Textbox(label="Train Data Path (pre-split)", value="")
                    test_data_path = gr.Textbox(label="Test Data Path (pre-split)", value="")
                with gr.Row():
                    split_cache_path = gr.Textbox(
                        label="Split Cache Path (.npz)", value="", placeholder="e.g. ./cache/split.npz")
                    split_cache_force_rebuild = gr.Checkbox(label="Force Rebuild Split Cache", value=False)

            # ── Cross-Validation ──
            with gr.Accordion("Cross-Validation", open=False):
                with gr.Row():
                    cv_strategy = gr.Dropdown(
                        label="CV Strategy",
                        choices=["", "random", "group", "grouped", "time", "stratified", "timeseries", "temporal"],
                        value="",
                        allow_custom_value=True,
                    )
                    cv_splits = gr.Number(label="CV Splits", value=0, precision=0)
                with gr.Row():
                    cv_group_col = gr.Textbox(label="CV Group Column", value="")
                    cv_time_col = gr.Textbox(label="CV Time Column", value="")
                    cv_time_ascending = gr.Checkbox(label="CV Time Ascending", value=True)

            # ── XGBoost Settings ──
            with gr.Accordion("XGBoost Settings", open=False):
                with gr.Row():
                    xgb_max_depth_max = gr.Number(label="Max Depth", value=25, precision=0)
                    xgb_n_estimators_max = gr.Number(label="Max Estimators", value=500, precision=0)
                    xgb_gpu_id = gr.Number(label="GPU ID", value=0, precision=0)
                with gr.Row():
                    xgb_use_dmatrix = gr.Checkbox(label="Use DMatrix", value=True)
                    xgb_chunk_size = gr.Number(label="Chunk Size (0=off)", value=0, precision=0)
                    xgb_cleanup_per_fold = gr.Checkbox(label="Cleanup Per Fold", value=False)
                    xgb_cleanup_synchronize = gr.Checkbox(label="Cleanup Synchronize", value=False)
                xgb_search_space_json = gr.Textbox(
                    label="Search Space (JSON)",
                    value=xgb_search_space_template,
                    lines=8,
                    max_lines=16,
                )

            # ── ResNet Settings ──
            with gr.Accordion("ResNet Settings", open=False):
                with gr.Row():
                    resn_use_lazy_dataset = gr.Checkbox(label="Lazy Dataset", value=True)
                    resn_predict_batch_size = gr.Number(
                        label="Predict Batch Size (0=auto)", value=0, precision=0)
                    resn_weight_decay = gr.Number(label="Weight Decay", value=0.0001)
                with gr.Row():
                    resn_cleanup_per_fold = gr.Checkbox(label="Cleanup Per Fold", value=False)
                    resn_cleanup_synchronize = gr.Checkbox(label="Cleanup Synchronize", value=False)
                resn_search_space_json = gr.Textbox(
                    label="Search Space (JSON)",
                    value=resn_search_space_template,
                    lines=8,
                    max_lines=16,
                )

            # ── FT-Transformer Settings ──
            with gr.Accordion("FT-Transformer Settings", open=False):
                with gr.Row():
                    ft_role = gr.Dropdown(
                        label="FT Role",
                        choices=["model", "embedding", "unsupervised_embedding"],
                        value="model",
                    )
                    ft_feature_prefix = gr.Textbox(label="Feature Prefix", value="ft_emb")
                    ft_num_numeric_tokens = gr.Number(
                        label="Num Numeric Tokens (0=auto)", value=0, precision=0)
                with gr.Row():
                    ft_use_lazy_dataset = gr.Checkbox(label="Lazy Dataset", value=True)
                    ft_predict_batch_size = gr.Number(
                        label="Predict Batch Size (0=auto)", value=0, precision=0)
                with gr.Row():
                    ft_oof_folds = gr.Number(label="OOF Folds (0=off)", value=0, precision=0)
                    ft_oof_strategy = gr.Dropdown(
                        label="OOF Strategy",
                        choices=["", "random", "group", "time", "stratified"],
                        value="",
                        allow_custom_value=True,
                    )
                    ft_oof_shuffle = gr.Checkbox(label="OOF Shuffle", value=True)
                with gr.Row():
                    ft_cleanup_per_fold = gr.Checkbox(label="Cleanup Per Fold", value=False)
                    ft_cleanup_synchronize = gr.Checkbox(label="Cleanup Synchronize", value=False)
                with gr.Row():
                    with gr.Column():
                        ft_search_space_json = gr.Textbox(
                            label="Supervised Search Space (JSON)",
                            value=ft_search_space_template,
                            lines=8,
                            max_lines=16,
                        )
                    with gr.Column():
                        ft_unsupervised_search_space_json = gr.Textbox(
                            label="Unsupervised Search Space (JSON)",
                            value=ft_unsupervised_search_space_template,
                            lines=8,
                            max_lines=16,
                        )

            # ── GNN Settings ──
            with gr.Accordion("GNN Settings", open=False):
                with gr.Row():
                    gnn_use_approx_knn = gr.Checkbox(label="Approx KNN", value=True)
                    gnn_approx_knn_threshold = gr.Number(label="Approx KNN Threshold", value=50000, precision=0)
                    gnn_max_gpu_knn_nodes = gr.Number(label="Max GPU KNN Nodes", value=200000, precision=0)
                with gr.Row():
                    gnn_knn_gpu_mem_ratio = gr.Slider(
                        label="KNN GPU Mem Ratio", minimum=0.1, maximum=1.0, value=0.9, step=0.05)
                    gnn_knn_gpu_mem_overhead = gr.Number(label="KNN GPU Mem Overhead (GB)", value=2.0)
                with gr.Row():
                    gnn_max_fit_rows = gr.Number(label="Max Fit Rows (0=unlimited)", value=0, precision=0)
                    gnn_max_predict_rows = gr.Number(label="Max Predict Rows (0=unlimited)", value=0, precision=0)
                    gnn_predict_chunk_rows = gr.Number(label="Predict Chunk Rows (0=auto)", value=0, precision=0)
                with gr.Row():
                    gnn_graph_cache = gr.Textbox(label="Graph Cache Path", value="")
                    gnn_cleanup_per_fold = gr.Checkbox(label="Cleanup Per Fold", value=False)
                    gnn_cleanup_synchronize = gr.Checkbox(label="Cleanup Synchronize", value=False)

            # ── Distributed Training ──
            with gr.Accordion("Distributed Training", open=False):
                with gr.Row():
                    nproc_per_node = gr.Number(label="Processes Per Node", value=2, precision=0)
                    ddp_min_rows = gr.Number(label="DDP Min Rows", value=50000, precision=0)
                with gr.Row():
                    use_resn_ddp = gr.Checkbox(label="ResNet DDP", value=True)
                    use_ft_ddp = gr.Checkbox(label="FT DDP", value=True)
                with gr.Row():
                    use_resn_data_parallel = gr.Checkbox(label="ResNet DataParallel", value=False)
                    use_ft_data_parallel = gr.Checkbox(label="FT DataParallel", value=False)
                    use_gnn_data_parallel = gr.Checkbox(label="GNN DataParallel", value=False)

            # ── Preprocessing ──
            with gr.Accordion("Preprocessing", open=False):
                with gr.Row():
                    build_oht = gr.Checkbox(label="Build OHT", value=True)
                    oht_sparse_csr = gr.Checkbox(label="OHT Sparse CSR", value=True)
                    keep_unscaled_oht = gr.Checkbox(label="Keep Unscaled OHT", value=False)
                with gr.Row():
                    infer_categorical_max_unique = gr.Number(
                        label="Infer Categorical Max Unique", value=50, precision=0)
                    infer_categorical_max_ratio = gr.Number(
                        label="Infer Categorical Max Ratio", value=0.05)

            # ── Geographic / Regional ──
            with gr.Accordion("Geographic / Regional", open=False):
                with gr.Row():
                    geo_feature_nmes = gr.Textbox(
                        label="Geo Feature Names (comma-separated)", value="")
                    region_province_col = gr.Textbox(label="Province Column", value="")
                    region_city_col = gr.Textbox(label="City Column", value="")
                    region_effect_alpha = gr.Number(label="Region Effect Alpha", value=0.0)
                with gr.Row():
                    geo_token_hidden_dim = gr.Number(label="Hidden Dim", value=32, precision=0)
                    geo_token_layers = gr.Number(label="Layers", value=2, precision=0)
                    geo_token_dropout = gr.Number(label="Dropout", value=0.1)
                with gr.Row():
                    geo_token_k_neighbors = gr.Number(label="K Neighbors", value=10, precision=0)
                    geo_token_learning_rate = gr.Number(label="Learning Rate", value=0.001)
                    geo_token_epochs = gr.Number(label="Epochs", value=50, precision=0)

            # ── Ensemble & Refit ──
            with gr.Accordion("Ensemble & Refit", open=False):
                with gr.Row():
                    final_ensemble = gr.Checkbox(label="Final Ensemble", value=False)
                    final_ensemble_k = gr.Number(label="Ensemble K", value=3, precision=0)
                    final_refit = gr.Checkbox(label="Final Refit", value=True)
                    reuse_best_params = gr.Checkbox(label="Reuse Best Params", value=False)

            # ── Output & Caching ──
            with gr.Accordion("Output & Caching", open=False):
                with gr.Row():
                    optuna_study_prefix = gr.Textbox(label="Optuna Study Prefix", value="pricing")
                    optuna_cleanup_synchronize = gr.Checkbox(
                        label="Optuna Cleanup Synchronize", value=False)
                with gr.Row():
                    cache_predictions = gr.Checkbox(label="Cache Predictions", value=False)
                    prediction_cache_format = gr.Dropdown(
                        label="Cache Format",
                        choices=["parquet", "csv"],
                        value="parquet",
                    )
                with gr.Row():
                    dataloader_workers = gr.Number(label="DataLoader Workers", value=0, precision=0)
                    bo_sample_limit = gr.Number(label="BO Sample Limit (0=unlimited)", value=0, precision=0)

            # ── Plot Settings ──
            with gr.Accordion("Plot Settings", open=False):
                with gr.Row():
                    plot_enable = gr.Checkbox(label="Enable Plot", value=False)
                    plot_n_bins = gr.Number(label="Plot Bins", value=10, precision=0)
                with gr.Row():
                    plot_oneway = gr.Checkbox(label="Oneway", value=False)
                    plot_oneway_pred = gr.Checkbox(label="Oneway Pred", value=False)
                    plot_pre_oneway = gr.Checkbox(label="Pre Oneway", value=False)
                    plot_double_lift = gr.Checkbox(label="Double Lift", value=False)

            # ── Calibration ──
            with gr.Accordion("Calibration", open=False):
                with gr.Row():
                    calibration_enable = gr.Checkbox(label="Enable Calibration", value=False)
                    calibration_method = gr.Dropdown(
                        label="Method",
                        choices=["sigmoid", "isotonic"],
                        value="sigmoid",
                    )
                    calibration_max_rows = gr.Number(label="Max Rows (0=all)", value=0, precision=0)
                    calibration_seed = gr.Number(label="Seed", value=13, precision=0)

            # ── Threshold ──
            with gr.Accordion("Threshold", open=False):
                with gr.Row():
                    threshold_enable = gr.Checkbox(label="Enable Threshold", value=False)
                    threshold_metric = gr.Dropdown(
                        label="Metric",
                        choices=["f1", "precision", "recall", "balanced_accuracy"],
                        value="f1",
                        allow_custom_value=True,
                    )
                    threshold_grid = gr.Number(label="Grid", value=99, precision=0)
                with gr.Row():
                    threshold_max_rows = gr.Number(label="Max Rows (0=all)", value=0, precision=0)
                    threshold_seed = gr.Number(label="Seed", value=13, precision=0)

            # ── Bootstrap ──
            with gr.Accordion("Bootstrap", open=False):
                with gr.Row():
                    bootstrap_enable = gr.Checkbox(label="Enable Bootstrap", value=False)
                    bootstrap_n_samples = gr.Number(label="N Samples", value=200, precision=0)
                    bootstrap_ci = gr.Number(label="CI", value=0.95)
                    bootstrap_seed = gr.Number(label="Seed", value=13, precision=0)

            # ── Advanced Overrides ──
            with gr.Accordion("Advanced Manual Overrides (JSON Deep-Merge)", open=False):
                gr.Markdown(
                    "Override any generated config field via JSON deep-merge, including "
                    "`runner`, `plot`, `env`, `report_*`, `psi_*`, `registry_*`, etc."
                )
                config_overrides_json = gr.Textbox(
                    label="Config Overrides JSON",
                    value="{}",
                    lines=8,
                    max_lines=20,
                    placeholder='{"runner":{"mode":"entry","nproc_per_node":1},"plot":{"enable":true}}',
                )

            # ── Build & Save ──
            with gr.Row():
                build_btn = gr.Button("Build Configuration", variant="primary", size="lg")
                save_config_btn = gr.Button("Save Configuration", variant="secondary", size="lg")

            build_status = gr.Textbox(label="Status", interactive=False)
            config_json = gr.Textbox(label="Generated Config (JSON)", lines=12, max_lines=24)

            with gr.Row():
                save_filename = gr.Textbox(
                    label="Save Filename", value="my_config.json", scale=3)
                save_status = gr.Textbox(
                    label="Save Status", interactive=False, scale=4)

        # ════════════════════════════════════════════════════════════════
        #  Tab: Run Task
        # ════════════════════════════════════════════════════════════════
        with gr.Tab("Run Task"):
            gr.Markdown(
                "### Run Model Task\n"
                "Task type is detected from `config.runner.mode`: "
                "**entry** (training), **explain**, **incremental**, **watchdog**."
            )

            with gr.Row():
                run_btn = gr.Button("Run Task", variant="primary", size="lg")
                run_status = gr.Textbox(label="Task Status", interactive=False)

            log_output = gr.Textbox(
                label="Task Logs",
                lines=25,
                max_lines=50,
                interactive=False,
                autoscroll=True,
            )

            with gr.Row():
                open_folder_btn = gr.Button("Open Results Folder", size="lg")
                folder_status = gr.Textbox(label="Status", interactive=False, scale=2)

        # ════════════════════════════════════════════════════════════════
        #  Tab: FT Two-Step Workflow
        # ════════════════════════════════════════════════════════════════
        with gr.Tab("FT Two-Step Workflow"):
            gr.Markdown(
                "### FT-Transformer Two-Step Training\n"
                "1. **Step 1**: Train FT as unsupervised embedding generator\n"
                "2. **Step 2**: Merge embeddings with raw data and train XGB/ResN"
            )

            with gr.Row():
                with gr.Column():
                    gr.Markdown("#### Step 1: FT Embedding Generation")
                    with gr.Row():
                        ft_use_ddp = gr.Checkbox(label="Use DDP for FT", value=True)
                        ft_nproc = gr.Number(
                            label="Number of Processes (DDP)", value=2, precision=0)
                    prepare_step1_btn = gr.Button("Prepare Step 1 Config", variant="primary")
                    step1_status = gr.Textbox(label="Status", interactive=False)
                    step1_config_display = gr.Textbox(
                        label="Step 1 Config (FT Embedding)", lines=15, max_lines=25)

                with gr.Column():
                    gr.Markdown("#### Step 2: Train XGB/ResN with Embeddings")
                    with gr.Row():
                        target_models_input = gr.Textbox(
                            label="Target Models (comma-separated)",
                            value="xgb, resn",
                        )
                        augmented_data_dir_input = gr.Textbox(
                            label="Augmented Data Directory",
                            value="./DataFTUnsupervised",
                        )
                    xgb_overrides_input = gr.Textbox(
                        label="XGB Step 2 Overrides (JSON)",
                        value=xgb_step2_overrides_template,
                        lines=8,
                        max_lines=16,
                    )
                    resn_overrides_input = gr.Textbox(
                        label="ResN Step 2 Overrides (JSON)",
                        value=resn_step2_overrides_template,
                        lines=8,
                        max_lines=16,
                    )
                    prepare_step2_btn = gr.Button("Prepare Step 2 Configs", variant="primary")
                    step2_status = gr.Textbox(label="Status", interactive=False)
                    with gr.Tab("XGB Config"):
                        xgb_config_display = gr.Textbox(
                            label="XGB Step 2 Config", lines=15, max_lines=25)
                    with gr.Tab("ResN Config"):
                        resn_config_display = gr.Textbox(
                            label="ResN Step 2 Config", lines=15, max_lines=25)

        # ════════════════════════════════════════════════════════════════
        #  Tab: Workflow Config
        # ════════════════════════════════════════════════════════════════
        with gr.Tab("Workflow Config"):
            gr.Markdown(
                "### Config-Driven Workflows\n"
                "Supported `workflow.mode`: `pre_oneway`, `plot_direct`, `plot_embed`, "
                "`predict_ft_embed`, `compare_xgb`, `compare_resn`, `compare`, `double_lift`"
            )
            with gr.Row():
                with gr.Column(scale=1):
                    workflow_file = gr.File(
                        label="Upload Workflow Config",
                        file_types=[".json"],
                        type="filepath",
                    )
                    workflow_load_btn = gr.Button("Load Workflow Config", variant="secondary")
                    workflow_load_status = gr.Textbox(label="Load Status", interactive=False)
                with gr.Column(scale=2):
                    workflow_config_json = gr.Textbox(
                        label="Workflow Config (JSON)",
                        value=workflow_template,
                        lines=16,
                        max_lines=28,
                    )
            workflow_run_btn = gr.Button("Run Workflow Config", variant="primary", size="lg")
            workflow_status = gr.Textbox(label="Workflow Status", interactive=False)
            workflow_log = gr.Textbox(
                label="Workflow Logs",
                lines=18,
                max_lines=40,
                interactive=False,
                autoscroll=True,
            )

        # ════════════════════════════════════════════════════════════════
        #  Tab: Plotting
        # ════════════════════════════════════════════════════════════════
        with gr.Tab("Plotting"):
            with gr.Tab("Pre Oneway"):
                with gr.Row():
                    with gr.Column(scale=3):
                        with gr.Row():
                            pre_data_path = gr.Textbox(
                                label="Data Path", value="./Data/od_bc.csv", scale=3)
                            pre_output_dir = gr.Textbox(
                                label="Output Dir (optional)", value="", scale=2)
                        with gr.Row():
                            pre_model_name = gr.Textbox(label="Model Name", value="od_bc")
                            pre_target = gr.Textbox(label="Target Column", value="response")
                            pre_weight = gr.Textbox(label="Weight Column", value="weights")
                        with gr.Row():
                            pre_cfg_path = gr.Textbox(
                                label="Plot Config (for Oneway Factors)",
                                value="config_plot.json",
                                scale=4,
                            )
                            pre_load_factors_btn = gr.Button(
                                "Load Factors", variant="secondary", scale=1)
                        pre_oneway_factors = gr.Dropdown(
                            label="Oneway Factors",
                            choices=pre_factor_choices_init,
                            value=pre_factor_value_init,
                            multiselect=True,
                        )
                        pre_factor_status = gr.Textbox(
                            label="Oneway Factor Source",
                            value=pre_factor_status_init,
                            interactive=False,
                        )
                        with gr.Accordion("Advanced: Split Data Override", open=False):
                            with gr.Row():
                                pre_train_data_path = gr.Textbox(
                                    label="Train Data Path (optional)", value="")
                                pre_test_data_path = gr.Textbox(
                                    label="Test Data Path (optional)", value="")
                    with gr.Column(scale=2):
                        pre_feature_list = gr.Textbox(
                            label="Fallback Feature List (comma-separated)",
                            lines=3,
                            placeholder="feature_1, feature_2, feature_3",
                        )
                        pre_categorical = gr.Textbox(
                            label="Categorical Features (comma-separated)",
                            lines=2,
                            placeholder="feature_2, feature_3",
                        )
                        with gr.Row():
                            pre_n_bins = gr.Number(label="Bins", value=10, precision=0)
                            pre_holdout = gr.Slider(
                                label="Holdout Ratio", minimum=0.0, maximum=0.5,
                                value=0.25, step=0.05)
                            pre_seed = gr.Number(label="Random Seed", value=13, precision=0)
                with gr.Row():
                    pre_run_btn = gr.Button("Run Pre Oneway", variant="primary", scale=1)
                    pre_status = gr.Textbox(label="Status", interactive=False, scale=3)
                pre_log = gr.Textbox(label="Logs", lines=15, max_lines=40, interactive=False)
                pre_gallery = gr.Gallery(
                    label="Generated Images", columns=4, object_fit="contain", height=320)

            with gr.Tab("Direct Plot"):
                with gr.Row():
                    direct_cfg_path = gr.Textbox(
                        label="Plot Config", value="config_plot.json")
                    direct_xgb_cfg = gr.Textbox(
                        label="XGB Config", value="config_xgb_direct.json")
                    direct_resn_cfg = gr.Textbox(
                        label="ResN Config", value="config_resn_direct.json")
                with gr.Row():
                    direct_oneway_factors = gr.Dropdown(
                        label="Oneway Factors",
                        choices=direct_factor_choices_init,
                        value=direct_factor_value_init,
                        multiselect=True,
                        scale=4,
                    )
                    direct_load_factors_btn = gr.Button(
                        "Load Factors", variant="secondary", scale=1)
                direct_factor_status = gr.Textbox(
                    label="Oneway Factor Source",
                    value=direct_factor_status_init,
                    interactive=False,
                )
                with gr.Accordion("Advanced: Manual Data/Model Overrides", open=False):
                    with gr.Row():
                        direct_train_data_path = gr.Textbox(
                            label="Train Data Path (optional)", value="")
                        direct_test_data_path = gr.Textbox(
                            label="Test Data Path (optional)", value="")
                    with gr.Row():
                        direct_xgb_model_file = gr.File(
                            label="Upload XGB Model File",
                            file_types=[".pkl", ".pth"],
                            type="filepath",
                        )
                        direct_resn_model_file = gr.File(
                            label="Upload ResN Model File",
                            file_types=[".pkl", ".pth"],
                            type="filepath",
                        )
                with gr.Row():
                    direct_run_btn = gr.Button("Run Direct Plot", variant="primary", scale=1)
                    direct_status = gr.Textbox(label="Status", interactive=False, scale=3)
                direct_log = gr.Textbox(
                    label="Logs", lines=15, max_lines=40, interactive=False)
                direct_gallery = gr.Gallery(
                    label="Generated Images", columns=4, object_fit="contain", height=320)

            with gr.Tab("Embed Plot"):
                with gr.Row():
                    embed_cfg_path = gr.Textbox(
                        label="Plot Config", value="config_plot.json")
                    embed_ft_cfg = gr.Textbox(
                        label="FT Embed Config", value="config_ft_unsupervised_ddp_embed.json")
                    embed_runtime = gr.Checkbox(label="Use Runtime FT Embedding", value=False)
                with gr.Row():
                    embed_xgb_cfg = gr.Textbox(
                        label="XGB Embed Config", value="config_xgb_from_ft_unsupervised.json")
                    embed_resn_cfg = gr.Textbox(
                        label="ResN Embed Config", value="config_resn_from_ft_unsupervised.json")
                with gr.Row():
                    embed_oneway_factors = gr.Dropdown(
                        label="Oneway Factors",
                        choices=embed_factor_choices_init,
                        value=embed_factor_value_init,
                        multiselect=True,
                        scale=4,
                    )
                    embed_load_factors_btn = gr.Button(
                        "Load Factors", variant="secondary", scale=1)
                embed_factor_status = gr.Textbox(
                    label="Oneway Factor Source",
                    value=embed_factor_status_init,
                    interactive=False,
                )
                with gr.Accordion("Advanced: Manual Data/Model Overrides", open=False):
                    with gr.Row():
                        embed_train_data_path = gr.Textbox(
                            label="Train Data Path (optional)", value="")
                        embed_test_data_path = gr.Textbox(
                            label="Test Data Path (optional)", value="")
                    with gr.Row():
                        embed_xgb_model_file = gr.File(
                            label="Upload XGB Model File",
                            file_types=[".pkl", ".pth"],
                            type="filepath",
                        )
                        embed_resn_model_file = gr.File(
                            label="Upload ResN Model File",
                            file_types=[".pkl", ".pth"],
                            type="filepath",
                        )
                    with gr.Row():
                        embed_ft_model_file = gr.File(
                            label="Upload FT Model File",
                            file_types=[".pth"],
                            type="filepath",
                        )
                with gr.Row():
                    embed_run_btn = gr.Button("Run Embed Plot", variant="primary", scale=1)
                    embed_status = gr.Textbox(label="Status", interactive=False, scale=3)
                embed_log = gr.Textbox(
                    label="Logs", lines=15, max_lines=40, interactive=False)
                embed_gallery = gr.Gallery(
                    label="Generated Images", columns=4, object_fit="contain", height=320)

            with gr.Tab("Double Lift"):
                with gr.Row():
                    dl_data_path = gr.Textbox(
                        label="Data Path (CSV)", value="./Data/od_bc.csv", scale=3)
                    dl_output_path = gr.Textbox(
                        label="Output Image Path (optional)", value="", scale=2)
                with gr.Row():
                    dl_pred_col_1 = gr.Textbox(label="Prediction Column 1", value="pred_xgb")
                    dl_pred_col_2 = gr.Textbox(label="Prediction Column 2", value="pred_resn")
                    dl_target_col = gr.Textbox(label="Target Column", value="reponse")
                    dl_weight_col = gr.Textbox(label="Weight Column", value="weights")
                with gr.Row():
                    dl_label1 = gr.Textbox(label="Label 1", value="Model 1")
                    dl_label2 = gr.Textbox(label="Label 2", value="Model 2")
                    dl_n_bins = gr.Number(label="Bins", value=10, precision=0)
                    dl_rand_seed = gr.Number(label="Random Seed", value=13, precision=0)
                with gr.Row():
                    dl_holdout_ratio = gr.Slider(
                        label="Holdout Ratio (0 = all data)",
                        minimum=0.0, maximum=0.5, value=0.0, step=0.05, scale=2)
                    dl_split_strategy = gr.Dropdown(
                        label="Split Strategy",
                        choices=["random", "stratified", "time", "group"],
                        value="random", scale=1)
                    dl_split_group_col = gr.Textbox(
                        label="Group Column (optional)", value="", scale=1)
                    dl_split_time_col = gr.Textbox(
                        label="Time Column (optional)", value="", scale=1)
                    dl_split_time_ascending = gr.Checkbox(
                        label="Time Ascending", value=True, scale=1)
                with gr.Row():
                    dl_pred1_weighted = gr.Checkbox(label="Prediction 1 Is Weighted", value=False)
                    dl_pred2_weighted = gr.Checkbox(label="Prediction 2 Is Weighted", value=False)
                    dl_actual_weighted = gr.Checkbox(label="Actual Is Weighted", value=False)
                with gr.Row():
                    dl_run_btn = gr.Button("Run Double Lift", variant="primary", scale=1)
                    dl_status = gr.Textbox(label="Status", interactive=False, scale=3)
                dl_log = gr.Textbox(
                    label="Logs", lines=15, max_lines=40, interactive=False)
                dl_gallery = gr.Gallery(
                    label="Generated Images", columns=4, object_fit="contain", height=320)

            with gr.Tab("FT-Embed Compare"):
                gr.Markdown("Compare Direct vs FT-Embed models and draw double-lift curves.")
                with gr.Row():
                    cmp_model_key = gr.Dropdown(
                        label="Model Key",
                        choices=["xgb", "resn"],
                        value="xgb",
                        scale=1,
                    )
                    cmp_direct_cfg = gr.Textbox(
                        label="Direct Model Config", value="config_xgb_direct.json", scale=2)
                    cmp_ft_cfg = gr.Textbox(
                        label="FT Config", value="config_ft_unsupervised_ddp_embed.json", scale=2)
                    cmp_embed_cfg = gr.Textbox(
                        label="FT-Embed Model Config",
                        value="config_xgb_from_ft_unsupervised.json", scale=2)
                with gr.Row():
                    cmp_label_direct = gr.Textbox(label="Direct Label", value="XGB_raw")
                    cmp_label_ft = gr.Textbox(label="FT Label", value="XGB_ft_embed")
                    cmp_runtime = gr.Checkbox(label="Use Runtime FT Embedding", value=False)
                    cmp_bins = gr.Number(label="Bins Override", value=10, precision=0)
                with gr.Accordion("Advanced: Manual Data/Model Overrides", open=False):
                    with gr.Row():
                        cmp_train_data_path = gr.Textbox(
                            label="Train Data Path (optional)", value="")
                        cmp_test_data_path = gr.Textbox(
                            label="Test Data Path (optional)", value="")
                    with gr.Row():
                        cmp_direct_model_file = gr.File(
                            label="Upload Direct Model File",
                            file_types=[".pkl", ".pth"],
                            type="filepath",
                        )
                        cmp_ft_embed_model_file = gr.File(
                            label="Upload FT-Embed Model File",
                            file_types=[".pkl", ".pth"],
                            type="filepath",
                        )
                    with gr.Row():
                        cmp_ft_model_file = gr.File(
                            label="Upload FT Model File",
                            file_types=[".pth"],
                            type="filepath",
                        )
                with gr.Row():
                    cmp_run_btn = gr.Button("Run Compare", variant="primary", scale=1)
                    cmp_status = gr.Textbox(label="Status", interactive=False, scale=3)
                cmp_log = gr.Textbox(
                    label="Logs", lines=15, max_lines=40, interactive=False)
                cmp_gallery = gr.Gallery(
                    label="Generated Images", columns=4, object_fit="contain", height=320)

        # ════════════════════════════════════════════════════════════════
        #  Tab: Prediction
        # ════════════════════════════════════════════════════════════════
        with gr.Tab("Prediction"):
            gr.Markdown("### FT Embed Prediction")
            with gr.Row():
                pred_ft_cfg = gr.Textbox(
                    label="FT Config", value="config_ft_unsupervised_ddp_embed.json")
                pred_xgb_cfg = gr.Textbox(
                    label="XGB Config (optional)",
                    value="config_xgb_from_ft_unsupervised.json")
                pred_resn_cfg = gr.Textbox(
                    label="ResN Config (optional)",
                    value="config_resn_from_ft_unsupervised.json")
            with gr.Row():
                pred_model_name = gr.Textbox(label="Model Name (optional)", value="")
                pred_model_keys = gr.Textbox(label="Model Keys", value="xgb, resn")
                pred_input = gr.Textbox(label="Input Data", value="./Data/od_bc_new.csv")
                pred_output = gr.Textbox(
                    label="Output CSV", value="./Results/predictions_ft_xgb.csv")
            with gr.Accordion("Advanced: Upload Model Files", open=False):
                with gr.Row():
                    pred_ft_model_file = gr.File(
                        label="Upload FT Model File",
                        file_types=[".pth"],
                        type="filepath",
                    )
                    pred_xgb_model_file = gr.File(
                        label="Upload XGB Model File",
                        file_types=[".pkl", ".pth"],
                        type="filepath",
                    )
                    pred_resn_model_file = gr.File(
                        label="Upload ResN Model File",
                        file_types=[".pth", ".pkl"],
                        type="filepath",
                    )
            with gr.Row():
                pred_run_btn = gr.Button("Run Prediction", variant="primary", scale=1)
                pred_status = gr.Textbox(label="Status", interactive=False, scale=3)
            pred_log = gr.Textbox(label="Logs", lines=15, max_lines=40, interactive=False)

        # ════════════════════════════════════════════════════════════════
        #  Event Handlers
        # ════════════════════════════════════════════════════════════════

        # Working directory
        set_working_dir_btn.click(
            fn=_set_working_dir_ui,
            inputs=[working_dir_input],
            outputs=[
                working_dir_status,
                working_dir_input,
                working_dir_browse_root,
                working_dir_picker,
            ],
        )
        refresh_working_dir_btn.click(
            fn=_refresh_working_dir_choices_ui,
            inputs=[working_dir_browse_root],
            outputs=[working_dir_status, working_dir_picker],
        )
        use_selected_working_dir_btn.click(
            fn=_set_working_dir_ui,
            inputs=[working_dir_picker],
            outputs=[
                working_dir_status,
                working_dir_input,
                working_dir_browse_root,
                working_dir_picker,
            ],
        )
        working_dir_picker.change(
            fn=lambda path: str(path or ""),
            inputs=[working_dir_picker],
            outputs=[working_dir_input],
        )

        # Config load / build / save
        load_btn.click(
            fn=app.load_json_config,
            inputs=[json_file],
            outputs=[load_status, config_display, config_json],
            show_api=False,
        )

        build_btn.click(
            fn=app.build_config_from_ui,
            inputs=[
                # ── Core Data & Task ──
                data_dir, model_list, model_categories, target, weight,
                feature_list, categorical_features, task_type, distribution,
                binary_resp_nme,
                prop_test, holdout_ratio, val_ratio, split_strategy, rand_seed,
                epochs, output_dir, use_gpu, model_keys, max_evals,
                plot_curves,
                # ── Split & Pre-split ──
                split_group_col, split_time_col, split_time_ascending,
                train_data_path, test_data_path,
                split_cache_path, split_cache_force_rebuild,
                # ── Cross-Validation ──
                cv_strategy, cv_splits, cv_group_col, cv_time_col,
                cv_time_ascending,
                # ── XGBoost ──
                xgb_max_depth_max, xgb_n_estimators_max, xgb_gpu_id,
                xgb_use_dmatrix, xgb_chunk_size,
                xgb_cleanup_per_fold, xgb_cleanup_synchronize,
                xgb_search_space_json,
                # ── ResNet ──
                resn_use_lazy_dataset, resn_predict_batch_size,
                resn_weight_decay,
                resn_cleanup_per_fold, resn_cleanup_synchronize,
                resn_search_space_json,
                # ── FT-Transformer ──
                ft_role, ft_feature_prefix, ft_num_numeric_tokens,
                ft_use_lazy_dataset, ft_predict_batch_size,
                ft_oof_folds, ft_oof_strategy, ft_oof_shuffle,
                ft_cleanup_per_fold, ft_cleanup_synchronize,
                ft_search_space_json, ft_unsupervised_search_space_json,
                # ── GNN ──
                gnn_use_approx_knn, gnn_approx_knn_threshold,
                gnn_max_gpu_knn_nodes,
                gnn_knn_gpu_mem_ratio, gnn_knn_gpu_mem_overhead,
                gnn_max_fit_rows, gnn_max_predict_rows, gnn_predict_chunk_rows,
                gnn_graph_cache,
                gnn_cleanup_per_fold, gnn_cleanup_synchronize,
                # ── Distributed ──
                nproc_per_node, ddp_min_rows,
                use_resn_ddp, use_ft_ddp,
                use_resn_data_parallel, use_ft_data_parallel, use_gnn_data_parallel,
                # ── Preprocessing ──
                build_oht, oht_sparse_csr, keep_unscaled_oht,
                infer_categorical_max_unique, infer_categorical_max_ratio,
                # ── Geographic ──
                geo_feature_nmes, region_province_col, region_city_col,
                region_effect_alpha,
                geo_token_hidden_dim, geo_token_layers, geo_token_dropout,
                geo_token_k_neighbors, geo_token_learning_rate, geo_token_epochs,
                # ── Ensemble & Refit ──
                final_ensemble, final_ensemble_k, final_refit,
                reuse_best_params,
                # ── Output & Caching ──
                optuna_study_prefix, optuna_cleanup_synchronize,
                cache_predictions, prediction_cache_format,
                dataloader_workers, bo_sample_limit,
                # ── Plot Settings ──
                plot_enable, plot_n_bins, plot_oneway, plot_oneway_pred,
                plot_pre_oneway, plot_double_lift,
                # ── Calibration ──
                calibration_enable, calibration_method,
                calibration_max_rows, calibration_seed,
                # ── Threshold ──
                threshold_enable, threshold_metric,
                threshold_grid, threshold_max_rows, threshold_seed,
                # ── Bootstrap ──
                bootstrap_enable, bootstrap_n_samples,
                bootstrap_ci, bootstrap_seed,
                # ── Advanced ──
                config_overrides_json,
            ],
            outputs=[build_status, config_json],
        )

        save_config_btn.click(
            fn=app.save_config,
            inputs=[config_json, save_filename],
            outputs=[save_status],
        )

        # Run task
        run_btn.click(
            fn=app.run_training,
            inputs=[config_json],
            outputs=[run_status, log_output],
        )
        open_folder_btn.click(
            fn=app.open_results_folder,
            inputs=[config_json],
            outputs=[folder_status],
        )

        # Workflow config
        workflow_load_btn.click(
            fn=app.load_workflow_config,
            inputs=[workflow_file],
            outputs=[workflow_load_status, workflow_config_json],
            show_api=False,
        )
        workflow_run_btn.click(
            fn=app.run_workflow_config_ui,
            inputs=[workflow_config_json],
            outputs=[workflow_status, workflow_log],
        )

        # FT two-step
        prepare_step1_btn.click(
            fn=app.prepare_ft_step1,
            inputs=[config_json, ft_use_ddp, ft_nproc],
            outputs=[step1_status, step1_config_display],
        )
        prepare_step2_btn.click(
            fn=app.prepare_ft_step2,
            inputs=[
                gr.State(lambda: app.current_step1_config or "temp_ft_step1_config.json"),
                target_models_input,
                augmented_data_dir_input,
                xgb_overrides_input,
                resn_overrides_input,
            ],
            outputs=[step2_status, xgb_config_display, resn_config_display],
        )

        # Oneway factor loaders
        def _bind_oneway_factor_loader(load_btn, cfg_input, status_output, factors_output):
            load_btn.click(
                fn=_load_oneway_factors_ui,
                inputs=[cfg_input],
                outputs=[status_output, factors_output],
                show_api=False,
            )
            cfg_input.change(
                fn=_load_oneway_factors_ui,
                inputs=[cfg_input],
                outputs=[status_output, factors_output],
                show_api=False,
            )

        _bind_oneway_factor_loader(
            pre_load_factors_btn, pre_cfg_path, pre_factor_status, pre_oneway_factors)
        _bind_oneway_factor_loader(
            direct_load_factors_btn, direct_cfg_path, direct_factor_status, direct_oneway_factors)
        _bind_oneway_factor_loader(
            embed_load_factors_btn, embed_cfg_path, embed_factor_status, embed_oneway_factors)

        # Plotting workflows
        pre_run_btn.click(
            fn=app.run_pre_oneway_ui,
            inputs=[
                pre_data_path, pre_model_name, pre_target, pre_weight,
                pre_feature_list, pre_oneway_factors, pre_categorical, pre_n_bins,
                pre_holdout, pre_seed, pre_output_dir,
                pre_train_data_path, pre_test_data_path,
            ],
            outputs=[pre_status, pre_log, pre_gallery],
            show_api=False,
        )
        direct_run_btn.click(
            fn=app.run_plot_direct_ui,
            inputs=[
                direct_cfg_path, direct_xgb_cfg, direct_resn_cfg,
                direct_oneway_factors,
                direct_train_data_path, direct_test_data_path,
                direct_xgb_model_file, direct_resn_model_file,
            ],
            outputs=[direct_status, direct_log, direct_gallery],
            show_api=False,
        )
        embed_run_btn.click(
            fn=app.run_plot_embed_ui,
            inputs=[
                embed_cfg_path, embed_xgb_cfg, embed_resn_cfg,
                embed_ft_cfg, embed_runtime, embed_oneway_factors,
                embed_train_data_path, embed_test_data_path,
                embed_xgb_model_file, embed_resn_model_file, embed_ft_model_file,
            ],
            outputs=[embed_status, embed_log, embed_gallery],
            show_api=False,
        )
        dl_run_btn.click(
            fn=app.run_double_lift_ui,
            inputs=[
                dl_data_path, dl_pred_col_1, dl_pred_col_2,
                dl_target_col, dl_weight_col,
                dl_n_bins, dl_label1, dl_label2,
                dl_pred1_weighted, dl_pred2_weighted, dl_actual_weighted,
                dl_holdout_ratio, dl_split_strategy, dl_split_group_col,
                dl_split_time_col, dl_split_time_ascending, dl_rand_seed,
                dl_output_path,
            ],
            outputs=[dl_status, dl_log, dl_gallery],
            show_api=False,
        )

        # Compare
        cmp_model_key.change(
            fn=_suggest_compare_defaults,
            inputs=[cmp_model_key],
            outputs=[cmp_direct_cfg, cmp_embed_cfg, cmp_label_direct, cmp_label_ft],
        )
        cmp_run_btn.click(
            fn=app.run_compare_ui,
            inputs=[
                cmp_model_key, cmp_direct_cfg, cmp_ft_cfg, cmp_embed_cfg,
                cmp_label_direct, cmp_label_ft, cmp_runtime, cmp_bins,
                cmp_train_data_path, cmp_test_data_path,
                cmp_direct_model_file, cmp_ft_embed_model_file, cmp_ft_model_file,
            ],
            outputs=[cmp_status, cmp_log, cmp_gallery],
            show_api=False,
        )

        # Prediction
        pred_run_btn.click(
            fn=app.run_predict_ui,
            inputs=[
                pred_ft_cfg, pred_xgb_cfg, pred_resn_cfg, pred_input,
                pred_output, pred_model_name, pred_model_keys,
                pred_ft_model_file, pred_xgb_model_file, pred_resn_model_file,
            ],
            outputs=[pred_status, pred_log],
            show_api=False,
        )

    return demo


if __name__ == "__main__":
    demo = create_ui()
    server_name = os.environ.get("GRADIO_SERVER_NAME", "127.0.0.1").strip() or "127.0.0.1"
    launch_kwargs = {
        "server_name": server_name,
        "share": False,
        "show_error": True,
    }
    server_port_env = os.environ.get("GRADIO_SERVER_PORT", "").strip()
    if server_port_env:
        try:
            launch_kwargs["server_port"] = int(server_port_env)
        except ValueError as exc:
            raise ValueError(
                f"Invalid GRADIO_SERVER_PORT: {server_port_env!r}. Must be an integer."
            ) from exc
    if "analytics_enabled" in inspect.signature(demo.launch).parameters:
        launch_kwargs["analytics_enabled"] = False
    demo.launch(**launch_kwargs)
