quantum_solver.py

"""
QAOA-based quantum solver for sawmill optimization.

Uses Quantum Approximate Optimization to solve the board-packing QUBO
directly. When hidden defects are discovered during simulated cutting,
QAOA re-optimizes the remaining volume globally — unlike the greedy
solver which just skips the defect-hit board and continues in its fixed
priority order.
"""

import time
import numpy as np
from qiskit.circuit import QuantumCircuit, Parameter
from qiskit.quantum_info import SparsePauliOp
from scipy.optimize import minimize as scipy_minimize

from data_generator import Log, _board_intersects_defect
from qubo_formulation import build_qubo, qubo_to_ising, _boards_overlap
from classical_solver import _boards_physically_overlap


MAX_QUBITS = 12


def _build_cost_hamiltonian(h, J, n):
    terms, coeffs = [], []
    for i, hi in h.items():
        if abs(hi) < 1e-12 or i >= n:
            continue
        label = ['I'] * n
        label[n - 1 - i] = 'Z'
        terms.append(''.join(label))
        coeffs.append(hi)
    for (i, j), jij in J.items():
        if abs(jij) < 1e-12 or i >= n or j >= n:
            continue
        label = ['I'] * n
        label[n - 1 - i] = 'Z'
        label[n - 1 - j] = 'Z'
        terms.append(''.join(label))
        coeffs.append(jij)
    if not terms:
        terms.append('I' * n)
        coeffs.append(0.0)
    return SparsePauliOp.from_list(list(zip(terms, coeffs)))


def _evaluate_qubo(bitstring, Q):
    n = Q.shape[0]
    x = np.array([int(b) for b in reversed(bitstring[-n:])])
    return float(x @ Q @ x)


def _solve_qaoa(Q, h, J, n, p=1, n_shots=1024):
    """Core QAOA: variational optimization + sampling."""
    from qiskit.primitives import StatevectorEstimator

    cost_op = _build_cost_hamiltonian(h, J, n)
    gammas = [Parameter(f'g{k}') for k in range(p)]
    betas = [Parameter(f'b{k}') for k in range(p)]

    est_qc = QuantumCircuit(n)
    est_qc.h(range(n))
    for k in range(p):
        for i, hi in h.items():
            if abs(hi) > 1e-12 and i < n:
                est_qc.rz(2 * hi * gammas[k], i)
        for (i, j), jij in J.items():
            if abs(jij) > 1e-12 and i < n and j < n:
                est_qc.rzz(2 * jij * gammas[k], i, j)
        for i in range(n):
            est_qc.rx(2 * betas[k], i)

    all_params = gammas + betas
    estimator = StatevectorEstimator()

    def cost_fn(params):
        pdict = {all_params[i]: params[i] for i in range(len(all_params))}
        bound = est_qc.assign_parameters(pdict)
        return float(np.real(estimator.run([(bound, cost_op)]).result()[0].data.evs))

    best_cost, best_params = float('inf'), None
    for _ in range(2):
        x0 = np.random.uniform(0, np.pi, 2 * p)
        res = scipy_minimize(cost_fn, x0, method='COBYLA',
                             options={'maxiter': 60, 'rhobeg': 0.5})
        if res.fun < best_cost:
            best_cost, best_params = res.fun, res.x

    # sample from optimized circuit
    sample_qc = QuantumCircuit(n)
    sample_qc.h(range(n))
    sg = [Parameter(f'sg{k}') for k in range(p)]
    sb = [Parameter(f'sb{k}') for k in range(p)]
    for k in range(p):
        for i, hi in h.items():
            if abs(hi) > 1e-12 and i < n:
                sample_qc.rz(2 * hi * sg[k], i)
        for (i, j), jij in J.items():
            if abs(jij) > 1e-12 and i < n and j < n:
                sample_qc.rzz(2 * jij * sg[k], i, j)
        for i in range(n):
            sample_qc.rx(2 * sb[k], i)
    sample_qc.measure_all()

    pdict = {sg[k]: best_params[k] for k in range(p)}
    pdict.update({sb[k]: best_params[p + k] for k in range(p)})

    from qiskit_aer import AerSimulator
    from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
    backend = AerSimulator()
    pm = generate_preset_pass_manager(backend=backend, optimization_level=1)
    transpiled = pm.run(sample_qc.assign_parameters(pdict))
    counts = backend.run(transpiled, shots=n_shots).result().get_counts()

    # evaluate all sampled bitstrings against QUBO
    return min(counts.keys(), key=lambda b: _evaluate_qubo(b, Q))


def _bitstring_to_selection(bitstring, candidates):
    selected = []
    for i, bit in enumerate(reversed(bitstring)):
        if bit == '1' and i < len(candidates):
            selected.append(candidates[i])
    return selected


def _validate_selection(selected):
    """Post-process: remove overlapping boards, keeping highest value."""
    valid = []
    for b in sorted(selected, key=lambda b: b.value, reverse=True):
        if not any(_boards_overlap(b, v) > 0 for v in valid):
            valid.append(b)
    return valid


def qaoa_solve(log: Log, p=1):
    """
    QAOA-based solve with defect re-optimization.

    Phase 1: Solve the full problem with QAOA (global optimization)
    Phase 2: Simulate cutting — discover hidden defects
    Phase 3: If defects found, re-optimize remaining space with QAOA
             (classical greedy would just skip and continue)
    """
    t0 = time.time()
    candidates = list(log.board_candidates)

    if not candidates:
        return {
            'selected_boards': [], 'total_value': 0.0,
            'waste_fraction': 1.0, 'solve_time': time.time() - t0,
            'defect_discoveries': 0, 'method': 'qaoa',
        }

    # pre-filter if needed
    if len(candidates) > MAX_QUBITS:
        for c in candidates:
            c._density = c.value / max(c.width * c.height, 1e-6)
        candidates.sort(key=lambda c: c._density, reverse=True)
        candidates = candidates[:MAX_QUBITS]

    n = len(candidates)
    for new_i, c in enumerate(candidates):
        c.index = new_i

    # Phase 1: QAOA on full problem (no hidden defect knowledge)
    # build QUBO ignoring hidden defects (they're unknown)
    visible_defects = [d for d in log.defects if not d.hidden]
    qaoa_log = Log(
        diameter_large=log.diameter_large,
        diameter_small=log.diameter_small,
        length=log.length, taper_rate=log.taper_rate,
        defects=visible_defects, board_candidates=candidates,
    )
    Q, _ = build_qubo(qaoa_log)
    h, J, _ = qubo_to_ising(Q)
    bits = _solve_qaoa(Q, h, J, n, p=p)
    selected = _bitstring_to_selection(bits, candidates)
    selected = _validate_selection(selected)

    # Phase 2: simulate cutting — discover hidden defects
    surviving = []
    defect_discoveries = 0
    for b in selected:
        if any(_board_intersects_defect(b, d) for d in log.defects if d.hidden):
            defect_discoveries += 1
        else:
            surviving.append(b)

    # Phase 3: re-optimize freed space if defects found
    if defect_discoveries > 0:
        surviving_idx = {b.index for b in surviving}
        remaining = [
            c for c in candidates
            if c.index not in surviving_idx
            and not any(_boards_physically_overlap(c, s) for s in surviving)
            and not any(_board_intersects_defect(c, d)
                        for d in log.defects if d.hidden)
        ]

        if remaining:
            for new_i, c in enumerate(remaining):
                c.index = new_i

            relog = Log(
                diameter_large=log.diameter_large,
                diameter_small=log.diameter_small,
                length=log.length, taper_rate=log.taper_rate,
                defects=log.defects,  # now include ALL defects
                board_candidates=remaining,
            )
            n2 = len(remaining)
            if n2 > MAX_QUBITS:
                remaining = remaining[:MAX_QUBITS]
                n2 = MAX_QUBITS

            Q2, _ = build_qubo(relog)
            h2, J2, _ = qubo_to_ising(Q2)
            bits2 = _solve_qaoa(Q2, h2, J2, n2, p=p)
            reopt = _bitstring_to_selection(bits2, remaining)
            reopt = _validate_selection(reopt)

            for rb in reopt:
                if not any(_boards_physically_overlap(rb, s) for s in surviving):
                    surviving.append(rb)

    total_value = sum(b.value for b in surviving)
    log_area = np.pi * (log.diameter_small / 2.0) ** 2
    board_area = sum(b.width * b.height for b in surviving)
    waste = 1.0 - min(board_area / max(log_area, 1e-6), 1.0)

    return {
        'selected_boards': surviving,
        'total_value': total_value,
        'waste_fraction': waste,
        'solve_time': time.time() - t0,
        'defect_discoveries': defect_discoveries,
        'method': 'qaoa',
    }