import time, warnings
import numpy as np
from qiskit import IBMQ
from qiskit.test.mock import FakeMumbai
from qiskit.providers.aer import AerSimulator
from qiskit.providers.aer.noise import NoiseModel
from qiskit_optimization.algorithms import CplexOptimizer
from utilities.helpers import create_qaoa_meo, create_quadratic_programs_from_paths
DIR = 'TEST_DATA' + "/" + time.strftime("%d_%m_%Y")
DIR

def job_callback(job_id, job_status, queue_position, job):
    # BUG ? This is not called and not set in quantum instance (see info logging)
    print(job_id)
    print(job_status)
    print(queue_position)
    print(job)

def qaoa_callback(eval_ct: int, opt_pars: np.ndarray, mean: float, stdev: float) -> None:
    """Print number of iteration in QAOA."""
    print("Evaluation count:", eval_ct)

# select linear programs 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']

# init local backend simulator with noise model
device = FakeMumbai()
local = AerSimulator.from_backend(device)
noise_model = NoiseModel.from_backend(device)
conf = device.configuration()

# init IBM Q Experience Simulator
IBMQ.load_account()
ibmq = IBMQ.get_provider(hub='ibm-q').get_backend('simulator_statevector')

# init Optimizers
cplex = CplexOptimizer()
quantum_instance_kwargs = {"shots": 4096, "noise_model": noise_model, 
                            "job_callback": job_callback, "optimization_level": 3}

qaoa_local_sim = create_qaoa_meo(backend=local, **quantum_instance_kwargs)
qaoa_ibmq_sim = create_qaoa_meo(backend=ibmq, coupling_map=conf.coupling_map, basis_gates=conf.basis_gates,   
                                  max_iter=5, qaoa_callback=qaoa_callback, **quantum_instance_kwargs)

# solve classically
cplex.solve(qp_sparse)

# solve by using noise model with local qasm sim
qaoa_local_sim.solve(qp_sparse)

cplex.solve(qp_dense)

qaoa_local_sim.solve(qp_dense)

# solve by using noise model with IBM Q Experience simulator
qp_sparse = qps_sparse['test_7']
warnings.filterwarnings("ignore", category=DeprecationWarning)
qaoa_ibmq_sim.solve(qp_sparse)

cplex.solve(qp_sparse)

# solve by using noise model with IBM Q Experience simulator
qp_dense = qps_dense['test_15']
warnings.filterwarnings("ignore", category=DeprecationWarning)
qaoa_ibmq_sim.solve(qp_dense)

cplex.solve(qp_dense)