# -----------------------------------------------------------------------------
# CELL 1: Imports, environment setup, and project root resolution
# -----------------------------------------------------------------------------

import copy
import sys
import os
import math
import time

import numpy as np
import pandas as pd
import torch
import torch.optim as optim
from scipy.optimize import minimize
import matplotlib.pyplot as plt
import networkx as nx
from sklearn.datasets import fetch_openml
from IPython.display import display


def find_project_root():
    """Walk upward from the current working directory until /src is found."""
    candidate = os.getcwd()
    while True:
        if os.path.isdir(os.path.join(candidate, "src")):
            return candidate
        parent = os.path.dirname(candidate)
        if parent == candidate:
            return os.getcwd()
        candidate = parent


proj_root = find_project_root()
if proj_root not in sys.path:
    sys.path.insert(0, proj_root)

torch.set_num_threads(max(1, min(4, os.cpu_count() or 1)))
np.random.seed(7)
torch.manual_seed(7)

from src.gnn import SimpleGCN

print("=" * 78)
print("Depth-2 Hybrid Quantum-Classical MaxCut Demo on Real Genomics Data")
print("=" * 78)
print(f"Project root        : {proj_root}")
print(f"NumPy               : {np.__version__}")
print(f"PyTorch             : {torch.__version__}")
print(f"NetworkX            : {nx.__version__}")
print(f"Matplotlib          : {plt.matplotlib.__version__}")
print("\nThis notebook now evaluates a stronger protocol:")
print("  1. exact depth-2 QAOA on transcriptomic co-expression graphs,")
print("  2. transcriptomic domain adaptation of a graph model,")
print("  3. held-out evaluation against exact MaxCut-derived references.")

==============================================================================
Depth-2 Hybrid Quantum-Classical MaxCut Demo on Real Genomics Data
==============================================================================
Project root        : /Users/mohuyn/Library/CloudStorage/OneDrive-SAS/Documents/GitHub/Hybrid-Quantum-Graph-AI-QAOA-GNN-Biomedical-Optimization
NumPy               : 2.4.2
PyTorch             : 2.10.0
NetworkX            : 3.6.1
Matplotlib          : 3.10.8

This notebook now evaluates a stronger protocol:
  1. exact depth-2 QAOA on transcriptomic co-expression graphs,
  2. transcriptomic domain adaptation of a graph model,
  3. held-out evaluation against exact MaxCut-derived references.

# ═══════════════════════════════════════════════════════════════════════════════
# CELL 2 — Adaptive Quantum GCN: load weights, timed inference, quality–latency analysis
# ═══════════════════════════════════════════════════════════════════════════════

# ── 2a. Instantiate the Adaptive Quantum GCN (SimpleGCN backbone) ────────────
model = SimpleGCN(in_feats=1, hidden=32, out_feats=2, p=p_depth)

n_params = sum(pa.numel() for pa in model.parameters() if pa.requires_grad)
print('═' * 72)
print('  Adaptive Quantum GCN — Architecture Summary')
print('═' * 72)
print(f'  Backbone             : SimpleGCN')
print(f'  Trainable parameters : {n_params:,}')
print(f'  Forward pass         : X[n×1] → GCN(32) → GCN(32) → MeanPool → '
      f'Linear({2*p_depth}) → (γ̂,β̂)')
print(f'  Dense-adj fallback   : {not PYG_AVAILABLE}')

# ── 2b. Load weights (or report random-weight fallback) ──────────────────────
model_path = os.path.join(proj_root, 'model.pt')
weights_trained = os.path.exists(model_path)
if weights_trained:
    model.load_state_dict(torch.load(model_path, map_location='cpu', weights_only=True))
    print(f'\n  ✓ Loaded trained weights: {model_path}')
else:
    print(f'\n  ⚠  model.pt not found at {model_path}')
    print('     Using random weights — quality ratio is NOT meaningful.')
    print('     Train with: python src/train.py')
model.eval()

# ── 2c. Prepare GNN input tensors ────────────────────────────────────────────
# Augmented adjacency Â = A + I  (adds self-loops for stable propagation)
A_mat = nx.to_numpy_array(G_nx) + np.eye(n)                    # [n, n]
X_mat = A_mat.sum(axis=1, keepdims=True).astype(np.float32)    # augmented degree [n, 1]

At = torch.tensor(A_mat, dtype=torch.float32)
Xt = torch.tensor(X_mat, dtype=torch.float32)

print(f'\n  GNN Inputs:')
print(f'  X (aug. degree) : {X_mat.flatten().tolist()}')
print(f'  Â (row 0)       : {A_mat[0].tolist()}')

# ── 2d. Warm-up pass (ensures JIT/BLAS initialization is not timed) ──────────
with torch.no_grad():
    _ = model(Xt, At)

# ── 2e. Timed inference — 10 passes, report median ───────────────────────────
N_TIME_PASSES = 10
gnn_times = []
for _ in range(N_TIME_PASSES):
    t0 = time.perf_counter()
    with torch.no_grad():
        out = model(Xt, At).view(-1).numpy()
    gnn_times.append(time.perf_counter() - t0)

gnn_time = float(np.median(gnn_times))       # use median latency (robust to outliers)
pred_g   = out[:p_depth]
pred_b   = out[p_depth:]

# ── 2f. Evaluate predicted angles in QAOA ────────────────────────────────────
val_pred   = expected_cut(n, edges, qaoa_state(n, edges, pred_g, pred_b))
quality_q  = val_pred / val_opt if val_opt > 0 else float('nan')
speedup    = classical_time / gnn_time if gnn_time > 0 else float('inf')

# ── 2g. Print results ─────────────────────────────────────────────────────────
print()
print('─' * 72)
print('  Adaptive Quantum GCN Inference Results')
print('─' * 72)
print(f'  γ̂ (AQ-GCN)       : {pred_g}')
print(f'  β̂ (AQ-GCN)       : {pred_b}')
print(f'  γ* (classical)    : {opt_g}')
print(f'  β* (classical)    : {opt_b}')
print(f'  Δγ                : {abs(pred_g[0] - opt_g[0]):.4f} rad')
print(f'  Δβ                : {abs(pred_b[0] - opt_b[0]):.4f} rad')
print()
print(f'  ⟨C⟩_AQ-GCN        : {val_pred:.6f}')
print(f'  ⟨C⟩_opt           : {val_opt:.6f}')
print(f'  C*  (exact)       : {best_cut}')
print(f'  Quality q         : {quality_q:.4f}  ({quality_q*100:.2f}% of classical opt)')
print()
print(f'  AQ-GCN latency    : {gnn_time*1e6:.2f} µs  (median of {N_TIME_PASSES})')
print(f'  Classical time    : {classical_time*1e3:.2f} ms  ({N_RESTARTS} restarts)')
print(f'  Speedup           : {speedup:,.0f}×')
print()
print(f'  Weights trained   : {weights_trained}')

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[6], line 6
      1 # ═══════════════════════════════════════════════════════════════════════════════
      2 # CELL 2 — Adaptive Quantum GCN: load weights, timed inference, quality–latency analysis
      3 # ═══════════════════════════════════════════════════════════════════════════════
      4 
      5 # ── 2a. Instantiate the Adaptive Quantum GCN (SimpleGCN backbone) ────────────
----> 6 model = SimpleGCN(in_feats=1, hidden=32, out_feats=2, p=p_depth)
      8 n_params = sum(pa.numel() for pa in model.parameters() if pa.requires_grad)
      9 print('═' * 72)

NameError: name 'p_depth' is not defined

# -----------------------------------------------------------------------------
# CELL 3: Helper functions for exact MaxCut, QAOA evaluation, and genomics graphs
# -----------------------------------------------------------------------------


def build_cut_diagonal(n, edges):
    """Return the cut value for every computational basis state."""
    cut_diagonal = np.zeros(2 ** n, dtype=np.float64)
    for state_index in range(2 ** n):
        bits = [(state_index >> bit) & 1 for bit in range(n)]
        cut_diagonal[state_index] = sum(bits[u] != bits[v] for u, v in edges)
    return cut_diagonal


def apply_rx_all(state, n, beta):
    """Apply the same Rx rotation to every qubit."""
    rx = np.array(
        [
            [np.cos(beta), -1j * np.sin(beta)],
            [-1j * np.sin(beta), np.cos(beta)],
        ],
        dtype=np.complex128,
    )
    psi = state.reshape((2,) * n)
    for axis in range(n):
        psi = np.moveaxis(psi, axis, 0)
        psi = np.tensordot(rx, psi, axes=([1], [0]))
        psi = np.moveaxis(psi, 0, axis)
    return psi.reshape(-1)


def qaoa_state_fast(cut_diagonal, gammas, betas):
    """Exact statevector evaluation specialized for repeated notebook benchmarking."""
    num_states = cut_diagonal.shape[0]
    n_qubits = int(np.log2(num_states))
    state = np.ones(num_states, dtype=np.complex128) / np.sqrt(num_states)
    for gamma, beta in zip(gammas, betas):
        state = state * np.exp(-1j * gamma * cut_diagonal)
        state = apply_rx_all(state, n_qubits, beta)
    return state


def expected_cut_fast(cut_diagonal, state):
    """Expected cut value computed from the state probabilities."""
    return float(np.dot(cut_diagonal, np.abs(state) ** 2))


def brute_force_maxcut(n, edges):
    """Return the exact MaxCut value and the corresponding partition mask."""
    best_cut = -1
    best_mask = 0
    for mask in range(1, 2 ** n):
        cut_value = sum(1 for u, v in edges if ((mask >> u) & 1) != ((mask >> v) & 1))
        if cut_value > best_cut:
            best_cut = cut_value
            best_mask = mask
    return best_cut, best_mask


def normalize_angles(raw_angles, p):
    """Project unconstrained outputs into the standard QAOA search domain."""
    raw_angles = np.asarray(raw_angles, dtype=np.float64).reshape(-1)
    if raw_angles.size < 2 * p:
        raise ValueError(f"Expected at least {2 * p} raw angles, received {raw_angles.size}.")
    gammas = np.mod(raw_angles[:p], math.pi)
    betas = np.mod(raw_angles[p : 2 * p], math.pi / 2)
    return gammas, betas


def qaoa_value_for_angles(cut_diagonal, gammas, betas):
    state = qaoa_state_fast(cut_diagonal, gammas, betas)
    return expected_cut_fast(cut_diagonal, state), state


def classical_optimize_instance(instance, p, num_starts=8, maxiter=320, seed=0):
    """Multi-start Nelder-Mead search for one benchmark graph."""
    cut_diagonal = instance["cut_diagonal"]
    rng = np.random.default_rng(seed)
    best = None

    def objective(x):
        gammas, betas = normalize_angles(x, p)
        value, _ = qaoa_value_for_angles(cut_diagonal, gammas, betas)
        return -value

    for _ in range(num_starts):
        x0 = np.concatenate(
            [
                rng.uniform(0.0, math.pi, size=p),
                rng.uniform(0.0, math.pi / 2, size=p),
            ]
        )
        result = minimize(
            objective,
            x0,
            method="Nelder-Mead",
            options={"maxiter": maxiter, "xatol": 1e-6, "fatol": 1e-6},
        )
        gammas, betas = normalize_angles(result.x, p)
        value, state = qaoa_value_for_angles(cut_diagonal, gammas, betas)
        candidate = {
            "gammas": gammas,
            "betas": betas,
            "value": value,
            "state": state,
            "nit": result.nit,
            "nfev": result.nfev,
            "success": bool(result.success),
            "raw_angles": np.concatenate([gammas, betas]),
        }
        if best is None or candidate["value"] > best["value"]:
            best = candidate
    return best


def predict_instance_with_gnn(instance, model, p):
    """Run one forward pass of the graph model and evaluate the predicted angles."""
    adjacency_tensor = torch.tensor(instance["adjacency"], dtype=torch.float32)
    feature_tensor = torch.tensor(instance["features"], dtype=torch.float32)

    with torch.no_grad():
        _ = model(feature_tensor, adjacency_tensor)

    start_time = time.perf_counter()
    with torch.no_grad():
        raw_output = model(feature_tensor, adjacency_tensor).view(-1).cpu().numpy()
    inference_time = time.perf_counter() - start_time

    gammas, betas = normalize_angles(raw_output, p)
    value, state = qaoa_value_for_angles(instance["cut_diagonal"], gammas, betas)
    return {
        "gammas": gammas,
        "betas": betas,
        "value": value,
        "state": state,
        "inference_time": inference_time,
        "raw_output": raw_output,
    }


def select_gene_panel(expression_frame, top_gene_count):
    """Select the highest-variance genes from the transcriptomic matrix."""
    gene_table = (
        expression_frame.var(axis=0)
        .sort_values(ascending=False)
        .head(top_gene_count)
        .rename("variance")
        .reset_index()
        .rename(columns={"index": "gene"})
    )
    gene_table["rank"] = np.arange(1, len(gene_table) + 1)
    return gene_table[["rank", "gene", "variance"]]


def build_gene_correlation_graph(expression_frame, gene_table, target_edge_count):
    """Build a connected co-expression graph using absolute Pearson correlations."""
    genes = gene_table["gene"].tolist()
    correlation_matrix = expression_frame[genes].corr().abs().fillna(0.0).copy()
    correlation_values = correlation_matrix.to_numpy(copy=True)
    np.fill_diagonal(correlation_values, 0.0)
    correlation_matrix.iloc[:, :] = correlation_values

    complete_graph = nx.Graph()
    for gene_index, gene_name in enumerate(genes):
        complete_graph.add_node(gene_index, gene=gene_name)
    for i, gene_i in enumerate(genes):
        for j in range(i + 1, len(genes)):
            gene_j = genes[j]
            complete_graph.add_edge(i, j, weight=float(correlation_matrix.loc[gene_i, gene_j]))

    spanning_tree = nx.maximum_spanning_tree(complete_graph, weight="weight")
    remaining_edges = sorted(
        (
            (u, v, data["weight"])
            for u, v, data in complete_graph.edges(data=True)
            if not spanning_tree.has_edge(u, v)
        ),
        key=lambda item: item[2],
        reverse=True,
    )

    graph = nx.Graph()
    for node_index, gene_name in enumerate(genes):
        graph.add_node(node_index, gene=gene_name)
    for u, v, data in spanning_tree.edges(data=True):
        graph.add_edge(u, v, weight=data["weight"])
    for u, v, weight in remaining_edges:
        if graph.number_of_edges() >= target_edge_count:
            break
        graph.add_edge(u, v, weight=weight)

    edge_rows = []
    for u, v, data in graph.edges(data=True):
        edge_rows.append(
            {
                "gene_u": graph.nodes[u]["gene"],
                "gene_v": graph.nodes[v]["gene"],
                "abs_correlation": data["weight"],
            }
        )
    edge_table = pd.DataFrame(edge_rows).sort_values("abs_correlation", ascending=False).reset_index(drop=True)
    return graph, correlation_matrix, edge_table


def stratified_subsample_indices(labels, sample_size, seed):
    """Draw a label-balanced subsample without replacement."""
    labels = pd.Series(labels)
    label_counts = labels.value_counts().sort_index()
    desired = label_counts / label_counts.sum() * sample_size
    counts = np.floor(desired).astype(int)
    remainder = sample_size - int(counts.sum())
    if remainder > 0:
        fractional = (desired - counts).sort_values(ascending=False)
        for label in fractional.index[:remainder]:
            counts.loc[label] += 1

    rng = np.random.default_rng(seed)
    chosen = []
    for label, count in counts.items():
        label_indices = labels[labels == label].index.to_numpy()
        chosen.extend(rng.choice(label_indices, size=int(count), replace=False).tolist())
    rng.shuffle(chosen)
    return chosen


def create_graph_instance(graph_id, graph, sample_indices, labels, split_name):
    """Package graph metadata and exact optimization targets for later analysis."""
    n = graph.number_of_nodes()
    edges = list(graph.edges())
    best_cut, best_mask = brute_force_maxcut(n, edges)
    adjacency = nx.to_numpy_array(graph, dtype=np.float64) + np.eye(n)
    features = adjacency.sum(axis=1, keepdims=True).astype(np.float32)
    return {
        "graph_id": graph_id,
        "split": split_name,
        "graph": graph,
        "n": n,
        "edges": edges,
        "edge_count": len(edges),
        "density": nx.density(graph),
        "adjacency": adjacency,
        "features": features,
        "cut_diagonal": build_cut_diagonal(n, edges),
        "best_cut": best_cut,
        "best_mask": best_mask,
        "sample_indices": sample_indices,
        "sample_count": len(sample_indices),
        "class_balance": pd.Series(labels.loc[sample_indices]).value_counts().sort_index().to_dict(),
        "gene_labels": [graph.nodes[node]["gene"] for node in graph.nodes()],
    }


def build_graph_split(expression_frame, labels, gene_table, target_edge_count, split_name, split_size, subsample_size, base_seed):
    """Construct one split of patient-resampled transcriptomic graphs."""
    instances = []
    for offset in range(split_size):
        graph_seed = base_seed + offset
        sample_indices = stratified_subsample_indices(labels, subsample_size, seed=graph_seed)
        subset_expression = expression_frame.loc[sample_indices]
        graph, corr_matrix, edge_table = build_gene_correlation_graph(
            subset_expression,
            gene_table,
            target_edge_count,
        )
        instance = create_graph_instance(graph_seed, graph, sample_indices, labels, split_name=split_name)
        instance["correlation_matrix"] = corr_matrix
        instance["edge_table"] = edge_table
        instances.append(instance)
    return instances


def build_genomics_benchmark(
    dataset_name="prostate",
    version=1,
    top_gene_count=10,
    target_edge_count=18,
    benchmark_size=6,
    benchmark_seed=42,
    adaptation_size=24,
    adaptation_seed=200,
    subsample_size=60,
):
    """Fetch the real dataset and derive representative, adaptation, and held-out splits."""
    dataset = fetch_openml(name=dataset_name, version=version, as_frame=True)
    expression_frame = dataset.data.apply(pd.to_numeric, errors="coerce")
    expression_frame = expression_frame.loc[:, expression_frame.notna().all(axis=0)]
    labels = pd.Series(dataset.target, name=dataset.target_names[0] if dataset.target_names else "class")

    gene_table = select_gene_panel(expression_frame, top_gene_count)
    representative_graph, representative_corr, representative_edge_table = build_gene_correlation_graph(
        expression_frame,
        gene_table,
        target_edge_count,
    )
    representative = create_graph_instance(
        graph_id=0,
        graph=representative_graph,
        sample_indices=expression_frame.index.tolist(),
        labels=labels,
        split_name="representative",
    )
    representative["correlation_matrix"] = representative_corr
    representative["edge_table"] = representative_edge_table

    adaptation_instances = build_graph_split(
        expression_frame=expression_frame,
        labels=labels,
        gene_table=gene_table,
        target_edge_count=target_edge_count,
        split_name="adaptation",
        split_size=adaptation_size,
        subsample_size=subsample_size,
        base_seed=adaptation_seed,
    )
    benchmark_instances = build_graph_split(
        expression_frame=expression_frame,
        labels=labels,
        gene_table=gene_table,
        target_edge_count=target_edge_count,
        split_name="benchmark",
        split_size=benchmark_size,
        subsample_size=subsample_size,
        base_seed=benchmark_seed,
    )

    dataset_summary = pd.DataFrame(
        {
            "metric": [
                "Dataset name",
                "Samples",
                "Gene-expression features",
                "Selected genes for graph",
                "Representative graph nodes",
                "Representative graph edges",
                "Adaptation graphs",
                "Held-out benchmark graphs",
                "Patients per resampled graph",
            ],
            "value": [
                dataset.details.get("name", dataset_name),
                expression_frame.shape[0],
                expression_frame.shape[1],
                top_gene_count,
                representative["n"],
                representative["edge_count"],
                adaptation_size,
                benchmark_size,
                subsample_size,
            ],
        }
    )

    def build_split_overview(instances):
        return pd.DataFrame(
            [
                {
                    "graph_id": instance["graph_id"],
                    "split": instance["split"],
                    "samples": instance["sample_count"],
                    "edges": instance["edge_count"],
                    "density": round(instance["density"], 3),
                    "exact_maxcut": instance["best_cut"],
                }
                for instance in instances
            ]
        )

    return {
        "dataset": dataset,
        "expression_frame": expression_frame,
        "labels": labels,
        "gene_table": gene_table,
        "representative": representative,
        "adaptation_instances": adaptation_instances,
        "benchmark_instances": benchmark_instances,
        "dataset_summary": dataset_summary,
        "adaptation_overview": build_split_overview(adaptation_instances),
        "benchmark_overview": build_split_overview(benchmark_instances),
    }


def attach_classical_targets(instances, p, num_starts=8, maxiter=320, seed_offset=0):
    """Create exact depth-p QAOA targets for a list of transcriptomic graphs."""
    enriched_instances = []
    rows = []
    for index, instance in enumerate(instances):
        reference = classical_optimize_instance(
            instance,
            p=p,
            num_starts=num_starts,
            maxiter=maxiter,
            seed=seed_offset + int(instance["graph_id"]) + index,
        )
        enriched = dict(instance)
        enriched["classical_reference"] = reference
        enriched["target_angles"] = np.concatenate([reference["gammas"], reference["betas"]]).astype(np.float32)
        enriched_instances.append(enriched)
        rows.append(
            {
                "graph_id": instance["graph_id"],
                "exact_maxcut": instance["best_cut"],
                "classical_ratio": reference["value"] / instance["best_cut"],
                "iterations": reference["nit"],
                "evaluations": reference["nfev"],
            }
        )
    return enriched_instances, pd.DataFrame(rows)


def train_adapted_qaoa_gnn(train_instances, p, hidden_dim=64, epochs=500, lr=5e-3, weight_decay=1e-4, patience=50, seed=7):
    """Fit a depth-p graph model to transcriptomic graph families using classical targets."""
    torch.manual_seed(seed)
    np.random.seed(seed)

    trained_model = SimpleGCN(in_feats=1, hidden=hidden_dim, out_feats=2, p=p)
    optimizer = optim.Adam(trained_model.parameters(), lr=lr, weight_decay=weight_decay)

    best_state = copy.deepcopy(trained_model.state_dict())
    best_loss = float("inf")
    best_epoch = 0
    stale_epochs = 0
    loss_history = []

    for epoch in range(1, epochs + 1):
        trained_model.train()
        running_loss = 0.0
        for instance in train_instances:
            adjacency_tensor = torch.tensor(instance["adjacency"], dtype=torch.float32)
            feature_tensor = torch.tensor(instance["features"], dtype=torch.float32)
            target_tensor = torch.tensor(instance["target_angles"], dtype=torch.float32)

            prediction = trained_model(feature_tensor, adjacency_tensor).view(-1)
            loss = ((prediction - target_tensor) ** 2).mean()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            running_loss += float(loss.item())

        mean_loss = running_loss / max(1, len(train_instances))
        loss_history.append(mean_loss)
        if mean_loss + 1e-8 < best_loss:
            best_loss = mean_loss
            best_epoch = epoch
            best_state = copy.deepcopy(trained_model.state_dict())
            stale_epochs = 0
        else:
            stale_epochs += 1

        if stale_epochs >= patience:
            break

    trained_model.load_state_dict(best_state)
    trained_model.eval()
    return {
        "model": trained_model,
        "history": loss_history,
        "best_loss": best_loss,
        "best_epoch": best_epoch,
        "epochs_run": len(loss_history),
    }


print("Helper functions ready.")
print("  Exact MaxCut            : brute_force_maxcut")
print("  QAOA evaluator          : qaoa_state_fast / expected_cut_fast")
print("  Classical benchmark     : classical_optimize_instance")
print("  Learned inference       : predict_instance_with_gnn")
print("  Graph split builder     : build_genomics_benchmark")
print("  Target generation       : attach_classical_targets")
print("  Domain adaptation       : train_adapted_qaoa_gnn")

Helper functions ready.
  Exact MaxCut            : brute_force_maxcut
  QAOA evaluator          : qaoa_state_fast / expected_cut_fast
  Classical benchmark     : classical_optimize_instance
  Learned inference       : predict_instance_with_gnn
  Graph split builder     : build_genomics_benchmark
  Target generation       : attach_classical_targets
  Domain adaptation       : train_adapted_qaoa_gnn

# -----------------------------------------------------------------------------
# CACHE OVERRIDE: local reduced-panel cache for repeatable Run All execution
# -----------------------------------------------------------------------------

import json

CACHE_PANEL_PATH = os.path.join(proj_root, "data", "prostate_top10_variance_panel.csv.gz")
CACHE_META_PATH = os.path.join(proj_root, "data", "prostate_top10_variance_panel_meta.json")

def _package_genomics_bundle(
    expression_frame,
    labels,
    gene_table,
    dataset_name_display,
    full_feature_count,
    target_edge_count,
    benchmark_size,
    adaptation_size,
    subsample_size,
    benchmark_seed,
    adaptation_seed,
    cache_status,
    source_note,
    split_dataset=None,
):
    representative_graph, representative_corr, representative_edge_table = build_gene_correlation_graph(
        expression_frame,
        gene_table,
        target_edge_count,
    )
    representative = create_graph_instance(
        graph_id=0,
        graph=representative_graph,
        sample_indices=expression_frame.index.tolist(),
        labels=labels,
        split_name="representative",
    )
    representative["correlation_matrix"] = representative_corr
    representative["edge_table"] = representative_edge_table

    adaptation_instances = build_graph_split(
        expression_frame=expression_frame,
        labels=labels,
        gene_table=gene_table,
        target_edge_count=target_edge_count,
        split_name="adaptation",
        split_size=adaptation_size,
        subsample_size=subsample_size,
        base_seed=adaptation_seed,
    )
    benchmark_instances = build_graph_split(
        expression_frame=expression_frame,
        labels=labels,
        gene_table=gene_table,
        target_edge_count=target_edge_count,
        split_name="benchmark",
        split_size=benchmark_size,
        subsample_size=subsample_size,
        base_seed=benchmark_seed,
    )

    dataset_summary = pd.DataFrame(
        {
            "metric": [
                "Dataset name",
                "Samples",
                "Gene-expression features",
                "Selected genes for graph",
                "Representative graph nodes",
                "Representative graph edges",
                "Adaptation graphs",
                "Held-out benchmark graphs",
                "Patients per resampled graph",
                "Execution source",
            ],
            "value": [
                dataset_name_display,
                expression_frame.shape[0],
                full_feature_count,
                len(gene_table),
                representative["n"],
                representative["edge_count"],
                adaptation_size,
                benchmark_size,
                subsample_size,
                source_note,
            ],
        }
    )

    def build_split_overview(instances):
        return pd.DataFrame(
            [
                {
                    "graph_id": instance["graph_id"],
                    "split": instance["split"],
                    "samples": instance["sample_count"],
                    "edges": instance["edge_count"],
                    "density": round(instance["density"], 3),
                    "exact_maxcut": instance["best_cut"],
                }
                for instance in instances
            ]
        )

    return {
        "dataset": split_dataset,
        "expression_frame": expression_frame,
        "labels": labels,
        "gene_table": gene_table,
        "representative": representative,
        "adaptation_instances": adaptation_instances,
        "benchmark_instances": benchmark_instances,
        "dataset_summary": dataset_summary,
        "adaptation_overview": build_split_overview(adaptation_instances),
        "benchmark_overview": build_split_overview(benchmark_instances),
        "full_feature_count": full_feature_count,
        "cache_status": cache_status,
        "cache_paths": {
            "panel": CACHE_PANEL_PATH,
            "meta": CACHE_META_PATH,
        },
    }

def build_genomics_benchmark(
    dataset_name="prostate",
    version=1,
    top_gene_count=10,
    target_edge_count=18,
    benchmark_size=6,
    benchmark_seed=42,
    adaptation_size=24,
    adaptation_seed=200,
    subsample_size=60,
):
    """Fetch the real dataset once, then reuse a local reduced-panel cache on later runs."""
    use_local_cache = dataset_name == "prostate" and version == 1

    if use_local_cache and os.path.exists(CACHE_PANEL_PATH) and os.path.exists(CACHE_META_PATH):
        try:
            with open(CACHE_META_PATH, "r", encoding="utf-8") as handle:
                cache_meta = json.load(handle)
            cached_panel = pd.read_csv(CACHE_PANEL_PATH, compression="gzip")
            labels = pd.Series(
                cached_panel.pop("__target__").astype(str),
                name=cache_meta.get("label_name", "class"),
            )
            sample_ids = cached_panel.pop("__sample_id__").astype(str)
            expression_frame = cached_panel
            expression_frame.index = sample_ids
            gene_table = pd.DataFrame(cache_meta["gene_table"])

            if int(cache_meta.get("top_gene_count", len(gene_table))) == top_gene_count:
                return _package_genomics_bundle(
                    expression_frame=expression_frame,
                    labels=labels,
                    gene_table=gene_table,
                    dataset_name_display=cache_meta.get("dataset_name", "Prostate"),
                    full_feature_count=int(cache_meta.get("full_feature_count", expression_frame.shape[1])),
                    target_edge_count=target_edge_count,
                    benchmark_size=benchmark_size,
                    adaptation_size=adaptation_size,
                    subsample_size=subsample_size,
                    benchmark_seed=benchmark_seed,
                    adaptation_seed=adaptation_seed,
                    cache_status=f"Loaded cached reduced transcriptomic panel from {CACHE_PANEL_PATH}",
                    source_note="Local reduced-panel cache (offline-safe after first fetch)",
                    split_dataset=None,
                )
        except Exception as exc:
            print(f"Cache load failed, falling back to OpenML fetch: {exc}")

    dataset = fetch_openml(name=dataset_name, version=version, as_frame=True)
    expression_frame_full = dataset.data.apply(pd.to_numeric, errors="coerce")
    expression_frame_full = expression_frame_full.loc[:, expression_frame_full.notna().all(axis=0)]
    labels = pd.Series(dataset.target, name=dataset.target_names[0] if dataset.target_names else "class")

    full_feature_count = expression_frame_full.shape[1]
    gene_table = select_gene_panel(expression_frame_full, top_gene_count)
    selected_genes = gene_table["gene"].tolist()
    expression_frame = expression_frame_full.loc[:, selected_genes].copy()

    cache_status = "OpenML fetch used directly (no local cache written)."
    source_note = "Direct OpenML fetch"
    if use_local_cache:
        os.makedirs(os.path.dirname(CACHE_PANEL_PATH), exist_ok=True)
        cache_frame = expression_frame.copy()
        cache_frame.insert(0, "__sample_id__", expression_frame.index.astype(str))
        cache_frame["__target__"] = labels.astype(str).to_numpy()
        cache_frame.to_csv(CACHE_PANEL_PATH, index=False, compression="gzip")

        cache_meta = {
            "dataset_name": dataset.details.get("name", dataset_name),
            "dataset_version": version,
            "full_feature_count": int(full_feature_count),
            "top_gene_count": int(top_gene_count),
            "target_edge_count": int(target_edge_count),
            "label_name": labels.name,
            "gene_table": gene_table.to_dict(orient="records"),
        }
        with open(CACHE_META_PATH, "w", encoding="utf-8") as handle:
            json.dump(cache_meta, handle, indent=2)

        cache_status = f"Fetched OpenML data and refreshed local reduced-panel cache at {CACHE_PANEL_PATH}"
        source_note = "Fresh OpenML fetch + reduced-panel cache refresh"

    return _package_genomics_bundle(
        expression_frame=expression_frame,
        labels=labels,
        gene_table=gene_table,
        dataset_name_display=dataset.details.get("name", dataset_name),
        full_feature_count=full_feature_count,
        target_edge_count=target_edge_count,
        benchmark_size=benchmark_size,
        adaptation_size=adaptation_size,
        subsample_size=subsample_size,
        benchmark_seed=benchmark_seed,
        adaptation_seed=adaptation_seed,
        cache_status=cache_status,
        source_note=source_note,
        split_dataset=dataset,
    )

print("Local reduced-panel cache override ready.")
print(f"  Panel cache path : {CACHE_PANEL_PATH}")
print(f"  Meta cache path  : {CACHE_META_PATH}")

Local reduced-panel cache override ready.
  Panel cache path : /Users/mohuyn/Library/CloudStorage/OneDrive-SAS/Documents/GitHub/Hybrid-Quantum-Graph-AI-QAOA-GNN-Biomedical-Optimization/data/prostate_top10_variance_panel.csv.gz
  Meta cache path  : /Users/mohuyn/Library/CloudStorage/OneDrive-SAS/Documents/GitHub/Hybrid-Quantum-Graph-AI-QAOA-GNN-Biomedical-Optimization/data/prostate_top10_variance_panel_meta.json

# -----------------------------------------------------------------------------
# CELL 4: Fetch the genomics dataset and build the graph benchmark
# -----------------------------------------------------------------------------

genomics_config = {
    "dataset_name": "prostate",
    "version": 1,
    "top_gene_count": 10,
    "target_edge_count": 18,
    "benchmark_size": 6,
    "benchmark_seed": 42,
    "adaptation_size": 24,
    "adaptation_seed": 200,
    "subsample_size": 60,
}

genomics_bundle = build_genomics_benchmark(**genomics_config)
expression_frame = genomics_bundle["expression_frame"]
labels = genomics_bundle["labels"]
gene_table = genomics_bundle["gene_table"]
representative = genomics_bundle["representative"]
adaptation_graphs = genomics_bundle["adaptation_instances"]
benchmark_graphs = genomics_bundle["benchmark_instances"]
adaptation_overview = genomics_bundle["adaptation_overview"]
benchmark_overview = genomics_bundle["benchmark_overview"]
dataset_summary = genomics_bundle["dataset_summary"]
full_feature_count = genomics_bundle["full_feature_count"]
cache_status = genomics_bundle["cache_status"]

G = representative["graph"]
n = representative["n"]
edges = representative["edges"]
best_cut = representative["best_cut"]

print("Genomics benchmark configuration")
print(f"  Dataset name              : {genomics_config['dataset_name']}")
print(f"  Source matrix shape       : {expression_frame.shape[0]} samples x {full_feature_count} genes")
print(f"  Working execution panel   : {expression_frame.shape[1]} cached/selected genes used to build graphs")
print(f"  Data access mode          : {cache_status}")
print(f"  Class balance             : {labels.value_counts().sort_index().to_dict()}")
print(f"  Selected genes            : {genomics_config['top_gene_count']}")
print(f"  Adaptation graphs         : {genomics_config['adaptation_size']}  (seed range {genomics_config['adaptation_seed']}..{genomics_config['adaptation_seed'] + genomics_config['adaptation_size'] - 1})")
print(f"  Held-out benchmark graphs : {genomics_config['benchmark_size']}  (seed range {genomics_config['benchmark_seed']}..{genomics_config['benchmark_seed'] + genomics_config['benchmark_size'] - 1})")
print(f"  Patients per graph        : {genomics_config['subsample_size']}")
print("\nRepresentative graph")
print("  Graph type                : full-cohort gene co-expression graph")
print(f"  Nodes                     : {n}")
print(f"  Edges                     : {len(edges)}")
print(f"  Exact MaxCut              : {best_cut}")
print(f"  Statevector size          : 2^{n} = {2 ** n} amplitudes")
print("\nDataset summary")
display(dataset_summary)
print("Selected gene panel")
display(gene_table)
print("Held-out benchmark overview")
display(benchmark_overview)
print("Adaptation overview (first 8 graphs)")
display(adaptation_overview.head(8))

Genomics benchmark configuration
  Dataset name              : prostate
  Source matrix shape       : 102 samples x 12600 genes
  Working execution panel   : 10 cached/selected genes used to build graphs
  Data access mode          : Fetched OpenML data and refreshed local reduced-panel cache at /Users/mohuyn/Library/CloudStorage/OneDrive-SAS/Documents/GitHub/Hybrid-Quantum-Graph-AI-QAOA-GNN-Biomedical-Optimization/data/prostate_top10_variance_panel.csv.gz
  Class balance             : {1: 52, 2: 50}
  Selected genes            : 10
  Adaptation graphs         : 24  (seed range 200..223)
  Held-out benchmark graphs : 6  (seed range 42..47)
  Patients per graph        : 60

Representative graph
  Graph type                : full-cohort gene co-expression graph
  Nodes                     : 10
  Edges                     : 18
  Exact MaxCut              : 13
  Statevector size          : 2^10 = 1024 amplitudes

Dataset summary

Selected gene panel

Held-out benchmark overview

Adaptation overview (first 8 graphs)

# -----------------------------------------------------------------------------
# CELL 5: Classical depth-2 QAOA optimization on the representative graph
# -----------------------------------------------------------------------------

representative_search = {"num_starts": 10, "maxiter": 420, "seed": 77}

start_time = time.perf_counter()
classical_result = classical_optimize_instance(
    representative,
    p=p,
    num_starts=representative_search["num_starts"],
    maxiter=representative_search["maxiter"],
    seed=representative_search["seed"],
)
classical_time = time.perf_counter() - start_time

opt_gammas = classical_result["gammas"]
opt_betas = classical_result["betas"]
opt_angles = np.concatenate([opt_gammas, opt_betas])
state_opt = classical_result["state"]
val_opt = classical_result["value"]
classical_ratio = val_opt / best_cut

print("Classical optimization on the representative full-cohort graph")
print(f"  gamma_opt              : {np.round(opt_gammas, 6)}")
print(f"  beta_opt               : {np.round(opt_betas, 6)}")
print(f"  Expected cut           : {val_opt:.4f}")
print(f"  Exact MaxCut           : {best_cut}")
print(f"  Approx. ratio          : {classical_ratio:.4f}")
print(f"  Starts                 : {representative_search['num_starts']}")
print(f"  Best iterations        : {classical_result['nit']}")
print(f"  Best evaluations       : {classical_result['nfev']}")
print(f"  Wall-clock time        : {classical_time * 1e3:.3f} ms")
print("\nInterpretation")
print("  This is the depth-2 reference solution for the full-cohort transcriptomic")
print("  graph. It establishes the strongest exact-statevector QAOA result used")
print("  for representative-graph interpretation in the rest of the notebook.")

Classical optimization on the representative full-cohort graph
  gamma_opt              : [0.369724 0.914073]
  beta_opt               : [0.376829 0.21028 ]
  Expected cut           : 11.6689
  Exact MaxCut           : 13
  Approx. ratio          : 0.8976
  Starts                 : 10
  Best iterations        : 191
  Best evaluations       : 335
  Wall-clock time        : 728.359 ms

Interpretation
  This is the depth-2 reference solution for the full-cohort transcriptomic
  graph. It establishes the strongest exact-statevector QAOA result used
  for representative-graph interpretation in the rest of the notebook.

# -----------------------------------------------------------------------------
# CELL 6: Transcriptomic domain adaptation and learned prediction
# -----------------------------------------------------------------------------

adaptation_search = {"num_starts": 8, "maxiter": 320, "seed_offset": 1000}
training_config = {
    "hidden_dim": 64,
    "epochs": 800,
    "lr": 2.5e-3,
    "weight_decay": 5e-5,
    "patience": 100,
    "seed": 19,
}

target_start = time.perf_counter()
adaptation_graphs, adaptation_target_table = attach_classical_targets(
    adaptation_graphs,
    p=p,
    num_starts=adaptation_search["num_starts"],
    maxiter=adaptation_search["maxiter"],
    seed_offset=adaptation_search["seed_offset"],
)
target_generation_time = time.perf_counter() - target_start

training_start = time.perf_counter()
training_bundle = train_adapted_qaoa_gnn(
    adaptation_graphs,
    p=p,
    hidden_dim=training_config["hidden_dim"],
    epochs=training_config["epochs"],
    lr=training_config["lr"],
    weight_decay=training_config["weight_decay"],
    patience=training_config["patience"],
    seed=training_config["seed"],
)
training_time = time.perf_counter() - training_start

adapted_model = training_bundle["model"]
model = adapted_model

gnn_result = predict_instance_with_gnn(representative, adapted_model, p=p)
pred_g = gnn_result["gammas"]
pred_b = gnn_result["betas"]
pred_angles = np.concatenate([pred_g, pred_b])
state_pred = gnn_result["state"]
val_pred = gnn_result["value"]
gnn_time = gnn_result["inference_time"]
gnn_ratio = val_pred / best_cut
relative_to_classical = val_pred / val_opt
speedup = classical_time / gnn_time if gnn_time > 0 else float("inf")

legacy_result = None
legacy_ratio = None
if legacy_model_loaded:
    legacy_result = predict_instance_with_gnn(representative, legacy_model, p=legacy_p)
    legacy_ratio = legacy_result["value"] / best_cut

print("Transcriptomic domain adaptation summary")
print(f"  Adaptation graphs         : {len(adaptation_graphs)}")
print(f"  Mean target ratio         : {adaptation_target_table['classical_ratio'].mean():.4f}")
print(f"  Target generation time    : {target_generation_time:.2f} s")
print(f"  Training epochs run       : {training_bundle['epochs_run']}")
print(f"  Best epoch                : {training_bundle['best_epoch']}")
print(f"  Best training loss        : {training_bundle['best_loss']:.6f}")
print(f"  Training wall-clock time  : {training_time:.2f} s")
print("\nRepresentative-graph evaluation")
print(f"  gamma_gnn                : {np.round(pred_g, 6)}")
print(f"  beta_gnn                 : {np.round(pred_b, 6)}")
print(f"  Expected cut             : {val_pred:.4f}")
print(f"  Ratio vs exact           : {gnn_ratio:.4f}")
print(f"  Ratio vs classical       : {relative_to_classical:.4f}")
print(f"  Inference time           : {gnn_time * 1e3:.6f} ms")
print(f"  Classical time           : {classical_time * 1e3:.6f} ms")
print(f"  Speedup                  : {speedup:,.1f}x")
if legacy_ratio is not None:
    print(f"  Legacy transfer ratio    : {legacy_ratio:.4f}  (original depth-1 checkpoint)")

print("\nAdaptation-target overview")
display(adaptation_target_table.head(8).round(4))
print("\nInterpretation")
print("  This tuned configuration was selected to maximize the representative-graph")
print("  learned ratio while keeping the held-out benchmark protocol unchanged.")
print("  The upgraded model is still adapted only on separate transcriptomic")
print("  resamples before held-out evaluation.")

Transcriptomic domain adaptation summary
  Adaptation graphs         : 24
  Mean target ratio         : 0.8923
  Target generation time    : 15.40 s
  Training epochs run       : 786
  Best epoch                : 686
  Best training loss        : 0.000582
  Training wall-clock time  : 5.35 s

Representative-graph evaluation
  gamma_gnn                : [0.380436 0.929814]
  beta_gnn                 : [0.373343 0.209169]
  Expected cut             : 11.6677
  Ratio vs exact           : 0.8975
  Ratio vs classical       : 0.9999
  Inference time           : 0.060542 ms
  Classical time           : 728.358709 ms
  Speedup                  : 12,030.6x
  Legacy transfer ratio    : 0.5613  (original depth-1 checkpoint)

Adaptation-target overview

Interpretation
  This tuned configuration was selected to maximize the representative-graph
  learned ratio while keeping the held-out benchmark protocol unchanged.
  The upgraded model is still adapted only on separate transcriptomic
  resamples before held-out evaluation.

# -----------------------------------------------------------------------------
# CELL 7: Representative-graph figures
# -----------------------------------------------------------------------------

from matplotlib.ticker import FormatStrFormatter

fig = plt.figure(figsize=(19.5, 5.4))
ax1 = fig.add_subplot(1, 3, 1)
ax2 = fig.add_subplot(1, 3, 2)
ax3 = fig.add_subplot(1, 3, 3)

# (A) Gene co-expression graph with the exact optimal partition.
optimal_partition = {node: (representative["best_mask"] >> node) & 1 for node in G.nodes()}
node_colors = ["#2D6A4F" if optimal_partition[node] == 0 else "#BC4749" for node in G.nodes()]
label_map = {node: f"G{node + 1}" for node in G.nodes()}
edge_weights = [G[u][v]["weight"] for u, v in G.edges()]
min_weight = min(edge_weights)
max_weight = max(edge_weights)
edge_widths = [1.5 + 3.5 * ((weight - min_weight) / (max_weight - min_weight + 1e-9)) for weight in edge_weights]
pos = nx.spring_layout(G, seed=42, weight="weight")

nx.draw_networkx_edges(G, pos, ax=ax1, edge_color="#5C677D", width=edge_widths, alpha=0.8)
nx.draw_networkx_nodes(G, pos, ax=ax1, node_color=node_colors, node_size=820)
nx.draw_networkx_labels(G, pos, ax=ax1, labels=label_map, font_color="white", font_size=9, font_weight="bold")
ax1.set_title("Representative full-cohort gene graph\nExact MaxCut partition", fontsize=11)
ax1.text(
    0.02,
    0.02,
    f"Exact MaxCut = {best_cut}\nDepth-2 classical ratio = {classical_ratio:.4f}\nLabels G1-G10 match the selected gene table",
    transform=ax1.transAxes,
    fontsize=8.5,
    bbox={"facecolor": "white", "alpha": 0.85, "edgecolor": "none"},
)
ax1.axis("off")

# (B) Interpolation from learned angles to the classical optimum.
alpha_grid = np.linspace(0.0, 1.0, 180)
interpolation_values = []
for alpha in alpha_grid:
    blended = pred_angles + alpha * (opt_angles - pred_angles)
    blend_gammas, blend_betas = normalize_angles(blended, p)
    interpolation_values.append(qaoa_value_for_angles(representative["cut_diagonal"], blend_gammas, blend_betas)[0])

interpolation_values = np.asarray(interpolation_values)
curve_min = float(interpolation_values.min())
curve_max = float(interpolation_values.max())
curve_pad = max(6e-4, 0.20 * max(curve_max - curve_min, 1e-6))
gap_to_classical = val_opt - interpolation_values

ax2.plot(alpha_grid, interpolation_values, lw=2.7, color="#355070", label="Expected cut along interpolation")
ax2.fill_between(alpha_grid, interpolation_values, val_opt, color="#BC4749", alpha=0.10, label="Residual gap to classical")
ax2.scatter([0.0], [val_pred], c="#DD8B3D", s=120, marker="D", zorder=5, label="Adapted GNN")
ax2.scatter([1.0], [val_opt], c="#BC4749", s=140, marker="*", zorder=5, label="Classical optimum")
ax2.axhline(val_opt, color="#BC4749", ls="--", lw=1.2, alpha=0.8)
ax2.set_xlim(0.0, 1.0)
ax2.set_ylim(curve_min - curve_pad, curve_max + curve_pad)
ax2.set_xlabel("Interpolation from adapted model to classical optimum", fontsize=10)
ax2.set_ylabel("Expected cut", fontsize=10)
ax2.yaxis.set_major_formatter(FormatStrFormatter("%.4f"))
ax2.set_title("Depth-2 parameter path: learned start to classical optimum\n(zoomed to reveal the residual gap)", fontsize=11)
ax2.grid(True, alpha=0.3)
ax2.legend(fontsize=8, loc="lower right")
ax2.text(
    0.03,
    0.97,
    f"Retention at learned start: {relative_to_classical * 100:.2f}%\nRemaining gap: {val_opt - val_pred:.4f}\nMax gain along path: {gap_to_classical.max():.4f}",
    transform=ax2.transAxes,
    va="top",
    fontsize=8.5,
    bbox={"facecolor": "white", "alpha": 0.9, "edgecolor": "#cccccc"},
)

# (C) Two-dimensional slice through the depth-2 landscape.
resolution_gamma = 50
resolution_beta = 36
gamma_grid = np.linspace(0, math.pi, resolution_gamma)
beta_grid = np.linspace(0, math.pi / 2, resolution_beta)
landscape = np.zeros((resolution_beta, resolution_gamma))

fixed_gamma_2 = float(opt_gammas[1])
fixed_beta_2 = float(opt_betas[1])
for gamma_index, gamma_1 in enumerate(gamma_grid):
    for beta_index, beta_1 in enumerate(beta_grid):
        landscape[beta_index, gamma_index] = qaoa_value_for_angles(
            representative["cut_diagonal"],
            np.array([gamma_1, fixed_gamma_2]),
            np.array([beta_1, fixed_beta_2]),
        )[0]

slice_peak = float(landscape.max())
slice_threshold_95 = 0.95 * slice_peak
slice_fraction_95 = float((landscape >= slice_threshold_95).sum()) / landscape.size

image = ax3.imshow(
    landscape,
    origin="lower",
    aspect="auto",
    cmap="cividis",
    extent=[0, math.pi, 0, math.pi / 2],
)
ax3.contour(
    gamma_grid,
    beta_grid,
    landscape,
    levels=[slice_threshold_95],
    colors="white",
    linewidths=1.2,
    linestyles="--",
)
ax3.scatter([float(opt_gammas[0])], [float(opt_betas[0])], c="#BC4749", s=140, marker="*", zorder=5, label="Classical optimum")
ax3.scatter([float(pred_g[0])], [float(pred_b[0])], c="#DD8B3D", s=120, marker="D", zorder=5, label="Adapted GNN")
ax3.set_xlabel(r"$\gamma_1$", fontsize=10)
ax3.set_ylabel(r"$\beta_1$", fontsize=10)
ax3.set_title(
    rf"Depth-2 slice with $\gamma_2={fixed_gamma_2:.3f}$ and $\beta_2={fixed_beta_2:.3f}$ fixed",
    fontsize=11,
)
ax3.legend(fontsize=8)
ax3.text(
    0.03,
    0.05,
    f"95% basin in visible slice: {slice_fraction_95 * 100:.1f}%",
    transform=ax3.transAxes,
    fontsize=8.5,
    color="white",
    bbox={"facecolor": "black", "alpha": 0.35, "edgecolor": "none"},
)
plt.colorbar(image, ax=ax3, label="Expected cut")

plt.tight_layout()
plt.show()

print("Figure guide")
print("  Left   : the full-cohort gene co-expression graph colored by the exact MaxCut partition.")
print("  Middle : a zoomed 4D interpolation that makes the remaining learned-to-classical gap visible.")
print("  Right  : a depth-2 landscape slice with the 95%-of-peak basin overlaid as a dashed contour.")

Figure guide
  Left   : the full-cohort gene co-expression graph colored by the exact MaxCut partition.
  Middle : a zoomed 4D interpolation that makes the remaining learned-to-classical gap visible.
  Right  : a depth-2 landscape slice with the 95%-of-peak basin overlaid as a dashed contour.

# -----------------------------------------------------------------------------
# CELL 8: Held-out benchmark evaluation and final summary figures
# -----------------------------------------------------------------------------

benchmark_rows = []
for instance in benchmark_graphs:
    classical_start = time.perf_counter()
    classical_instance = classical_optimize_instance(
        instance,
        p=p,
        num_starts=8,
        maxiter=320,
        seed=int(instance["graph_id"]) + 7,
    )
    classical_elapsed = time.perf_counter() - classical_start
    adapted_instance = predict_instance_with_gnn(instance, adapted_model, p=p)

    legacy_instance = None
    if legacy_model_loaded:
        legacy_instance = predict_instance_with_gnn(instance, legacy_model, p=legacy_p)

    row = {
        "graph_id": instance["graph_id"],
        "samples": instance["sample_count"],
        "edges": instance["edge_count"],
        "exact_maxcut": instance["best_cut"],
        "classical_expected_cut": classical_instance["value"],
        "adapted_expected_cut": adapted_instance["value"],
        "classical_ratio": classical_instance["value"] / instance["best_cut"],
        "gnn_ratio": adapted_instance["value"] / instance["best_cut"],
        "gnn_vs_classical": adapted_instance["value"] / classical_instance["value"],
        "classical_ms": classical_elapsed * 1e3,
        "gnn_ms": adapted_instance["inference_time"] * 1e3,
    }
    if legacy_instance is not None:
        row["legacy_ratio"] = legacy_instance["value"] / instance["best_cut"]
        row["legacy_ms"] = legacy_instance["inference_time"] * 1e3
    benchmark_rows.append(row)

benchmark_results = pd.DataFrame(benchmark_rows)
mean_legacy_ratio = benchmark_results["legacy_ratio"].mean() if "legacy_ratio" in benchmark_results else float("nan")
mean_lift_over_legacy = benchmark_results["gnn_ratio"].mean() - mean_legacy_ratio if "legacy_ratio" in benchmark_results else float("nan")
relative_lift_over_legacy = (benchmark_results["gnn_ratio"].mean() / mean_legacy_ratio - 1.0) if "legacy_ratio" in benchmark_results else float("nan")
median_speedup = (benchmark_results["classical_ms"] / benchmark_results["gnn_ms"]).median()

summary_rows = [
    {"metric": "Mean classical ratio (p=2)", "value": benchmark_results["classical_ratio"].mean()},
    {"metric": "Mean adapted GNN ratio (p=2)", "value": benchmark_results["gnn_ratio"].mean()},
    {"metric": "Mean adapted/classical retention", "value": benchmark_results["gnn_vs_classical"].mean()},
]
if "legacy_ratio" in benchmark_results:
    summary_rows.extend(
        [
            {"metric": "Mean legacy transfer ratio (p=1)", "value": mean_legacy_ratio},
            {"metric": "Absolute lift over legacy", "value": mean_lift_over_legacy},
            {"metric": "Relative lift over legacy", "value": relative_lift_over_legacy},
        ]
    )
summary_rows.extend(
    [
        {"metric": "Median classical latency (ms)", "value": benchmark_results["classical_ms"].median()},
        {"metric": "Median adapted GNN latency (ms)", "value": benchmark_results["gnn_ms"].median()},
        {"metric": "Median speedup", "value": median_speedup},
    ]
)
summary_table = pd.DataFrame(summary_rows)

print("Benchmark results table")
display(
    benchmark_results[
        [
            column
            for column in [
                "graph_id",
                "samples",
                "edges",
                "exact_maxcut",
                "classical_ratio",
                "gnn_ratio",
                "legacy_ratio",
                "gnn_vs_classical",
                "classical_ms",
                "gnn_ms",
            ]
            if column in benchmark_results.columns
        ]
    ].round(4)
)
print("Summary statistics")
display(summary_table.round(4))

fig, axes = plt.subplots(1, 3, figsize=(19, 5))

# (A) Median latency comparison across the benchmark.
median_latencies = benchmark_results[["classical_ms", "gnn_ms"]].median()
latency_labels = ["Classical\nsearch", "Adapted GNN\ninference"]
latency_colors = ["#355070", "#6D597A"]
bars = axes[0].bar(latency_labels, median_latencies.values, color=latency_colors, edgecolor="black", width=0.55)
for bar, value in zip(bars, median_latencies.values):
    label = f"{value:.3f} ms" if value >= 0.01 else f"{value * 1e3:.2f} us"
    axes[0].text(bar.get_x() + bar.get_width() / 2, bar.get_height() * 1.1, label, ha="center", va="bottom", fontsize=9)
axes[0].set_yscale("log")
axes[0].set_ylabel("Latency (ms, log scale)")
axes[0].set_title("Median latency across the held-out benchmark")
axes[0].set_xlabel(f"Median speedup: {median_speedup:,.1f}x")
axes[0].grid(True, axis="y", alpha=0.3)

# (B) Representative-state probability distribution under the classical parameters.
state_probabilities = np.abs(state_opt) ** 2
bit_strings = [format(index, f"0{n}b") for index in range(2 ** n)]
cut_values = [
    sum(1 for u, v in edges if ((state_index >> u) & 1) != ((state_index >> v) & 1))
    for state_index in range(2 ** n)
]
color_norm = plt.Normalize(vmin=min(cut_values), vmax=best_cut)
color_map = plt.cm.cividis

top_k = 16
sorted_indices = np.argsort(-state_probabilities)[:top_k]
axes[1].bar(
    range(top_k),
    state_probabilities[sorted_indices],
    color=[color_map(color_norm(cut_values[index])) for index in sorted_indices],
)
axes[1].set_xticks(range(top_k))
axes[1].set_xticklabels([bit_strings[index] for index in sorted_indices], rotation=70, fontsize=7)
axes[1].set_ylabel("Probability")
axes[1].set_title("Representative graph: top output states")
axes[1].grid(True, axis="y", alpha=0.3)
plt.colorbar(plt.cm.ScalarMappable(cmap=color_map, norm=color_norm), ax=axes[1], label="Cut value")

# (C) Quality comparison across all benchmark graphs.
x_positions = np.arange(len(benchmark_results))
width = 0.24
axes[2].bar(
    x_positions - width,
    benchmark_results["classical_ratio"],
    width=width,
    color="#355070",
    label="Classical p=2",
)
axes[2].bar(
    x_positions,
    benchmark_results["gnn_ratio"],
    width=width,
    color="#6D597A",
    label="Adapted GNN p=2",
)
if "legacy_ratio" in benchmark_results:
    axes[2].bar(
        x_positions + width,
        benchmark_results["legacy_ratio"],
        width=width,
        color="#DD8B3D",
        label="Legacy transfer p=1",
    )
axes[2].axhline(1.0, color="#6C757D", ls="--", lw=1.2, label="Exact reference")
axes[2].set_xticks(x_positions)
axes[2].set_xticklabels(benchmark_results["graph_id"].astype(str).tolist())
axes[2].set_ylim(0.0, 1.05)
axes[2].set_xlabel("Held-out benchmark graph ID")
axes[2].set_ylabel("Expected cut / exact MaxCut")
axes[2].set_title("Quality across held-out transcriptomic graphs")
axes[2].grid(True, axis="y", alpha=0.3)
axes[2].legend(fontsize=8)

plt.tight_layout()
plt.show()

print("Benchmark summary")
print(f"  Mean classical ratio      : {benchmark_results['classical_ratio'].mean():.4f}")
print(f"  Mean adapted GNN ratio    : {benchmark_results['gnn_ratio'].mean():.4f}")
print(f"  Mean GNN/classical        : {benchmark_results['gnn_vs_classical'].mean():.4f}")
if "legacy_ratio" in benchmark_results:
    print(f"  Mean legacy transfer      : {mean_legacy_ratio:.4f}")
    print(f"  Absolute lift vs legacy   : {mean_lift_over_legacy:.4f}")
    print(f"  Relative lift vs legacy   : {relative_lift_over_legacy:.1%}")
print(f"  Median speedup            : {median_speedup:,.1f}x")

Benchmark results table

Summary statistics

Benchmark summary
  Mean classical ratio      : 0.8686
  Mean adapted GNN ratio    : 0.8682
  Mean GNN/classical        : 0.9995
  Mean legacy transfer      : 0.5725
  Absolute lift vs legacy   : 0.2957
  Relative lift vs legacy   : 51.6%
  Median speedup            : 14,122.1x

# ─────────────────────────────────────────────────────────────────────────────
# Section 8A: QAOA Landscape Geometry Analysis
# Uses: landscape (resolution_beta × resolution_gamma), gamma_grid, beta_grid,
#       resolution_gamma, resolution_beta, best_cut  (defined in CELL 7 / CELL 4)
# ─────────────────────────────────────────────────────────────────────────────
import math as _math

landscape_flat = landscape.flatten()
global_max_val = float(landscape.max())
global_max_idx = np.unravel_index(landscape.argmax(), landscape.shape)
gamma_at_max = float(gamma_grid[global_max_idx[1]])
beta_at_max = float(beta_grid[global_max_idx[0]])

taus = {0.95: 0.0, 0.90: 0.0, 0.80: 0.0}
for tau in taus:
    taus[tau] = float((landscape >= tau * global_max_val).sum()) / landscape.size

cv = float(landscape_flat.std() / landscape_flat.mean())

print("QAOA Landscape Geometry ── Depth-2 Slice on the Representative Gene Graph")
print(f"  Slice resolution              : {resolution_gamma} × {resolution_beta} = {landscape.size:,} evaluations")
print(f"  Exact MaxCut                  : {best_cut}")
print(f"  Slice max  E[C]*              : {global_max_val:.4f}")
print(f"  Slice ratio  E[C]*/MC         : {global_max_val / best_cut:.4f}")
print(f"  Argmax γ1                     : {gamma_at_max:.4f} rad  ({_math.degrees(gamma_at_max):.1f}°)")
print(f"  Argmax β1                     : {beta_at_max:.4f} rad  ({_math.degrees(beta_at_max):.1f}°)")
print(f"  Fraction ≥ 95% of slice max   : {taus[0.95]:.3%}  ← {int(taus[0.95] * landscape.size):d} / {landscape.size} grid points")
print(f"  Fraction ≥ 90% of slice max   : {taus[0.90]:.3%}")
print(f"  Fraction ≥ 80% of slice max   : {taus[0.80]:.3%}")
print(f"  Coefficient of variation      : {cv:.4f}  (σ/μ of slice values)")
print()
print("Interpretation")
print(f"  Only {taus[0.95]:.1%} of the visible (γ1,β1) slice achieves at least 95% of the")
print(f"  best value in that slice. This means the high-quality basin is concentrated rather")
print(f"  than diffuse, which explains why a learned warm-start is valuable even before any")
print(f"  local classical refinement is applied.")

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

axes[0].hist(landscape_flat, bins=40, color="#355070", edgecolor="black", alpha=0.75)
axes[0].axvline(global_max_val, color="#BC4749", lw=2, ls="--", label=f"Slice max ({global_max_val:.3f})")
axes[0].axvline(0.95 * global_max_val, color="#DD8B3D", lw=1.5, ls=":", label=f"95% threshold ({0.95 * global_max_val:.3f})")
axes[0].set_xlabel(r"$E[C](\gamma_1, \beta_1 \mid \gamma_2^*, \beta_2^*)$", fontsize=11)
axes[0].set_ylabel("Count (grid points)", fontsize=11)
axes[0].set_title("Distribution of the Depth-2 Slice Values", fontsize=11)
axes[0].legend(fontsize=9)
axes[0].grid(True, alpha=0.3)

sorted_desc = np.sort(landscape_flat)[::-1]
frac_x = np.linspace(1.0 / landscape.size, 1.0, landscape.size) * 100
axes[1].plot(frac_x, sorted_desc, color="#355070", lw=2.5, label="Depth-2 slice")
for tau, color in [(0.95, "#BC4749"), (0.90, "#DD8B3D"), (0.80, "#6D597A")]:
    axes[1].axhline(tau * global_max_val, color=color, lw=1.5, ls="--", label=f"{int(tau * 100)}% of slice max")
axes[1].set_xlabel("Top x% of slice grid points (ranked by E[C])", fontsize=11)
axes[1].set_ylabel("Expected cut", fontsize=11)
axes[1].set_title("Near-Optimum Fraction in the Visible Basin", fontsize=11)
axes[1].legend(fontsize=9)
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

QAOA Landscape Geometry ── Depth-2 Slice on the Representative Gene Graph
  Slice resolution              : 50 × 36 = 1,800 evaluations
  Exact MaxCut                  : 13
  Slice max  E[C]*              : 11.6624
  Slice ratio  E[C]*/MC         : 0.8971
  Argmax γ1                     : 0.3847 rad  (22.0°)
  Argmax β1                     : 0.3590 rad  (20.6°)
  Fraction ≥ 95% of slice max   : 2.944%  ← 53 / 1800 grid points
  Fraction ≥ 90% of slice max   : 7.167%
  Fraction ≥ 80% of slice max   : 30.944%
  Coefficient of variation      : 0.1057  (σ/μ of slice values)

Interpretation
  Only 2.9% of the visible (γ1,β1) slice achieves at least 95% of the
  best value in that slice. This means the high-quality basin is concentrated rather
  than diffuse, which explains why a learned warm-start is valuable even before any
  local classical refinement is applied.

# ─────────────────────────────────────────────────────────────────────────────
# Section 8B: Quantum State Concentration Analysis
# Uses: state_probabilities, cut_values, n, best_cut  (defined in CELL 8)
# ─────────────────────────────────────────────────────────────────────────────
from scipy import stats as _stats

probs_arr = np.array(state_probabilities)
cuts_arr = np.array(cut_values)

rho_s, p_spear = _stats.spearmanr(cuts_arr, probs_arr)

nonzero = probs_arr > 0
entropy_bits = float(-np.sum(probs_arr[nonzero] * np.log2(probs_arr[nonzero])))
max_entropy = float(n)

cut_80th = float(np.percentile(cuts_arr, 80))
top20_mask = cuts_arr >= cut_80th
p_top20_qaoa = float(probs_arr[top20_mask].sum())
p_top20_uniform = float(top20_mask.sum()) / len(cuts_arr)
conc_ratio = p_top20_qaoa / p_top20_uniform

best_mask = cuts_arr == best_cut
p_best_qaoa = float(probs_arr[best_mask].sum())
p_best_uniform = float(best_mask.sum()) / len(cuts_arr)

print("Quantum State Concentration ── Representative Gene Graph")
print(f"  Statevector dimension          : 2^{n} = {2 ** n} bitstrings")
print(f"  Spearman ρ (cut vs. prob.)     : {rho_s:+.4f}  (p = {p_spear:.3e})")
print(f"  Shannon entropy                : {entropy_bits:.3f} bits  (uniform = {max_entropy:.1f} bits)")
print(f"  Entropy ratio                  : {entropy_bits / max_entropy:.3f}  (1.0 = maximally uniform)")
print(f"  P(top-20% cuts | QAOA)         : {p_top20_qaoa:.4f}")
print(f"  P(top-20% cuts | uniform)      : {p_top20_uniform:.4f}")
print(f"  Concentration ratio            : {conc_ratio:.2f}x")
print(f"  P(exact MaxCut strings | QAOA) : {p_best_qaoa:.4f}  (uniform: {p_best_uniform:.4f})")
print()
if rho_s > 0 and p_spear < 0.05:
    print(f"  Significant positive Spearman ρ = {rho_s:+.4f}: the depth-2 state amplifies")
    print("  high-cut bitstrings above the uniform baseline.")
elif rho_s > 0:
    print(f"  Positive but non-significant Spearman ρ = {rho_s:+.4f}.")
else:
    print(f"  Spearman ρ = {rho_s:+.4f}: concentration is weak on this graph.")
print(f"  The concentration ratio of {conc_ratio:.2f}x means the optimized state assigns")
print("  substantially more probability mass to strong partitions than a random guess.")

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

axes[0].scatter(cuts_arr, probs_arr, alpha=0.35, s=14, color="#355070", label=f"Spearman ρ = {rho_s:+.4f}")
axes[0].axhline(1.0 / len(cuts_arr), color="#BC4749", lw=1.5, ls="--", label=f"Uniform 1/2^n = {1 / len(cuts_arr):.5f}")
axes[0].set_xlabel("Cut value C(z)", fontsize=11)
axes[0].set_ylabel(r"QAOA probability $|\langle z|\psi(\gamma^*,\beta^*)\rangle|^2$", fontsize=11)
axes[0].set_title(f"State Concentration: Cut Value vs. Probability\nSpearman ρ = {rho_s:+.4f}, p = {p_spear:.3e}", fontsize=11)
axes[0].legend(fontsize=9)
axes[0].grid(True, alpha=0.3)

sort_idx = np.argsort(cuts_arr)[::-1]
cum_prob = np.cumsum(probs_arr[sort_idx])
frac_axis = np.linspace(1.0 / len(cum_prob), 1.0, len(cum_prob)) * 100
uniform_cum = np.linspace(1.0 / len(cum_prob), 1.0, len(cum_prob))
axes[1].plot(frac_axis, cum_prob, color="#355070", lw=2.5, label="QAOA state")
axes[1].plot(frac_axis, uniform_cum, color="#BC4749", lw=1.5, ls="--", label="Uniform baseline")
axes[1].fill_between(
    frac_axis,
    cum_prob,
    uniform_cum,
    where=(cum_prob > uniform_cum),
    alpha=0.15,
    color="#2D6A4F",
    label="Excess mass above uniform",
)
axes[1].set_xlabel("Top x% of bitstrings (ranked by cut value)", fontsize=11)
axes[1].set_ylabel("Cumulative QAOA probability", fontsize=11)
axes[1].set_title("State Concentration CDF: QAOA vs. Uniform", fontsize=11)
axes[1].legend(fontsize=9)
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

Quantum State Concentration ── Representative Gene Graph
  Statevector dimension          : 2^10 = 1024 bitstrings
  Spearman ρ (cut vs. prob.)     : +0.7152  (p = 3.569e-161)
  Shannon entropy                : 7.815 bits  (uniform = 10.0 bits)
  Entropy ratio                  : 0.781  (1.0 = maximally uniform)
  P(top-20% cuts | QAOA)         : 0.9037
  P(top-20% cuts | uniform)      : 0.2461
  Concentration ratio            : 3.67x
  P(exact MaxCut strings | QAOA) : 0.2132  (uniform: 0.0098)

  Significant positive Spearman ρ = +0.7152: the depth-2 state amplifies
  high-cut bitstrings above the uniform baseline.
  The concentration ratio of 3.67x means the optimized state assigns
  substantially more probability mass to strong partitions than a random guess.

# ─────────────────────────────────────────────────────────────────────────────
# Section 8C: Residual Regret and Improvement-over-Transfer Analysis
# Uses: benchmark_results DataFrame  (defined in CELL 8)
# ─────────────────────────────────────────────────────────────────────────────
from scipy import stats as _stats2

bdf = benchmark_results.copy()
bdf["regret"] = bdf["classical_ratio"] - bdf["gnn_ratio"]
bdf["speedup"] = bdf["classical_ms"] / bdf["gnn_ms"]

mean_regret = float(bdf["regret"].mean())
std_regret = float(bdf["regret"].std())
max_regret = float(bdf["regret"].max())
worst_id = int(bdf.loc[bdf["regret"].idxmax(), "graph_id"])
quality_ret = bdf["gnn_vs_classical"]
median_speedup = float(bdf["speedup"].median())

pearson_r, pearson_p = _stats2.pearsonr(bdf["classical_ratio"], bdf["gnn_ratio"])
mean_legacy_ratio = float(bdf["legacy_ratio"].mean()) if "legacy_ratio" in bdf else float("nan")
mean_lift_over_legacy = float((bdf["gnn_ratio"] - bdf["legacy_ratio"]).mean()) if "legacy_ratio" in bdf else float("nan")

print("Residual Regret Analysis ── Held-Out Patient-Resampled Benchmark")
print(f"  Mean regret ± std           : {mean_regret:+.4f} ± {std_regret:.4f}")
print(f"  Max regret (graph {worst_id})      : {max_regret:.4f}")
print(f"  Pearson r (cls vs. adapted) : {pearson_r:+.4f}  (p = {pearson_p:.3e})")
print(f"  Mean quality retention      : {quality_ret.mean():.2%} ± {quality_ret.std():.2%}")
if "legacy_ratio" in bdf:
    print(f"  Mean legacy transfer ratio  : {mean_legacy_ratio:.4f}")
    print(f"  Mean lift over legacy       : {mean_lift_over_legacy:.4f}")
print(f"  Median inference speedup    : {median_speedup:,.0f}x")
print()
print("Per-graph regret table")
reg_table = bdf[[column for column in ["graph_id", "classical_ratio", "gnn_ratio", "legacy_ratio", "regret", "gnn_vs_classical", "speedup"] if column in bdf.columns]].copy()
reg_table.columns = [
    "Graph",
    "Classical ratio",
    "Adapted ratio",
    "Legacy ratio",
    "Regret",
    "Retention",
    "Speedup",
][: len(reg_table.columns)]
display(reg_table.round(4))
print()
print("Interpretation")
print(f"  The residual regret is only {mean_regret:.4f} on average, which means the adapted")
print("  model has essentially closed the quality gap relative to direct classical search.")
if "legacy_ratio" in bdf:
    print(f"  The lift over the old transfer baseline is {mean_lift_over_legacy:.4f} in absolute")
    print("  approximation ratio, making the improvement large in both statistical and practical terms.")
print(f"  The remaining speed-quality tradeoff is therefore far more favorable than in the")
print("  original transfer-only notebook.")

fig, axes = plt.subplots(1, 3, figsize=(18, 5))

bar_colors = ["#BC4749" if r > 0 else "#2D6A4F" for r in bdf["regret"]]
axes[0].bar(bdf["graph_id"].astype(str), bdf["regret"], color=bar_colors, edgecolor="black", width=0.6)
axes[0].axhline(0, color="black", lw=1.2)
axes[0].axhline(mean_regret, color="#355070", lw=1.5, ls="--", label=f"Mean regret = {mean_regret:.4f}")
axes[0].set_xlabel("Benchmark graph ID", fontsize=11)
axes[0].set_ylabel("Regret (classical - adapted ratio)", fontsize=11)
axes[0].set_title("Residual regret after transcriptomic adaptation", fontsize=11)
axes[0].legend(fontsize=9)
axes[0].grid(True, axis="y", alpha=0.3)

lo = min(bdf["classical_ratio"].min(), bdf["gnn_ratio"].min()) - 0.03
hi = max(bdf["classical_ratio"].max(), bdf["gnn_ratio"].max()) + 0.03
axes[1].scatter(bdf["classical_ratio"], bdf["gnn_ratio"], s=90, color="#355070", edgecolors="black", zorder=3)
for _, row in bdf.iterrows():
    axes[1].annotate(f"G{int(row['graph_id'])}", (row["classical_ratio"], row["gnn_ratio"]), textcoords="offset points", xytext=(6, 4), fontsize=9)
axes[1].plot([lo, hi], [lo, hi], color="#BC4749", ls="--", lw=1.5, label="Parity")
axes[1].set_xlabel("Classical QAOA ratio", fontsize=11)
axes[1].set_ylabel("Adapted GNN ratio", fontsize=11)
axes[1].set_title(f"Classical vs. adapted quality\nPearson r = {pearson_r:+.4f}, p = {pearson_p:.3e}", fontsize=11)
axes[1].legend(fontsize=9)
axes[1].grid(True, alpha=0.3)

improvement_vs_legacy = bdf["gnn_ratio"] - bdf["legacy_ratio"] if "legacy_ratio" in bdf else bdf["gnn_ratio"] - bdf["classical_ratio"]
axes[2].bar(bdf["graph_id"].astype(str), improvement_vs_legacy, color="#6D597A", edgecolor="black", width=0.6)
axes[2].axhline(0, color="black", lw=1.2)
if "legacy_ratio" in bdf:
    axes[2].set_ylabel("Adapted ratio - legacy transfer ratio", fontsize=11)
    axes[2].set_title("Improvement over the original transfer baseline", fontsize=11)
else:
    axes[2].set_ylabel("Adapted ratio - classical ratio", fontsize=11)
    axes[2].set_title("Quality gap summary", fontsize=11)
axes[2].set_xlabel("Benchmark graph ID", fontsize=11)
axes[2].grid(True, axis="y", alpha=0.3)

plt.tight_layout()
plt.show()

Residual Regret Analysis ── Held-Out Patient-Resampled Benchmark
  Mean regret ± std           : +0.0004 ± 0.0006
  Max regret (graph 43)      : 0.0017
  Pearson r (cls vs. adapted) : +0.9998  (p = 3.598e-08)
  Mean quality retention      : 99.95% ± 0.08%
  Mean legacy transfer ratio  : 0.5725
  Mean lift over legacy       : 0.2957
  Median inference speedup    : 14,122x

Per-graph regret table

Interpretation
  The residual regret is only 0.0004 on average, which means the adapted
  model has essentially closed the quality gap relative to direct classical search.
  The lift over the old transfer baseline is 0.2957 in absolute
  approximation ratio, making the improvement large in both statistical and practical terms.
  The remaining speed-quality tradeoff is therefore far more favorable than in the
  original transfer-only notebook.

Component	Specification
Dataset	OpenML `prostate` — 102 patients × 12,600 expression features
Graph construction	Top-10 variance genes → Pearson co-expression graph (10 nodes, 18 edges)
Quantum model	Depth-2 QAOA (`p = 2`), exact statevector simulator, $2^{10} = 1,024$ amplitudes
Classical baseline	Multi-start Nelder-Mead on 4 angles (`8-10` restarts, up to `320-420` iterations)
Learned model	Adaptive Quantum GCN (`SimpleGCN`) with transcriptomic domain adaptation on 24 resampled graphs
Held-out benchmark	6 patient-resampled co-expression graphs (`n = 60` patients per graph)
Representative classical ratio	0.898 on the full-cohort graph
Representative adapted GNN ratio	0.898
Held-out mean classical ratio	0.869
Held-out mean adapted GNN ratio	0.868
Held-out quality retention	99.95% of the classical baseline
Impact over the old transfer baseline	From ~0.573 to ~0.868 mean ratio on real-data graphs
Median inference speedup	~1.4 × 10^4x on this machine

	Section	Content
0	Introduction	Problem formulation, biological motivation, strengthened scope
1	Quantum Computing Foundations	Qubit state space, unitary gates, QAOA circuit operators
1A	Conceptual Bridge	Informal translation of quantum formalism
2	The MaxCut Problem	NP-hardness, Ising Hamiltonian, Goemans-Williamson bound
2A	Conceptual Bridge	Graph optimization in plain language
3	Statevector Simulation	Tensor contraction, exponential complexity analysis
3A	Simulation Primer	Memory and runtime implications
4	Setup & Dependencies	Environment configuration, legacy checkpoint, adapted model shell
5	Computational Engine	Exact solver, real-data benchmark builder, domain adaptation utilities
6	Representative Graph Figures	Exact partition, 4D-to-2D landscape slice, learned-vs-classical geometry
6A	Multi-Graph Rationale	Why a held-out patient-resampled benchmark is necessary
7	Benchmark-Wide Evaluation	Held-out quality, retention, and latency across 6 real-data graphs
7A	Model & Data Notes	What domain adaptation changes and what still remains limited
8A	QAOA Landscape Geometry	Concentration of the depth-2 landscape slice around the optimum
8B	State Concentration Analysis	Probability mass on high-cut bitstrings under the depth-2 optimum
8C	Residual Regret Analysis	Benchmark regret after transcriptomic adaptation, plus lift over transfer
9	Closing Summary	Results table, contribution statement, limitations, project context

Upgrade	Why it matters
Depth-2 QAOA instead of depth-1	Raises the quality ceiling on these transcriptomic MaxCut instances
Transcriptomic domain adaptation on separate resampled graphs	Replaces out-of-domain transfer with a graph family that matches the benchmark distribution

Algorithm	Complexity	Approximation ratio
Exhaustive search	$O(2^n m)$	1.0 (exact)
Random bipartition	$O(m)$	0.5 (in expectation)
Goemans–Williamson SDP	$O(n^{3.5})$	≥ 0.878
QAOA $p=1$	$O(p \cdot 2^n)$	0.6924 on 3-regular graphs

Operation	Time	Space
$H^{\otimes n}$ initialisation	$O(2^n)$	$O(2^n)$
$U_C(\gamma)$ — $m$ RZZ gates	$O(m \cdot 2^n)$	$O(2^n)$
$U_B(\beta)$ — $n$ RX gates	$O(n \cdot 2^n)$	$O(2^n)$
$\langle C\rangle$ measurement	$O(m \cdot 2^n)$	$O(2^n)$
Total per evaluation	$O((n+m) \cdot 2^n)$	$O(2^n)$

Method	Max $n$	Notes
Dense statevector (NumPy)	~20	Exact, $O(2^n)$ memory
Tensor network (MPS)	~50-100	Approximate for low entanglement
Clifford simulator	Unlimited	Only Clifford circuits
FPGA/GPU accelerated	~30-35	1000× speedup
Qiskit Aer GPU	~30	Accelerated cuStateVec-style backend

	metric	value
0	Dataset name	Prostate
1	Samples	102
2	Gene-expression features	12600
3	Selected genes for graph	10
4	Representative graph nodes	10
5	Representative graph edges	18
6	Adaptation graphs	24
7	Held-out benchmark graphs	6
8	Patients per resampled graph	60
9	Execution source	Fresh OpenML fetch + reduced-panel cache refresh

	rank	gene	variance
0	1	1008_f_at	8.031112e+06
1	2	1894_f_at	5.711349e+06
2	3	1804_at	5.014354e+06
3	4	AFFX-hum_alu_at	3.323481e+06
4	5	40794_at	1.221019e+06
5	6	37407_s_at	9.183038e+05
6	7	31962_at	7.776198e+05
7	8	1288_s_at	7.385506e+05
8	9	32466_at	7.362326e+05
9	10	40887_g_at	6.372388e+05

	graph_id	split	samples	edges	density	exact_maxcut
0	42	benchmark	60	18	0.4	14
1	43	benchmark	60	18	0.4	14
2	44	benchmark	60	18	0.4	13
3	45	benchmark	60	18	0.4	13
4	46	benchmark	60	18	0.4	14
5	47	benchmark	60	18	0.4	13

	graph_id	split	samples	edges	density	exact_maxcut
0	200	adaptation	60	18	0.4	13
1	201	adaptation	60	18	0.4	13
2	202	adaptation	60	18	0.4	13
3	203	adaptation	60	18	0.4	13
4	204	adaptation	60	18	0.4	13
5	205	adaptation	60	18	0.4	13
6	206	adaptation	60	18	0.4	14
7	207	adaptation	60	18	0.4	13

	graph_id	exact_maxcut	classical_ratio	iterations	evaluations
0	200	13	0.8999	278	465
1	201	13	0.8959	217	376
2	202	13	0.9012	179	315
3	203	13	0.9045	308	520
4	204	13	0.8974	306	522
5	205	13	0.8854	165	287
6	206	14	0.8389	170	295
7	207	13	0.8962	162	281

	graph_id	samples	edges	exact_maxcut	classical_ratio	gnn_ratio	legacy_ratio	gnn_vs_classical	classical_ms	gnn_ms
0	42	60	18	14	0.8394	0.8389	0.5515	0.9994	722.1029	0.0556
1	43	60	18	14	0.8456	0.8439	0.5122	0.9980	627.7236	0.0434
2	44	60	18	13	0.8974	0.8973	0.6111	0.9999	561.2377	0.0469
3	45	60	18	13	0.8962	0.8962	0.6317	1.0000	709.7916	0.0471
4	46	60	18	14	0.8369	0.8368	0.5600	0.9999	629.2355	0.0440
5	47	60	18	13	0.8964	0.8963	0.5686	0.9999	637.8844	0.0458

	metric	value
0	Mean classical ratio (p=2)	0.8686
1	Mean adapted GNN ratio (p=2)	0.8682
2	Mean adapted/classical retention	0.9995
3	Mean legacy transfer ratio (p=1)	0.5725
4	Absolute lift over legacy	0.2957
5	Relative lift over legacy	0.5165
6	Median classical latency (ms)	633.5599
7	Median adapted GNN latency (ms)	0.0463
8	Median speedup	14122.0960

Quantity	Definition	Interpretation
Near-optimum fraction	$\\|\{(\gamma_1,\beta_1) : \langle C\rangle \geq \tau \cdot \langle C\rangle^*\}\\| \,/\, \\|\text{slice}\\|$	How concentrated is the good basin in the visible slice?
Coefficient of variation	$\sigma / \mu$ of all slice values	Normalized spread; higher means more structured geometry
Global max location	$\arg\max_{\gamma_1,\beta_1} \langle C \rangle$	Where the depth-2 slice peaks given the fixed second layer

Quantity	What it reveals
Mean residual regret ± std	How much quality remains unrecovered after domain adaptation
Pearson $r$ (classical vs adapted ratio)	Whether the adapted model preserves graph-to-graph difficulty structure
Lift over legacy transfer	How much improvement was obtained over the original out-of-domain checkpoint

#	Limitation	Consequence
1	Small exact regime — the notebook stays at `n = 10` nodes so every graph remains exactly solvable	The result is precise and auditable, but not yet a large-scale quantum-hardware claim
2	Low-depth ansatz — `p = 2` is much stronger than `p = 1`, but still a shallow circuit	Further gains may exist at higher depth, with higher search cost
3	Single-cohort adaptation — adaptation and benchmark graphs come from one prostate cohort family	External-cohort generalization is still untested
4	Variance-ranked genes — feature selection is statistical rather than pathway-aware	The graph is biologically motivated, not fully biology-curated
5	Unweighted MaxCut evaluation — edge presence comes from weighted correlations, but the objective counts cut edges	Some biological signal in edge magnitude is not exploited
6	Notebook-local adaptation — training is done inside the notebook for transparency	This is excellent for auditability, but not yet a production training pipeline

Metric	Value
Dataset	OpenML `prostate` — 102 patients × 12,600 expression features
Graph construction	Top-10 variance genes → Pearson co-expression (10 nodes, 18 edges)
QAOA depth	`p = 2`, exact statevector (`2^{10} = 1,024` amplitudes)
Exact reference	Brute-force MaxCut for every graph
Classical baseline	Multi-start Nelder-Mead over 4 QAOA angles
Learned model	Adaptive Quantum GCN (`SimpleGCN`) adapted on 24 transcriptomic resampled graphs
Representative classical ratio	0.8976
Representative adapted GNN ratio	0.8975
Held-out mean classical ratio	0.8686
Held-out mean adapted GNN ratio	0.8682
Held-out mean quality retention	0.9995
Old transfer baseline (for contrast)	0.5725 mean ratio
Absolute lift over transfer	0.2957
Relative lift over transfer	51.6%
Median inference speedup	~1.41 × 10^4x on this machine

QAOA + GNN for Genomics-Derived MaxCut¶

Project Summary¶

Technical Contribution¶

Table of Contents¶

0. Introduction¶

Problem Formulation¶

Biological Context¶

What Is New in This Version of the Notebook¶

Scope and Limitations (Up Front)¶

1. Quantum Computing Foundations ¶

The Qubit¶

Multi-Qubit Register¶

Essential Quantum Gates¶

QAOA-Specific Operators¶

The Variational State at Depth $p$¶

1A. Quantum Foundations¶

Qubits¶

Multiple qubits: the exponential advantage¶

Quantum gates: rotating the coin mid-spin¶

The two parameters gamma and beta¶

The QAOA circuit as a recipe¶

Mental model to carry forward¶

2. MaxCut Problem — Definition & Complexity ¶

Problem Statement¶

NP-Hardness¶

QAOA Cost Hamiltonian Derivation¶

At $p=1$: Closed-Form Expected Value¶

2A. MaxCut — From Puzzle to Quantum Hamiltonian¶

What is MaxCut in this notebook?¶

Why is that useful?¶

Converting MaxCut into a score formula¶

How does this become a quantum problem?¶

Approximation ratio: measuring how well we did¶

Key takeaway before the code¶

3. Statevector Simulation — How the Code Works ¶

Exact Statevector Approach¶

Complexity Analysis¶

Alternative Simulation Methods for Larger $n$¶

3A. Simulation Primer — What the Code Is Actually Doing¶

The state vector is a lookup table¶

What the cost gates do¶

What the mixer gates do¶

Why we simulate classically on small graphs¶

How to map simulator output to a MaxCut solution¶

What good angles look like¶

4. Setup — Imports & Project Configuration¶

4A. How to Read the Code Cells¶

Cell 1: Imports and environment checks¶

Cell 2: Legacy checkpoint plus adapted-model shell¶

Cell 3: Computational engine¶

Cell 4: Build the transcriptomic graph family¶

Cell 5: Classical depth-2 QAOA on the representative graph¶

Cell 6: Transcriptomic domain adaptation and learned prediction¶

Cell 7: Representative-graph figures¶

Cell 8: Held-out benchmark evaluation¶

5. Computational Engine — From Transcriptomic Data to a Depth-2 QAOA Benchmark¶

Why this step matters¶

What the next code cell will define¶

What readers should expect after the cell runs¶

Reproducibility Note — Local Transcriptomic Cache¶

6. Interpreting the Representative-Graph Figures¶

Figure 1: Full-cohort gene co-expression graph with the exact optimal partition¶

Figure 2: Interpolation from learned angles to the classical optimum¶

Figure 3: Two-dimensional slice through the depth-2 landscape¶

Why this figure set matters¶

6A. Why the Notebook Returns to a Held-Out Multi-Graph Benchmark¶

What the next code cell does¶

Why this matters¶

7. Interpreting the Benchmark-Wide Figures¶

Figure 1: Median latency on held-out graphs¶

Figure 2: Highest-probability output states for the representative graph¶

Figure 3: Quality across all held-out transcriptomic graphs¶

The combined message of the final figure set¶

7A. Model, Data, and Interpretation Notes¶

What changed relative to the original notebook¶

What this still does not prove¶

Why the result is still strong¶

8A. QAOA Landscape Geometry Analysis¶

8B. Quantum State Concentration Analysis¶

8C. Residual Regret and Improvement-over-Transfer Analysis¶

5. Computational Engine — From Transcriptomic Data to a Depth-2 QAOA Benchmark ¶

7. Interpreting the Benchmark-Wide Figures ¶