from qiskit.utils import QuantumInstance
from qiskit.algorithms.minimum_eigen_solvers import VQE
from qiskit.circuit.library import RealAmplitudes
from qiskit.algorithms.optimizers import SPSA
from qiskit_aer import Aer

from qrao import (
    QuantumRandomAccessOptimizer,
    QuantumRandomAccessEncoding,
    SemideterministicRounding,
    MagicRounding,
)
from qrao.utils import get_random_maxcut_qp

# Generate a random QUBO in the form of a QuadraticProgram
problem = get_random_maxcut_qp(degree=3, num_nodes=6, seed=1, draw=True)

# Create and encode the problem using the (3,1,p)-QRAC
encoding = QuantumRandomAccessEncoding(max_vars_per_qubit=3)
encoding.encode(problem)

backend = Aer.get_backend("aer_simulator")

# Create a QuantumInstance for solving the relaxed Hamiltonian using VQE
relaxed_qi = QuantumInstance(backend=backend, shots=500)

# Create a QuantumInstance for magic rounding
rounding_qi = QuantumInstance(backend=backend, shots=1000)

# Set up the variational quantum eigensolver (ansatz width is determined by the encoding)
vqe = VQE(
    ansatz=RealAmplitudes(encoding.num_qubits),
    optimizer=SPSA(maxiter=50),
    quantum_instance=relaxed_qi,
)

# Use magic rounding
rounding_scheme = MagicRounding(rounding_qi)

# Construct the optimizer
qrao = QuantumRandomAccessOptimizer(
    encoding=encoding, min_eigen_solver=vqe, rounding_scheme=rounding_scheme
)

# Solve the optimization problem
results = qrao.solve()
print(results)

# Calculate number of consolidated samples
len(results.samples)

# The first solution sample (which is not optimal in this case since the optimal
# function value is 14.0)
results.samples[0]

print(f"The optimal function value is {results.fval} at {results.x}.")

# Solve the relaxed problem manually
relaxed_results, rounding_context = qrao.solve_relaxed()

# Print the MinimumEigensolverResult object, from Terra
print(relaxed_results)

mr_results = MagicRounding(rounding_qi).round(rounding_context)
print(
    f"Collected {len(mr_results.samples)} samples in {mr_results.time_taken} seconds."
)

sdr_results = SemideterministicRounding().round(rounding_context)
print(f"Performed semideterministic rounding ({len(sdr_results.samples)} sample).")
print(f"Result: {sdr_results.samples[0]}")

import qiskit.tools.jupyter  # pylint: disable=unused-import,wrong-import-order

%qiskit_version_table
%qiskit_copyright