import time, os, logging, 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 qiskit_optimization.algorithms.optimization_algorithm import OptimizationResultStatus
from utilities.helpers import create_qaoa_meo, create_quadratic_programs_from_paths

TEST_DIR = 'TEST_DATA' + "/" + time.strftime("%d_%m_%Y") + "/DENSE/"
RES_DIR = 'RESULTS' + "/" + time.strftime("%d_%m_%Y") + "/DENSE/"
os.makedirs(RES_DIR, exist_ok=True)
TEST_DIR

logger = logging.getLogger()
max_iter = 5 # set maximum number of optimizer iterations in qaoa online simulation
def qaoa_callback(eval_ct: int, opt_pars: np.ndarray, mean: float, stdev: float) -> None:
    """Print number of iteration in QAOA and log."""
    print("Evaluation count:", eval_ct)
    logger.info(f"\n Evaluation count {eval_ct} with parameters {opt_pars} and mean {mean:.2f}")
    if eval_ct == max_iter:
        logger.info(" reached maximum number of evaluations \n")

# select linear programs to solve
qps = create_quadratic_programs_from_paths(TEST_DIR)
# 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, "optimization_level": 3}
qaoa_local_sim = create_qaoa_meo(backend=local, max_iter=max_iter, **quantum_instance_kwargs)
qaoa_ibmq_sim = create_qaoa_meo(backend=ibmq, coupling_map=conf.coupling_map, basis_gates=conf.basis_gates,   
                                  max_iter=max_iter, qaoa_callback=qaoa_callback, **quantum_instance_kwargs)

# set online job informations in IBM Q Experience
qaoa_ibmq_sim.min_eigen_solver.quantum_instance.backend_options["job_tags"] = ["qaoa", "dense", "mumbai"]
def set_job_name(qp):
    qaoa_ibmq_sim.min_eigen_solver.quantum_instance.backend_options["job_name"] = "qaoa_dense_" \
        + qp.name + "_" + str(qp.get_num_vars())

warnings.filterwarnings("ignore", category=DeprecationWarning)
logger.setLevel(logging.INFO)
count_fail_qaoa = 0

for qp_name, qp in qps.items() :
    print(qp_name)
    print("number of qubits: ", qp.get_num_vars())
    output_file_handler = logging.FileHandler(filename=RES_DIR + qp.name + ".log")
    logger.addHandler(output_file_handler)
    if qp.get_num_vars() >= 15:
        set_job_name(qp)
        qaoa = qaoa_ibmq_sim
        logstring = "ibmq_simulator_statevector"
    else:
        qaoa = qaoa_local_sim
        logstring = "local_qasm_simulator"
        
    with open(RES_DIR + qp.name + '.res', 'w') as file:
        file.write(str("Start " + qp.name + "\n " + str(qp.get_num_vars()) + " qubits needed"))  
        file.write("\n " + logstring + " \n")
        logger.info("\n " + logstring + " \n")
        print(logstring)
        try:
            res_classic = cplex.solve(qp)
            res_quantum = qaoa.solve(qp)
            if res_quantum.status != OptimizationResultStatus.SUCCESS:
                print("QAOA minimum eigen solver found no solution.")
                file.write("\n QAOA minimum eigen solver found no solution. \n")
                count_fail_qaoa = count_fail_qaoa + 1
                if count_fail_qaoa == 3:
                    file.write("\n Stop testing \n")
                    break
            else:
                print("QAOA minimum eigen solver successful!")
                count_fail_qaoa = 0
                if res_quantum.fval == res_classic.fval:
                    file.write("\n QAOA found optimal solution\n")
                else:
                    print("\n optimal value QAOA "+str(res_quantum.fval) \
                                           + "  ,  cplex:"+ str(res_classic.fval))
                    file.write("\n QAOA:\n")
                    file.write(str(res_quantum))
                    file.write("\n CPLEX:\n")
                    file.write(str(res_classic))

        except Exception as ex:
            print(qp_name, " ", type(ex).__name__, " : ", ex)
            file.write("\n Cplex or QAOA solver produced an exception:\n")
            file.write(str(ex))
            count_fail_qaoa = count_fail_qaoa + 1
            if count_fail_qaoa == 3:
                file.write("\n Stop testing because of Exception! \n")
                break
                
    logger.removeHandler(output_file_handler)