#!/usr/bin/env python3

import util.script as script
import util.queries as queries
import util.graph as graph

import networkx as nx
import dwave_networkx as dnx
from neal import SimulatedAnnealingSampler
import dimod
from dwave.system.composites import FixedEmbeddingComposite
from dwave.system.samplers import DWaveSampler

from tqdm import tqdm

def main():
    mode = __get_mode()
    
    if mode == "SIMAN":
        __run_siman()
    elif mode == "QPU":
        __run_qpu()
    
def __get_mode():
    print("choose mode:")
    print("(1) simulated annealing")
    print("(2) qpu")
    mode = int(input())
    
    if mode == 1:
        return "SIMAN"
    elif mode == 2:
        return "QPU"

def __run_siman():
    db = script.connect_to_instance_pool()
    target_graph =  dnx.chimera_graph(16, 16, 4)

    solver_input_query = __get_solver_input_query(db, target_graph)
    
    base_sampler = SimulatedAnnealingSampler()
    chimera_sampler = dimod.StructureComposite(base_sampler,
                                               target_graph.nodes(),
                                               target_graph.edges())
    
    __run_on_scope(solver_input_query, db["wmis_siman_results"], base_sampler)


def __run_qpu():
    db = script.connect_to_instance_pool()
        
    base_solver = DWaveSampler()
    
    solver_input_query = __get_solver_input_query(db, solver_graph)
    
    solver_args = {}
    solver_args["annealing_time"] = int(input("annealing time (in us): "))
    
    __run_on_scope(solver_input_query, db["wmis_qpu_results"], base_solver, solver_args)
    
        

def __get_solver_graph_id(db, solver):
    solver_graph = graph.create_qpu_solver_nxgraph(solver)
    return queries.get_id_of_solver_graph(db["solver_graphs"],
                                          nx.node_link_data(solver_graph))

def __get_solver_input_query(db, solver):
    solver_graph_id = __get_solver_graph_id(db, solver)
    
    scope = input("scope: ")
    
    solver_input_query = queries.WMIS_solver_input_scope_query(db)
    solver_input_query.query(scope, solver_graph_id)
    
def __run_on_scope(solver_input_query, result_collection, base_solver, solver_args={}):
    run = int(input("save as run (numbered): "))
    
    for solver_input in tqdm(solver_input_query):
        
        embedding = solver_input["embeddings"][0]
        qubo = solver_input["qubo"]
        
        solver = FixedEmbeddingComposite(base_solver, embedding)
        
        res = solver.sample_qubo(qubo, **solver_args)
        
        script.save_result(result_collection,
                           res,
                           solver_input,
                           emb_list_index = 0,
                           run = run)
    
    
    
if __name__ == "__main__":
    main()