demo_callbacks.py

# Copyright 2026 D-Wave
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software distributed under the License
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
# or implied. See the License for the specific language governing permissions and limitations under
# the License.
#
# The use of code in the quantum-blockchain repository with a quantum computing system is protected
# by the intellectual property rights of D-Wave Quantum Inc. and its affiliates.
#
# The use of the quantum blockchain implementations below (including the Miner, Block, and Hash
# methods) with D-Wave's quantum computing system will require access to D-Wave’s LeapTM quantum
# cloud service and will be governed by the Leap Cloud Subscription Agreement available at:
# https://cloud.dwavesys.com/leap/legal/cloud_subscription_agreement/

from __future__ import annotations

import os
import time
from typing import NamedTuple

import dash
import plotly.graph_objects as go
from dash import ALL, MATCH, set_props
from dash.dependencies import Input, Output, State
from dash.exceptions import PreventUpdate
import numpy as np

from demo_configs import (
    MIN_SIMULATION_LOOP_TIME,
    QUANTUM_HASH_LENGTH,
    ALLOWABLE_ERR,
    N_ZEROES,
    RANDOM_SEED,
)
from demo_interface import BUTTONS, QPU_SOLVER_OPTS, SIMULATED_SOLVER_OPTS
from src.agents.trial_manager import TrialManager
from src.demo_enums import InterfaceButton, SolverMode, ViewOpt
from src.protocols.simulation_identification import (
    generate_simulation_id,
    get_simulation_params_from_id,
    simulation_parameter_fields,
)
from src.utilities.display_update import change_button_visibility, render_miner_status
from src.utilities.save_simulation_data import get_save_data_filename, save_simulation_data
from src.utilities.spiral_plotter import SpiralPlotter
from src.values import MAX_RNG_SEED_LEN, OUTPUTS_PATH

graph_layout_dict = dict(
    autosize=False,
    showlegend=False,
    xaxis=dict(showticklabels=False),
    yaxis=dict(showticklabels=False),
    margin=dict(l=0, r=0, b=0, t=0, pad=4),
    paper_bgcolor="white",
    plot_bgcolor="white",
)

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# =====================================================================================================
#                             SECTION: Mining Round Steps                                             |
# =====================================================================================================
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~


@dash.callback(
    Output({"type": "button", "index": ALL}, "style", allow_duplicate=True),
    Output(BUTTONS["SAVE"].id, "children", allow_duplicate=True),
    Output(BUTTONS["SAVE"].id, "disabled", allow_duplicate=True),
    inputs=[
        Input("simulation-save-filename", "data"),
        State("blockchain-structure-data", "data"),
    ],
    progress=[
        Output("current-block-data", "data"),
    ],
    cancel=[Input("simulation-pause-target", "data")],
    prevent_initial_call=True,
    background=True,
)
def simulation(
    set_progress_miner_table,  # set_progress function for 'progress' argument
    save_filename: str,
    stored_blockchain_data: list,
) -> tuple[list[dict], str, bool]:
    """Manages a single run of the blockchain simulation.

    This callback is triggered (indirectly) by the "run" and "resume" buttons. When triggered
    by the "run" button, it will complete a full run of the blockchain simulation with the solver
    scheme and the numbers of blocks and miners each defined by their respective input fields. The
    execution of this function is internally divided up into 'rounds,' where one round covers a
    single action (mining or validation) from each miner in the simulation. Each action takes
    roughly 1 second, so a round will last for roughly as many seconds as there are miners. It
    will run a number of rounds equal to the 'num_blocks' parameter. As this callback runs, it
    calls the 'set_progress_miner_table' function to provide regular updates to the miner table
    (approximately 1 per second). It will also use the 'set_props' function to keep the
    'blockchain-structure-data' dcc.Store up-to-date with the same frequency, and to update each of
    the graph views once per round.

    Args:
        set_progress_miner_table: the progress function for passing data out to update the miner
            status table
        save_filename: Setting this triggers the callback to run. The filename also encodes the
            parameters of the simulation, which will be parsed to determine solver, miner, blocks
            and random seed settings for the simulation.
        stored_blockchain_data: The data structure storing the current blockchain data.
            This will be empty when starting a new simulation, but if the simulation has been
            paused, it will contain all the data about the blocks mined up to this point, allowing
            the simulation to be restarted from the same state.

    Returns:
        Most of the output of this function is passed by 'set_progress_miner_table' or with the
        'set_props' function, so the only actual return values simply toggle button visibility

        button_visibility: makes the 'save', and 'reset' buttons visible and hides the 'pause' button.
        save_button_text: changes text to "Save Data"
        save_button_disabled: enables the 'save' button"""

    if not save_filename:
        raise PreventUpdate

    simulation_id = save_filename[: sum(simulation_parameter_fields.values())]
    manager = TrialManager(**get_simulation_params_from_id(simulation_id))
    if len(stored_blockchain_data) > 0:
        manager.restart_simulation(stored_blockchain_data)
        print("Simulation Restarts")
    else:
        print("Simulation Begins")

    while manager.blocks_mined < manager.max_blocks or manager.round_progress > 0:
        iter_start_time = time.time()
        miner_id = manager.single_step()

        set_props("blockchain-structure-data", {"data": manager.chain_data})
        time.sleep(0.2)

        set_progress_miner_table(manager.chain_data[-1])
        time.sleep(0.2)

        plotter = SpiralPlotter()

        if manager.round_progress == 0:  # round_progress resets to 0 at the end of a round
            mining_hashes = manager.mining_hashes
            global_fig = plotter.create_plot_from_tree(manager.global_tree, mining_hashes, True)
            global_fig.update_layout(**graph_layout_dict)
            set_props({"type": "view-graph", "index": 0}, {"figure": global_fig})
        elif miner_id in ViewOpt._member_names_:
            view_miner = manager.miners[miner_id]
            mining_hash = [manager.mining_hashes[0]]
            miner_fig = plotter.create_plot_from_tree(view_miner.blockchain, mining_hash)
            miner_fig.update_layout(**graph_layout_dict)
            view_idx = ViewOpt[miner_id].value - 1  # ViewOpt vals are off by 1 from miner names
            set_props({"type": "view-graph", "index": view_idx}, {"figure": miner_fig})

        iter_end_time = time.time()
        iter_total_time = iter_end_time - iter_start_time
        if iter_total_time < MIN_SIMULATION_LOOP_TIME:
            time.sleep(MIN_SIMULATION_LOOP_TIME - iter_total_time)

    set_props("blockchain-structure-data", {"data": manager.chain_data})
    time.sleep(0.3)  # ensure final data update is processed before Save Button is enabled

    print("Simulation Ends")

    button_visibility = change_button_visibility(
        buttons_to_show=[InterfaceButton.RESET, InterfaceButton.SAVE],
        buttons_to_hide=[InterfaceButton.PAUSE],
    )
    return button_visibility, "Save Data", False


# ======================================================================================================


@dash.callback(
    Output("miner-graph-and-table", "className", allow_duplicate=True),
    Output("view-select-and-block-status", "className", allow_duplicate=True),
    Output("prelim-text", "className"),
    Output("block-status", "children", allow_duplicate=True),
    Output("miner-status-table", "children", allow_duplicate=True),
    Output("graph-loading", "display", allow_duplicate=True),
    inputs=[
        Input("current-block-data", "data"),
        State("miner-slider", "value"),
        State("solver-mode-select", "value"),
        State("qpu-solver-select", "value"),
        State("simulated-solver-select", "value"),
        State("blocks-input", "value"),
    ],
    prevent_initial_call=True,
)
def update_miner_display(
    current_block_data: dict,
    num_miners: int,
    solver_mode: str,
    qpu_select: str,
    simulated_select: str,
    num_blocks: int,
):
    """Processes blockchain structure data and uses it to update the miner status
    table and the graph display."""

    solver_mode = SolverMode(int(solver_mode))
    show_solvers = (solver_mode is SolverMode.QPU and int(qpu_select) == -1) or (
        solver_mode is SolverMode.SIMULATED and int(simulated_select) == -1
    )

    block_number = current_block_data["block_number"]
    block_progress = len(current_block_data["scores"])

    block_status_text = (
        f"Currently mining block {block_number}"
        if block_number < num_blocks or block_progress < num_miners
        else "Finished Simulation"
    )

    miner_table_body = render_miner_status(current_block_data, num_miners, show_solvers)

    graph_loading = "auto" if block_number > 1 else dash.no_update

    return "", "", "display-none", block_status_text, miner_table_body, graph_loading


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# =====================================================================================================
#                             SECTION: User Controls                                                |
# =====================================================================================================
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~


class StartSimulationReturn(NamedTuple):
    """Return type for the ``run_simulation`` callback function."""

    save_filename: str
    button_styles: list = change_button_visibility(
        buttons_to_show=[InterfaceButton.PAUSE], buttons_to_hide=[InterfaceButton.START]
    )
    miner_slider_disabled: bool = True
    blocks_input_disabled: bool = True
    qpu_solver_select_disabled: bool = True
    simulated_solver_select_disabled: bool = True
    solver_mode_select_disabled: bool = True
    graph_loading_display: str = "show"


@dash.callback(
    Output("simulation-save-filename", "data", allow_duplicate=True),
    Output({"type": "button", "index": ALL}, "style", allow_duplicate=True),
    Output("miner-slider", "disabled", allow_duplicate=True),
    Output("blocks-input", "disabled", allow_duplicate=True),
    Output("qpu-solver-select", "disabled", allow_duplicate=True),
    Output("simulated-solver-select", "disabled", allow_duplicate=True),
    Output("solver-mode-select", "disabled", allow_duplicate=True),
    Output("graph-loading", "display", allow_duplicate=True),
    inputs=[
        Input(BUTTONS["START"].id, "n_clicks"),
        State("miner-slider", "value"),
        State("blocks-input", "value"),
        State("solver-mode-select", "value"),
        State("qpu-solver-select", "value"),
        State("simulated-solver-select", "value"),
        State("simulation-save-filename", "data"),
    ],
    prevent_initial_call=True,
)
def start_simulation(
    start_click: int,
    miner_slider_val: int,
    blocks_input_val: int,
    solver_mode_val: int,
    qpu_solver_select_val: int,
    simulated_solver_select_val: int,
    save_filename: str,
) -> StartSimulationReturn:
    """Runs a simulation with the selected number of miners and blocks.

    Args:
        start_click: unused
        miner_slider_val: value of the miner slider. Sets the number of miners in the simulation.
        blocks_input_val: value of the blocks input. Sets how many blocks the simulation will mine.
        solver_mode_val: value of the solver mode select. Determines whether the simulation will
            use the QPU solvers or the simulated solvers.
        qpu_solver_select_val: value of the QPU solver select dropdown. Determines which QPU solver
            will be used, if any.
        simulated_solver_select_val: value of the simulated solver select dropdown. Determines
            which simulated solver will be used, if any.
        save_filename: The name of the file to save the blockchain data. Only relevant if the
            REPLICATION_ID line in demo_configs is set to a specific simulation ID, in which
            case that simulation ID will be extracted from the save_filename and used to
            determine the parameters of the simulation.
    """

    if save_filename:
        print(
            f"Replicating Simulation with ID {save_filename[:sum(simulation_parameter_fields.values())]}."
        )
        return StartSimulationReturn(save_filename)

    init_rng = np.random.default_rng(RANDOM_SEED)
    manager_prng_seed = int(init_rng.integers(0, 16**MAX_RNG_SEED_LEN - 1))

    solver_mode = SolverMode(int(solver_mode_val))
    if solver_mode is SolverMode.QPU:
        dropdown_idx = int(qpu_solver_select_val)
        solver_opts = QPU_SOLVER_OPTS
    else:
        dropdown_idx = int(simulated_solver_select_val)
        solver_opts = SIMULATED_SOLVER_OPTS

    solver_lookup = {idx: solver_name for solver_name, idx in solver_opts.items()}
    solver_selection = (
        list(solver_opts.keys()) if dropdown_idx == -1 else [solver_lookup[dropdown_idx]]
    )
    params_dict = dict(
        max_blocks=blocks_input_val,
        num_miners=miner_slider_val,
        solvers=solver_selection,
        quantum_hash_length=QUANTUM_HASH_LENGTH,
        n_zeroes=N_ZEROES,
        allowable_err=ALLOWABLE_ERR,
        prng_seed=manager_prng_seed,
    )

    simulation_id = generate_simulation_id(**params_dict)
    save_filename = get_save_data_filename(simulation_id)

    print(
        f"Starting simulation with ID {simulation_id} with {miner_slider_val} miners and {blocks_input_val} blocks."
    )

    return StartSimulationReturn(save_filename)


# ========================================================================================


@dash.callback(
    Output({"type": "button", "index": ALL}, "style", allow_duplicate=True),
    Output(BUTTONS["SAVE"].id, "children", allow_duplicate=True),
    Output(BUTTONS["SAVE"].id, "disabled", allow_duplicate=True),
    Output("simulation-pause-target", "data"),
    inputs=[
        Input(BUTTONS["PAUSE"].id, "n_clicks"),
    ],
    prevent_initial_call=True,
)
def pause_simulation(pause_click: int):
    """This callback will pause the current, in-progress simulation. In reality, the
    'simulation' callback is cancelled, but the data defining its current state is
    still stored in the blockchain_structure_data dcc.Store object, so the simulation
    can be restarted by reconstructing the state.

    Args:
        pause_click: Unused. The pause button just needs to trigger the callback,
            its value is irrelevant.

    Returns:
        button_visibility_change: makes the 'reset', 'resume', and 'save' buttons visible and
            hides the 'pause' button.
        save_button_text: changes text to "Save Data"
        save_button_disabled: enables the 'save' button
        simulation-pause-target: Altering this Store (even from True to True) triggers the
            cancellation of the 'simulation' callback, effectively pausing the simulation."""

    visible_buttons = change_button_visibility(
        buttons_to_show=[InterfaceButton.RESET, InterfaceButton.RESUME, InterfaceButton.SAVE],
        buttons_to_hide=[InterfaceButton.PAUSE],
    )

    print("Simulation Paused")

    return visible_buttons, "Save Data", False, True


# ========================================================================================


@dash.callback(
    Output({"type": "button", "index": ALL}, "style", allow_duplicate=True),
    Output("simulation-save-filename", "data", allow_duplicate=True),
    inputs=[
        Input(BUTTONS["RESUME"].id, "n_clicks"),
        State("simulation-save-filename", "data"),
    ],
    prevent_initial_call=True,
)
def resume_simulation(pause_click: int, save_file: str):
    """Resumes a paused simulation. In practice, this means starting a new instance of the
    'simulation' callback, but without resetting the blockchain data. The simulation
    will then reconstruct its previous state and pick up where it left off.

    Args:
        pause_click: Unused.

    Returns:
        button_visibility_change: makes the 'pause' button visible and hides the 'reset', 'resume',
            and 'save' buttons.
        simulation-save-filename: The name of the file to save the blockchain data. This callback
            will merely return the same value already present in this store, but doing so will
            trigger the 'simulation' callback to fire.
    """

    visible_buttons = change_button_visibility(
        buttons_to_show=[InterfaceButton.PAUSE],
        buttons_to_hide=[InterfaceButton.RESET, InterfaceButton.RESUME, InterfaceButton.SAVE],
    )

    print(f"Resuming simulation with ID {save_file[:sum(simulation_parameter_fields.values())]}")

    return visible_buttons, save_file


# =========================================================================================


@dash.callback(
    Output(BUTTONS["SAVE"].id, "children", allow_duplicate=True),
    Output(BUTTONS["SAVE"].id, "disabled", allow_duplicate=True),
    inputs=[
        Input(BUTTONS["SAVE"].id, "n_clicks"),
        State("blockchain-structure-data", "data"),
        State("simulation-save-filename", "data"),
    ],
    prevent_initial_call=True,
)
def save_data(n_clicks: int, blockchain_data: list, save_filename: str) -> tuple[str, bool]:
    """Saves the current simulation data to a file when the 'save' button is clicked. Disables
        the 'save' button (changing its text to "Data Saved") to avoid redundant writes.

    Args:
        n_clicks: number of clicks on the 'save' button. Trigger for the callback, value is
            otherwise irrelevant.
        blockchain_data: The data structure storing the current blockchain data.
        save_filename: The name of the file to save the blockchain data to, including the .csv extension.

    Returns:
        save_button_text: changes text to "Data Saved"
        save_button_disabled: disables the 'save' button to prevent redundant saves.
    """

    save_simulation_data(blockchain_data, save_filename)
    print(f"Simulation data saved to {os.path.join(OUTPUTS_PATH, save_filename)}.")
    # Change button text to indicate data has been saved and disable it to prevent multiple clicks
    return "Data Saved", True


# ========================================================================================
class ResetSimulationReturn(NamedTuple):
    """Return type for the ``reset_simulation`` callback function."""

    intro_text_classname: str = ""
    loading_text_classname: str = "display-none"
    miner_graph_and_table_classname: str = "display-none"
    view_select_and_block_status_classname: str = "visibility-hidden"
    button_styles: list[dict] = change_button_visibility(
        buttons_to_show=[InterfaceButton.START],
        buttons_to_hide=[InterfaceButton.RESET, InterfaceButton.RESUME, InterfaceButton.SAVE],
    )
    prelim_text_classname: str = ""
    miner_slider_disabled: bool = False
    blocks_input_disabled: bool = False
    qpu_solver_select_disabled: bool = False
    simulated_solver_select_disabled: bool = False
    solver_mode_select_disabled: bool = False
    blockchain_data: list = []
    simulation_save_filename: str = ""
    graphs: list[go.Figure] = dash.no_update


@dash.callback(
    Output("intro-text", "className", allow_duplicate=True),
    Output("loading-text", "className", allow_duplicate=True),
    Output("miner-graph-and-table", "className", allow_duplicate=True),
    Output("view-select-and-block-status", "className", allow_duplicate=True),
    Output({"type": "button", "index": ALL}, "style", allow_duplicate=True),
    Output("prelim-text", "className", allow_duplicate=True),
    Output("miner-slider", "disabled", allow_duplicate=True),
    Output("blocks-input", "disabled", allow_duplicate=True),
    Output("qpu-solver-select", "disabled", allow_duplicate=True),
    Output("simulated-solver-select", "disabled", allow_duplicate=True),
    Output("solver-mode-select", "disabled", allow_duplicate=True),
    Output("blockchain-structure-data", "data", allow_duplicate=True),
    Output("simulation-save-filename", "data", allow_duplicate=True),
    Output({"type": "view-graph", "index": ALL}, "figure"),
    inputs=[
        Input(BUTTONS["RESET"].id, "n_clicks"),
        State({"type": "view-graph", "index": ALL}, "figure"),
    ],
    prevent_initial_call=True,
)
def reset_simulation(reset_click: int, graphs: list) -> ResetSimulationReturn:
    """Resets the simulation, allowing a new simulation to be started."""

    return ResetSimulationReturn(graphs=[go.Figure()] * len(graphs))


# =======================================================================================


@dash.callback(
    Output({"type": "view-wrapper", "index": ALL}, "className"),
    inputs=[
        Input("view-select", "value"),
        State({"type": "view-wrapper", "index": ALL}, "className"),
    ],
    prevent_initial_call=True,
)
def toggle_graph_display(selected_view: str, graphs: list[str]):
    """Toggles the visibility of the four different graph displays. Will default to showing the Global
    View graph. When triggered, will hide three of the graph displays and make the selected one visible.
    """

    return_tuple = ["display-none"] * len(graphs)
    return_tuple[int(selected_view) - 1] = ""

    return return_tuple


# =========================================================================================


@dash.callback(
    Output("qpu-dropdown", "className"),
    Output("simulated-dropdown", "className"),
    inputs=[
        Input("solver-mode-select", "value"),
    ],
)
def toggle_solver_mode(solver_mode):
    """Toggles between QPU Solver mode and Simulated Solver Mode"""
    solver_mode = SolverMode(int(solver_mode))

    if solver_mode is SolverMode.QPU:
        return "", "display-none"
    elif solver_mode is SolverMode.SIMULATED:
        return "display-none", ""

    raise Exception("Invalid solver select option")


@dash.callback(
    Output({"type": "to-collapse-class", "index": MATCH}, "className"),
    Output({"type": "collapse-trigger", "index": MATCH}, "aria-expanded"),
    inputs=[
        Input({"type": "collapse-trigger", "index": MATCH}, "n_clicks"),
        State({"type": "to-collapse-class", "index": MATCH}, "className"),
    ],
    prevent_initial_call=True,
)
def toggle_left_column(collapse_trigger: int, to_collapse_class: str) -> tuple[str, str]:
    """Toggles a 'collapsed' class that hides and shows some aspect of the UI.

    Args:
        collapse_trigger: The (total) number of times a collapse button has been clicked.
        to_collapse_class: Current class name of the thing to collapse, 'collapsed' if not
            visible, empty string if visible.

    Returns:
        str: The new class name of the thing to collapse.
        str: The aria-expanded value.
    """

    classes = to_collapse_class.split(" ") if to_collapse_class else []
    if "collapsed" in classes:
        classes.remove("collapsed")
        return " ".join(classes), "true"
    return to_collapse_class + " collapsed" if to_collapse_class else "collapsed", "false"