import time, warnings
from typing import Tuple
warnings.simplefilter("ignore")
import numpy as np

from qiskit import IBMQ
from qiskit.test.mock import FakeMumbai
from qiskit.test.mock.fake_qasm_backend import FakeQasmBackend
from qiskit.providers.aer import AerSimulator
from qiskit.providers.aer.noise import NoiseModel
from qiskit_optimization import QuadraticProgram
from qiskit_optimization.algorithms import MinimumEigenOptimizer as MEO
from qiskit_optimization.algorithms import MinimumEigenOptimizationResult as MEORes
from qiskit_optimization.algorithms import CplexOptimizer

import dimod
from dwave.system import DWaveSampler, EmbeddingComposite, DWaveCliqueSampler


from utilities.helpers import create_qaoa_meo, create_quadratic_programs_from_paths
DIR = 'TEST_DATA' + "/" + time.strftime("%d_%m_%Y")
DIR

# select QUBOs to solve
qps_dense = create_quadratic_programs_from_paths(DIR + "/DENSE/")
qp_dense = qps_dense['test_3']
qps_sparse = create_quadratic_programs_from_paths(DIR + "/SPARSE/")
qp_sparse = qps_sparse['test_3']
qps_multiple = create_quadratic_programs_from_paths(DIR + "/MULTIPLE/")
qp_mult = qps_multiple["test_20"]

qp_cplex = create_quadratic_programs_from_paths(DIR + "/CPLEX/")['test_cplex']

# set parameters and arguments
num_reads = 1024
classic_iters = 5

def qaoa_callback(eval_ct: int, opt_pars: np.ndarray, mean: float, stdev: float) -> None:
    """Print parameterset of last iteration."""
    if eval_ct == classic_iters:
        print("Evaluation count reached ", eval_ct, "with pars:", opt_pars)

device = FakeMumbai()
qaoa_kwargs = {"shots": num_reads, "noise_model": NoiseModel.from_backend(device),
                "max_iter": classic_iters, "qaoa_callback": qaoa_callback, "optimization_level": 3}

dwave_parameters = {'num_reads': num_reads, 'annealing_time': 3}

def init_qaoa_optimizers(device: FakeQasmBackend) -> Tuple[MEO, MEO]:
    # init local backend simulator with noise model
    local = AerSimulator.from_backend(device)
    qaoa_local = create_qaoa_meo(backend=local, **qaoa_kwargs)
    # init IBM Q Experience Simulator
    IBMQ.load_account()
    ibmq = IBMQ.get_provider(hub='ibm-q').get_backend('simulator_statevector')
    conf = device.configuration()
    qaoa_ibmq_sim = create_qaoa_meo(backend=ibmq, coupling_map=conf.coupling_map, basis_gates=conf.basis_gates,   
                                **qaoa_kwargs)
    return qaoa_local, qaoa_ibmq_sim

# init Optimizers
cplex = CplexOptimizer()
qaoa_local_sim, qaoa_ibmq_sim = init_qaoa_optimizers(device)

dwave_dense = DWaveCliqueSampler()
dwave_sparse = EmbeddingComposite(DWaveSampler()) # TODO rename to dwave_auto

def dwave_solve_qp(sampler: dimod.Sampler, qubo: QuadraticProgram):
    bqm = dimod.as_bqm(qubo.objective.linear.to_array(), qubo.objective.quadratic.to_array(), dimod.BINARY)
    sampleset = sampler.sample(bqm, **dwave_parameters)
    return sampleset

def eval_qaoa_result(res: MEORes) -> Tuple[float, float]:
    prob = res.samples[0].probability
    fval = res.samples[0].fval
    return prob, fval

def eval_bqm_sampleset(sampleset) -> Tuple[float, float]:
    prob = sampleset.record.num_occurrences[0] / num_reads
    fval = sampleset.record.energy[0]
    return prob, fval

# solve classically
for qp in (qp_dense, qp_sparse):
    print(qp.name)
    print(cplex.solve(qp))

# solve by using Mumbai noise model with local qasm sim
res_dense_qaoa = qaoa_local_sim.solve(qp_dense)
eval_qaoa_result(res_dense_qaoa)

# solve by using D-Wave Pegasus Advantage system
res_dense_dwave = dwave_solve_qp(dwave_dense, qp_dense)
eval_bqm_sampleset(res_dense_dwave)

qp_sparse = qps_sparse['test_8']
cplex.solve(qp_sparse)

# solve by using Mumbai noise model with IBM Q Experience simulator
warnings.filterwarnings("ignore", category=DeprecationWarning)
res_sparse_qaoa = qaoa_ibmq_sim.solve(qp_sparse)
eval_qaoa_result(res_sparse_qaoa)

cplex.solve(qp_mult) # takes just a few minutes

# solve by using D-Wave Pegasus Advantage system
res_mult_dwave = dwave_solve_qp(dwave_sparse, qp_mult)
eval_bqm_sampleset(res_mult_dwave)

start_time = time.perf_counter()
res_cplex = cplex.solve(qp_cplex)
cplex_runtime = time.perf_counter() - start_time
print(res_cplex)
cplex_runtime

start_time = time.perf_counter()
res_cplex_dwave = dwave_solve_qp(dwave_sparse, qp_cplex)
dwave_runtime = time.perf_counter() - start_time
print(eval_bqm_sampleset(res_cplex_dwave))
dwave_runtime

