import numpy as np

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 (
    QuantumRandomAccessEncoding,
    QuantumRandomAccessOptimizer,
    SemideterministicRounding,
)

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=3, draw=True)
print(problem.export_as_lp_string())

# Create an encoding object with a maximum of 3 variables per qubit, aka a (3,1,p)-QRAC
encoding = QuantumRandomAccessEncoding(max_vars_per_qubit=3)

# Encode the QUBO problem into an encoded Hamiltonian operator
encoding.encode(problem)

# This is our encoded operator
print(f"Our encoded Hamiltonian is:\n( {encoding.qubit_op} ).\n")
print(
    "We achieve a compression ratio of "
    f"({encoding.num_vars} binary variables : {encoding.num_qubits} qubits) "
    f"≈ {encoding.compression_ratio}.\n"
)

# Create a QuantumInstance for solving the relaxed Hamiltonian using VQE
relaxed_qi = QuantumInstance(backend=Aer.get_backend("aer_simulator"), shots=1024)

# 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 semideterministic rounding, known as "Pauli rounding"
# in https://arxiv.org/pdf/2111.03167v2.pdf
# (This is the default if no rounding scheme is specified.)
rounding_scheme = SemideterministicRounding()

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

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

# relaxed function value
results.relaxed_fval

# optimal function value
results.fval

# optimal value
results.x

# status
results.status

print(
    f"The obtained solution places a partition between nodes {np.where(results.x == 0)[0]} "
    f"and nodes {np.where(results.x == 1)[0]}."
)

results.relaxed_results

results.rounding_results.samples

# pylint: disable=wrong-import-order
from qiskit_optimization.algorithms import CplexOptimizer

cplex_optimizer = CplexOptimizer()
cplex_results = cplex_optimizer.solve(problem)
exact_fval = cplex_results.fval

cplex_results

approximation_ratio = qrao_fval / exact_fval
print("QRAO Approximate Optimal Function Value:", qrao_fval)
print("Exact Optimal Function Value:", exact_fval)
print(f"Approximation Ratio: {approximation_ratio:.2f}")

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

%qiskit_version_table
%qiskit_copyright