from collections import OrderedDict
import time
import os
import logging, warnings

from qiskit.optimization.algorithms import CplexOptimizer
from qiskit.optimization.algorithms.optimization_algorithm import OptimizationResultStatus

from dwave.system import DWaveCliqueSampler, LeapHybridSampler

from random_lp.lp_random_gen import create_models
from utilities.helpers import create_dwave_meo

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()
logger.setLevel(logging.INFO)
warnings.filterwarnings("ignore", category=DeprecationWarning)

qps = create_models(TEST_DIR)

# init Optimizers
cplex = CplexOptimizer()
hybrid = create_dwave_meo(LeapHybridSampler())
clique = create_dwave_meo(DWaveCliqueSampler(), num_reads=4096)
dwave_auto = create_dwave_meo(num_reads=4096)

results = OrderedDict()
count_fail_hybrid, count_fail_dwave = 0, 0

for qp_name, qp in qps.items() :
    
    print(qp_name)
    print("number of qubits: ", qp.qubo.get_num_vars())
    output_file_handler = logging.FileHandler(filename=RES_DIR + qp.name + ".log")
    logger.addHandler(output_file_handler)

    # set problem labels
    for solver, label in zip([hybrid, clique, dwave_auto], ["hybrid_", "clique_", "auto_emb_"]):
        solver.min_eigen_solver.sampler.set_label(label + qp_name + "_" + str(qp.complexity()))

    if qp.qubo.get_num_vars() > clique.min_eigen_solver.sampler.child.largest_clique_size:
        dwave = dwave_auto
    else:
        dwave = clique

    with open(RES_DIR + qp.name + '.res', 'w') as file:
        file.write(str("Start " + qp.name + "\n " + str(qp.qubo.get_num_vars()) + " qubits needed"))  
        file.write("\n Leap: \n")
        logger.info("\n Leap: \n ")
        res_classic = cplex.solve(qp)
        
        res_hybrid = hybrid.solve(qp)
        logger.info(str(res_hybrid.min_eigen_solver_result.sampleset.info))
        if res_hybrid.status != OptimizationResultStatus.SUCCESS:
            file.write("No solution found with DWave Hybrid Sampler Leap.")
            file.write("\n CPLEX:\n")
            file.write(str(res_classic))
            count_fail_hybrid = count_fail_hybrid + 1
            if count_fail_hybrid > 2 :
                break
        else:
            print("Leap successful!")
            count_fail_hybrid = 0
            if res_hybrid.fval == res_classic.fval:
                file.write("\n Leap found optimal solution\n")

            else:
                print("\n optimal value Leap   "+str(res_hybrid.fval) + "  ,  cplex:"+ str(res_classic.fval))
                file.write("\n Leap: \n")
                file.write(str(res_hybrid))
                file.write("\n CPLEX:\n")
                file.write(str(res_classic))
            if count_fail_dwave <= 2:
                file.write("\n DWave Quantum: \n")
                logger.info("\n DWave Quantum: \n ")
                try:
                    
                    res_quantum = dwave.solve(qp)
                    problem_id = str(res_quantum.min_eigen_solver_result.sampleset.info['problem_id'])
                    logger.info("\n problem id " + problem_id)
                    file.write("\n problem id: " + problem_id)
                    
                    if res_quantum.status != OptimizationResultStatus.SUCCESS:
                        print("Dwave quantum solver found no solution.")
                        file.write("\n No solution found with DWave Quantum Solver \n")
                        count_fail_dwave = count_fail_dwave + 1
                        if count_fail_dwave == 3:
                            file.write("\n Stop testing DWave Quantum Solver \n")
                    else:
                        print("Dwave Quantum successful!")
                        count_fail_dwave = 0
                        if res_quantum.fval == res_classic.fval:
                            file.write("\n DWave quantum solver found optimal solution\n")
                        else:
                            print("\n optimal value Dwave quantum "+str(res_quantum.fval) + "  ,  cplex:"+ str(res_classic.fval))
                            file.write("\n DWave quantum solver:\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 DWave Solver produced an exception:\n")
                    file.write(str(ex))
                
    logger.removeHandler(output_file_handler)