Local Simulated and Quantum Annealing

# QUBONOTEBOOKS_COLAB_IMPORT_CELL
Base.invokelatest(
    QUBONotebooksBootstrap.warm_notebook_packages!,
    "11-Annealing",
);

@assert pkgversion(DWave) == v"0.7.6"
@assert DWave.OCEAN_SDK_VERSION == v"9.3.0"
println(
    "Local annealing runtime ready: DWave $(pkgversion(DWave)), " *
    "Ocean $(DWave.OCEAN_SDK_VERSION).",
)

const LOCAL_SEED = 96096
const LOCAL_READS = 256
const LOCAL_SWEEPS = 1_000

local_config = (
    label = "DWave.Neal",
    optimizer = DWave.Neal.Optimizer,
    attributes = Dict{String,Any}(
        "num_reads" => LOCAL_READS,
        "num_sweeps" => LOCAL_SWEEPS,
        "seed" => LOCAL_SEED,
        "beta_schedule_type" => "geometric",
    ),
)

all_binary_states(n::Integer) = [
    collect(state)
    for state in Iterators.product(ntuple(_ -> (0, 1), n)...)
]

"""
    exact_baseline(n, energy)

Enumerate a tutorial-sized binary model using its independent application
energy function.
"""
function exact_baseline(n::Integer, energy)
    states = all_binary_states(n)
    energies = energy.(states)
    best_energy = minimum(energies)
    optima = [
        states[i]
        for i in eachindex(states)
        if isapprox(energies[i], best_energy; atol = 1e-9, rtol = 0)
    ]
    return (
        energy = Float64(best_energy),
        optima = optima,
        degeneracy = length(optima),
    )
end

"""
    quadratic_density(n, energy)

Recover the nonzero pair-interaction density from an independent quadratic
energy evaluator. The finite-difference coefficient is zero exactly when a
pair does not interact.
"""
function quadratic_density(n::Integer, energy)
    n < 2 && return 0.0
    zero_bits = zeros(Int, n)
    zero_energy = energy(zero_bits)
    singleton_energy = [
        begin
            bits = zeros(Int, n)
            bits[i] = 1
            energy(bits)
        end
        for i in 1:n
    ]
    interactions = 0
    for i in 1:(n - 1), j in (i + 1):n
        bits = zeros(Int, n)
        bits[i] = 1
        bits[j] = 1
        coefficient = energy(bits) - singleton_energy[i] -
            singleton_energy[j] + zero_energy
        interactions += !isapprox(coefficient, 0; atol = 1e-12, rtol = 0)
    end
    return interactions / binomial(n, 2)
end

partition_weights = [1, 3, 4, 8]
weighted_edges = [
    (1, 2, 2),
    (1, 3, 1),
    (2, 3, 2),
    (2, 4, 1),
    (3, 4, 3),
]
cover_edges = [(1, 2), (1, 3), (2, 3), (3, 4)]
cover_penalty = 2
order_values = [2, 3, 4, 5, 6, 8]
risk_exposures = [
    1   1  2  5   4  5
    4  -2  3  6  -3  0
]
value_weight = 2
risk_weight = 1

"""Build the number-partitioning QUBO from Notebook 7."""
function build_partition_model()
    model = Model()
    @variable(model, x[1:length(partition_weights)], Bin)
    @objective(
        model,
        Min,
        (
            sum(
                partition_weights[i] * (1 - 2 * x[i])
                for i in eachindex(partition_weights)
            )
        )^2,
    )
    return model, x
end

"""Build the weighted Max-Cut QUBO from Notebook 7."""
function build_maxcut_model()
    model = Model()
    @variable(model, x[1:4], Bin)
    @objective(
        model,
        Min,
        -sum(
            weight * (x[i] + x[j] - 2 * x[i] * x[j])
            for (i, j, weight) in weighted_edges
        ),
    )
    return model, x
end

"""Build the penalized minimum-vertex-cover QUBO from Notebook 7."""
function build_cover_model()
    model = Model()
    @variable(model, x[1:4], Bin)
    @objective(
        model,
        Min,
        sum(x) + cover_penalty * sum(
            (1 - x[i]) * (1 - x[j]) for (i, j) in cover_edges
        ),
    )
    return model, x
end

"""Build the grouped-square order-partitioning QUBO from Notebook 8."""
function build_order_model()
    model = Model()
    @variable(model, x[1:length(order_values)], Bin)
    @objective(
        model,
        Min,
        value_weight * (
            sum(order_values) -
            2 * sum(order_values[j] * x[j] for j in eachindex(order_values))
        )^2 +
        risk_weight * sum(
            (
                sum(
                    risk_exposures[i, j] * (2 * x[j] - 1)
                    for j in eachindex(order_values)
                )
            )^2
            for i in axes(risk_exposures, 1)
        ),
    )
    return model, x
end

partition_energy(bits) = sum(
    partition_weights[i] * (1 - 2 * bits[i])
    for i in eachindex(partition_weights)
)^2

maxcut_weight(bits) = sum(
    bits[i] == bits[j] ? 0 : weight
    for (i, j, weight) in weighted_edges
)
maxcut_energy(bits) = -maxcut_weight(bits)

uncovered_cover_edges(bits) = [
    (i, j) for (i, j) in cover_edges if bits[i] == 0 && bits[j] == 0
]
cover_energy(bits) = sum(bits) +
    cover_penalty * length(uncovered_cover_edges(bits))

function order_imbalances(bits)
    value = sum(
        order_values[j] * (1 - 2 * bits[j])
        for j in eachindex(order_values)
    )
    risk = [
        sum(
            risk_exposures[i, j] * (2 * bits[j] - 1)
            for j in eachindex(order_values)
        )
        for i in axes(risk_exposures, 1)
    ]
    return value, risk
end

function order_energy(bits)
    value, risk = order_imbalances(bits)
    return value_weight * value^2 + risk_weight * sum(abs2, risk)
end

decode_partition(bits) = (
    application_score = "imbalance=$(abs(sum(partition_weights .* (1 .- 2 .* bits))))",
    feasible = true,
    details = "two complementary group labels",
)
decode_maxcut(bits) = (
    application_score = "cut weight=$(maxcut_weight(bits))",
    feasible = true,
    details = "sides=$(findall(==(0), bits))/$(findall(==(1), bits))",
)
decode_cover(bits) = (
    application_score = "cover size=$(sum(bits))",
    feasible = isempty(uncovered_cover_edges(bits)),
    details = "selected=$(findall(==(1), bits))",
)
function decode_order(bits)
    value, risk = order_imbalances(bits)
    return (
        application_score = "value Δ=$(abs(value)); risk²=$(sum(abs2, risk))",
        feasible = true,
        details = "risk Δ=$(risk)",
    )
end

split_csv_lines(path) = [split(line, ',') for line in readlines(path)]

"""
    load_tcga_aml_aggregate(repo_root)

Load and validate the committed aggregate used by Notebook 9 without network
access or patient-level identifiers.
"""
function load_tcga_aml_aggregate(repo_root)
    data_dir = joinpath(repo_root, "notebooks_data")
    coverage_rows = split_csv_lines(
        joinpath(data_dir, "9-CancerGenomics_coverage.csv")
    )
    matrix_rows = split_csv_lines(
        joinpath(data_dir, "9-CancerGenomics_comutation.csv")
    )
    provenance = JSON.parsefile(
        joinpath(data_dir, "9-CancerGenomics_provenance.json")
    )

    @assert coverage_rows[1] == ["gene", "patient_count"]
    genes = String[row[1] for row in coverage_rows[2:end]]
    coverage = parse.(Int, [row[2] for row in coverage_rows[2:end]])
    @assert genes == String.(matrix_rows[1][2:end])

    A = zeros(Int, length(genes), length(genes))
    for (i, row) in enumerate(matrix_rows[2:end])
        @assert row[1] == genes[i]
        A[i, :] = parse.(Int, row[2:end])
    end

    patient_count = Int(provenance["aggregate"]["patient_count"])
    @assert length(genes) == 33
    @assert patient_count == 196
    @assert A == A'
    @assert iszero(diag(A))
    @assert all(
        A[i, j] <= min(coverage[i], coverage[j])
        for i in eachindex(genes), j in eachindex(genes)
    )
    return (
        genes = genes,
        coverage = coverage,
        A = A,
        patient_count = patient_count,
    )
end

const CANCER_ALPHA = 0.45
aml = load_tcga_aml_aggregate(QUBONOTEBOOKS_REPO_DIR)

function cancer_energy(bits)
    return dot(bits, aml.A * bits) -
        CANCER_ALPHA * dot(aml.coverage, bits)
end

"""Build the aggregate-only cancer-genomics QUBO from Notebook 9."""
function build_cancer_model()
    model = Model()
    @variable(model, x[1:length(aml.genes)], Bin)
    @objective(
        model,
        Min,
        sum(
            aml.A[i, j] * x[i] * x[j]
            for i in eachindex(aml.genes), j in eachindex(aml.genes)
        ) - CANCER_ALPHA * sum(
            aml.coverage[i] * x[i] for i in eachindex(aml.genes)
        ),
    )
    return model, x
end

function distinct_patient_bounds(bits)
    selected = findall(==(1), bits)
    isempty(selected) && return (lower = 0, upper = 0)
    coverage_sum = sum(aml.coverage[selected])
    pair_overlap = sum(
        (aml.A[i, j] for i in selected for j in selected if i < j);
        init = 0,
    )
    return (
        lower = max(maximum(aml.coverage[selected]), coverage_sum - pair_overlap),
        upper = min(aml.patient_count, coverage_sum),
    )
end

function decode_cancer(bits)
    selected = findall(==(1), bits)
    pair_overlap = sum(
        (aml.A[i, j] for i in selected for j in selected if i < j);
        init = 0,
    )
    bounds = distinct_patient_bounds(bits)
    return (
        application_score = "coverage bounds=$(bounds.lower)–$(bounds.upper)",
        feasible = true,
        details = "size=$(length(selected)); pair co-mutation=$pair_overlap; aggregate-only",
    )
end

partition_baseline = exact_baseline(4, partition_energy)
maxcut_baseline = exact_baseline(4, maxcut_energy)
cover_baseline = exact_baseline(4, cover_energy)
order_baseline = exact_baseline(6, order_energy)

@assert partition_baseline.energy == 0 && partition_baseline.degeneracy == 2
@assert maxcut_baseline.energy == -7 && maxcut_baseline.degeneracy == 4
@assert cover_baseline.energy == 2 && cover_baseline.degeneracy == 2
@assert order_baseline.energy == 28 && order_baseline.degeneracy == 2

partition_problem = (
    name = "number partitioning",
    n = 4,
    build_model = build_partition_model,
    energy = partition_energy,
    decode = decode_partition,
    exact_optimum_known = true,
    exact = partition_baseline,
)
maxcut_problem = (
    name = "weighted Max-Cut",
    n = 4,
    build_model = build_maxcut_model,
    energy = maxcut_energy,
    decode = decode_maxcut,
    exact_optimum_known = true,
    exact = maxcut_baseline,
)
cover_problem = (
    name = "minimum vertex cover",
    n = 4,
    build_model = build_cover_model,
    energy = cover_energy,
    decode = decode_cover,
    exact_optimum_known = true,
    exact = cover_baseline,
)
order_problem = (
    name = "order partitioning",
    n = 6,
    build_model = build_order_model,
    energy = order_energy,
    decode = decode_order,
    exact_optimum_known = true,
    exact = order_baseline,
)
cancer_problem = (
    name = "cancer-genomics aggregate",
    n = length(aml.genes),
    build_model = build_cancer_model,
    energy = cancer_energy,
    decode = decode_cancer,
    exact_optimum_known = false,
    exact = nothing,
)

annealing_problems = (
    partition_problem,
    maxcut_problem,
    cover_problem,
    order_problem,
    cancer_problem,
)
println("Prepared five solver-interchangeable models.")

"""
    run_annealing(problem, config)

Build a problem once through its JuMP formulation, attach the selected
DWave.jl optimizer, validate every returned energy against the independent
application evaluator, and preserve occurrence counts and timing metadata.
"""
function run_annealing(problem, config)
    model, variables = problem.build_model()
    set_optimizer(model, config.optimizer)
    for (name, value) in config.attributes
        set_optimizer_attribute(model, name, value)
    end

    user_wall_seconds = @elapsed optimize!(model)
    @assert termination_status(model) == JuMP.MOI.LOCALLY_SOLVED

    sampleset = DWave.QUBOTools.solution(DWave.QUBOTools.backend(model))
    sample_reads = DWave.QUBOTools.reads.(sampleset)
    metadata = DWave.QUBOTools.metadata(sampleset)
    execution_mode = metadata["execution"]["mode"]
    dwave_timing = get(
        get(metadata, "dwave_info", Dict{String,Any}()),
        "timing",
        Dict{String,Any}(),
    )
    qpu_access_time_microseconds = execution_mode == "qpu" ?
        get(dwave_timing, "qpu_access_time", missing) : missing
    @assert length(sample_reads) == result_count(model)

    rows = [
        begin
            bits = round.(Int, value.(variables; result = result))
            reported_energy = objective_value(model; result = result)
            recomputed_energy = problem.energy(bits)
            @assert isapprox(
                reported_energy,
                recomputed_energy;
                atol = 1e-8,
                rtol = 0,
            )
            decoded = problem.decode(bits)
            @assert decoded.feasible isa Bool
            (
                bits = bits,
                reported_energy = reported_energy,
                recomputed_energy = recomputed_energy,
                reads = sample_reads[result],
                decoded = decoded,
            )
        end
        for result in 1:result_count(model)
    ]
    sort!(rows; by = row -> (row.reported_energy, Tuple(row.bits)))

    best_raw_energy = first(rows).recomputed_energy
    total_reads = sum(row.reads for row in rows)
    @assert total_reads == config.attributes["num_reads"]
    success_count = if problem.exact_optimum_known
        sum(
            row.reads
            for row in rows
            if isapprox(
                row.recomputed_energy,
                problem.exact.energy;
                atol = 1e-8,
                rtol = 0,
            )
        )
    else
        missing
    end
    success_probability = ismissing(success_count) ?
        missing : success_count / total_reads

    return (
        problem = problem.name,
        solver = config.label,
        model_size = problem.n,
        density = quadratic_density(problem.n, problem.energy),
        requested_reads = config.attributes["num_reads"],
        returned_states = length(rows),
        validated_states = length(rows),
        best_raw_energy = best_raw_energy,
        decoded = first(rows).decoded,
        feasibility = first(rows).decoded.feasible,
        success_count = success_count,
        success_probability = success_probability,
        algorithm_effective_seconds = metadata["time"]["effective"],
        user_wall_seconds = user_wall_seconds,
        execution_mode = execution_mode,
        qpu_access_time_microseconds = qpu_access_time_microseconds,
        rows = rows,
        metadata = metadata,
    )
end

local_results = [
    run_annealing(problem, local_config)
    for problem in annealing_problems
]

@assert all(
    result.returned_states == result.validated_states
    for result in local_results
)
@assert all(result.success_count > 0 for result in local_results[1:4])
@assert all(
    isapprox(
        row.reported_energy,
        row.recomputed_energy;
        atol = 1e-8,
        rtol = 0,
    )
    for result in local_results for row in result.rows
)

println(
    "Seeded DWave.Neal comparison: seed=$LOCAL_SEED, " *
    "reads=$LOCAL_READS, sweeps=$LOCAL_SWEEPS.",
)
println("Validated every returned state for all five models.")

function format_comparison_row(result)
    success = ismissing(result.success_count) ?
        "n/a (exact optimum not claimed)" :
        @sprintf(
            "%d/%.3f",
            result.success_count,
            result.success_probability,
        )
    local_effective = result.execution_mode == "qpu" ?
        "n/a" : @sprintf("%.6f", result.algorithm_effective_seconds)
    qpu_access = ismissing(result.qpu_access_time_microseconds) ?
        "n/a" : @sprintf("%.0f", result.qpu_access_time_microseconds)
    return @sprintf(
        "%s | %s | %d | %.3f | %d | %.3f | %s | %s | %s | %s | %.6f | %s",
        result.solver,
        result.problem,
        result.model_size,
        result.density,
        result.requested_reads,
        result.best_raw_energy,
        result.decoded.application_score,
        string(result.feasibility),
        success,
        local_effective,
        result.user_wall_seconds,
        qpu_access,
    )
end

function print_comparison_table(results)
    println(
        "solver | problem | n | density | reads | best raw energy | " *
        "decoded score | feasible | exact successes/probability | " *
        "local effective s | user wall s | qpu access μs",
    )
    foreach(result -> println(format_comparison_row(result)), results)
    return nothing
end

# Credential-free regression check for the live-QPU reporting branch.
synthetic_qpu_result = (
    solver = "D-Wave QPU (synthetic)",
    problem = "weighted Max-Cut",
    model_size = 4,
    density = 5 / 6,
    requested_reads = 100,
    best_raw_energy = -7.0,
    decoded = (application_score = "cut weight=7",),
    feasibility = true,
    success_count = 80,
    success_probability = 0.8,
    algorithm_effective_seconds = 0.001234,
    user_wall_seconds = 0.25,
    execution_mode = "qpu",
    qpu_access_time_microseconds = 1_234,
)
synthetic_qpu_row = format_comparison_row(synthetic_qpu_result)
@assert synthetic_qpu_row == "D-Wave QPU (synthetic) | weighted Max-Cut | " *
    "4 | 0.833 | 100 | -7.000 | cut weight=7 | true | 80/0.800 | " *
    "n/a | 0.250000 | 1234"

print_comparison_table(local_results)

function qpu_opt_in_requested()
    raw_value = get(ENV, "QUBONOTEBOOKS_ANNEALING_ENABLE_QPU", "0")
    raw_value in ("0", "1") || error(
        "Set QUBONOTEBOOKS_ANNEALING_ENABLE_QPU to 0 or 1.",
    )
    return raw_value == "1"
end

function require_qpu_credentials()
    token = get(ENV, "DWAVE_API_TOKEN", "")
    isempty(strip(token)) && error(
        "QPU execution was requested, but DWAVE_API_TOKEN is missing. " *
        "Set it in the process environment or hosted secret store, then rerun.",
    )
    return nothing
end

"""
    sanitized_qpu_metadata(metadata)

Return a small allowlist of reviewable solver, timing, and embedding fields.
Opaque problem identifiers and credentials are intentionally omitted.
"""
function sanitized_qpu_metadata(metadata)
    dwave_info = get(metadata, "dwave_info", Dict{String,Any}())
    timing = get(dwave_info, "timing", Dict{String,Any}())
    chip = get(dwave_info, "chip_info", Dict{String,Any}())
    context = get(dwave_info, "embedding_context", Dict{String,Any}())
    embedding = DWave.embedding(metadata)
    return Dict{String,Any}(
        "solver_name" => get(chip, "solver_name", "unavailable"),
        "topology" => get(chip, "topology", "unavailable"),
        "num_qubits" => get(chip, "num_qubits", "unavailable"),
        "qpu_access_time" => get(timing, "qpu_access_time", "unavailable"),
        "qpu_access_time_unit" => "microseconds",
        "chain_strength" => get(context, "chain_strength", "unavailable"),
        "embedding_parameters" => get(
            context,
            "embedding_parameters",
            "unavailable",
        ),
        "embedded_variable_count" => isnothing(embedding) ?
            0 : length(embedding),
    )
end

qpu_requested = qpu_opt_in_requested()
qpu_hardware_required = get(
    ENV,
    "QUBONOTEBOOKS_ANNEALING_REQUIRE_QPU",
    "0",
) == "1"
qpu_hardware_submitted = false
qpu_result = nothing
qpu_summary = nothing

if qpu_requested
    require_qpu_credentials()
    qpu_config = (
        label = "D-Wave QPU",
        optimizer = DWave.Optimizer,
        attributes = Dict{String,Any}(
            "num_reads" => 100,
            "annealing_time" => 20.0,
            "return_embedding" => true,
        ),
    )
    qpu_result = run_annealing(maxcut_problem, qpu_config)
    qpu_hardware_submitted = qpu_result.execution_mode == "qpu"
    qpu_summary = sanitized_qpu_metadata(qpu_result.metadata)
    println("Neal/QPU comparison for weighted Max-Cut:")
    print_comparison_table((local_results[2], qpu_result))
    println("Sanitized QPU metadata:")
    display(qpu_summary)
else
    println(
        "D-Wave QPU execution is disabled. Set " *
        "QUBONOTEBOOKS_ANNEALING_ENABLE_QPU=1 and provide " *
        "DWAVE_API_TOKEN through a secret store to opt in.",
    )
end

if qpu_hardware_required && !qpu_hardware_submitted
    error("D-Wave QPU verification was required, but no job was submitted.")
end

if isnothing(qpu_result)
    println("Topology and embedding plots skipped because QPU execution is disabled.")
else
    working_graph = DWave.WorkingGraph(qpu_result.metadata)
    display(DWave.draw_topology(working_graph; node_size = 1, with_labels = false))
    display(DWave.draw_embedding(qpu_result.metadata; node_size = 2))
end

# EXERCISE 1: Compare the Metropolis acceptance probability for ΔE=2 at
# temperatures T=4 and T=0.5. Which run explores uphill moves more readily?
nothing

# EXERCISE 2: For the four exact-checked models, print the number of successful
# reads and verify that probability equals successes divided by total reads.
nothing

# EXERCISE 3: Name the three timing fields in this notebook and state why the
# local and QPU values should not be compared as equivalent runtime.
nothing