#!/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

__QUBO__ = 1
__ISING__ = 2

def main():
    mode = __get_mode()
    
    ising_qubo_collection = input("ising/qubo collection: ")
    
    model_type = __get_model_type()
    
    result_collection = input("result collection: ")
    
    if mode == "SIMAN":
        __run_siman(ising_qubo_collection, model_type, result_collection)
    elif mode == "QPU":
        __run_qpu(ising_qubo_collection, model_type, result_collection)
    
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 __get_model_type():
    print("model types:")
    print("(q) qubo")
    print("(i) ising")
    model_type = input("choose: ")
    
    if model_type == "q":
        return __QUBO__
    if model_type == "i":
        return __ISING__
    
def __run_siman(ising_qubo_collection, model_type, result_collection):
    db = script.connect_to_instance_pool()
    target_graph =  dnx.chimera_graph(16, 16, 4)
    target_graph_id = queries.get_id_of_solver_graph(db["solver_graphs"],
                                                     nx.node_link_data(target_graph))
    
    solver_input_query = __get_solver_input_query(db,
                                                  target_graph_id,
                                                  ising_qubo_collection)
    
    base_sampler = SimulatedAnnealingSampler()
    chimera_sampler = dimod.StructureComposite(base_sampler,
                                               target_graph.nodes(),
                                               target_graph.edges())
    
    __run_on_scope(solver_input_query, db[result_collection], chimera_sampler, model_type)


def __run_qpu(ising_qubo_collection, model_type, result_collection):
    db = script.connect_to_instance_pool()
        
    base_solver = DWaveSampler()
    solver_graph_id = __get_solver_graph_id(db, base_solver.solver)
    
    solver_input_query = __get_solver_input_query(db,
                                                  solver_graph_id,
                                                  ising_qubo_collection)

    solver_args = {}
    solver_args["annealing_time"] = int(input("annealing time (in us): "))
    solver_args["num_reads"] = int(input("number of reads per sample: "))
    
    __run_on_scope(solver_input_query,
                   db[result_collection],
                   base_solver,
                   model_type,
                   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_graph_id, ising_qubo_collection):
    scope = input("scope: ")
    
    solver_input_query = queries.WMIS_solver_input_scope_query(db, ising_qubo_collection)
    solver_input_query.query(scope, solver_graph_id)
    
    return solver_input_query
    
def __run_on_scope(solver_input_query,
                   result_collection,
                   base_solver,
                   model_type,
                   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)
        
        if model_type == __QUBO__:
            res = solver.sample_qubo(qubo, **solver_args)
        elif model_type == __ISING__:
            h, J = graph.split_ising(qubo)
            
            res = solver.sample_ising(h, J, **solver_args)
        
        script.save_sample_set(result_collection,
                               res,
                               solver_input,
                               emb_list_index = 0,
                               run = run)
    
    
    
if __name__ == "__main__":
    main()