from . import script
from . import queries
import dimod
import random 
import numpy as np

from bson.objectid import ObjectId
import pandas as pd

from tqdm import tqdm

def extract_wmis_qpu_results():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    scope = input("scope: ")
    result_collection = input("result collection: ")
    
    table_name = input("table name: ")
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query(scope, result_collection)
    
    insert_row = ("INSERT INTO {} "
                  "(result_id, "
                  " run, "
                  " instance_id, "
                  " number_of_found_assignments, "
                  " chain_break_fraction, "
                  " num_occurrences, "
                  " energy, "
                  " satisfiable, "
                  " anneal_time, "
                  " energy_reach, "
                  " sample_size) "
                  "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) ").format(table_name)
    
    run = int(input("run: "))
    
    for result in tqdm(q):
        if result["run"] == run:
            sample_set = queries.read_raw_wmis_sample_set(result["data"])
        
            data = script.analyze_wmis_sample(sample_set.first)
        
            sat = queries.get_instance_by_id(idb["instances"], result["instance"])
        
            model = script.majority_vote_sample(sample_set.first.sample)
        
            isSatisfiable = sat.checkAssignment(model)
        
            anneal_time = result["data"]["info"]["timing"]["qpu_anneal_time_per_sample"]
    
            energy_reach = abs(sample_set.record["energy"].max() - 
                               sample_set.record["energy"].min())
            
            sample_size = np.sum(sample_set.record["num_occurrences"])
            
            edb_cursor.execute(insert_row, (str(result["_id"]),
                                            int(result["run"]),
                                            str(result["instance"]),
                                            int(data["number_of_assignments"]),
                                            float(data["chain_break_fraction"]),
                                            int(data["num_occurrences"]),
                                            int(data["energy"]),
                                            isSatisfiable,
                                            int(anneal_time),
                                            int(energy_reach),
                                            int(sample_size)))
        
    edb.commit()
    edb_cursor.close()
    edb.close()
    
def extract_wmis_2_qpu_results():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    scope = input("scope: ")
    result_collection = input("result collection: ")
    table_name = input("table name: ")
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query(scope, result_collection)
    
    insert_row = ("INSERT INTO {} "
                  "(result_id, "
                  " run, "
                  " instance_id, "
                  " chain_break_fraction, "
                  " num_occurrences, "
                  " energy, "
                  " satisfiable, "
                  " anneal_time, "
                  " energy_reach, "
                  " sample_size) "
                  "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) ").format(table_name)
    
        
    run = int(input("run: "))
    
    for result in tqdm(q):
        if result["run"] == run:
            sample_set = queries.read_raw_primitive_ising_sample_set(result["data"])
            
            data = script.analyze_wmis_sample(sample_set.first)
            
            sat = queries.get_instance_by_id(idb["instances"], result["instance"])
            
            post_process_results = __evaluate_wmis_2_sample(sat,
                                                            sample_set.first.sample)
            
            anneal_time = result["data"]["info"]["timing"]["qpu_anneal_time_per_sample"]
            
            energy_reach = abs(sample_set.record["energy"].max() - 
                               sample_set.record["energy"].min())
            sample_size = np.sum(sample_set.record["num_occurrences"])
            
            edb_cursor.execute(insert_row, (str(result["_id"]),
                                            int(result["run"]),
                                            str(result["instance"]),
                                            float(data["chain_break_fraction"]),
                                            int(data["num_occurrences"]),
                                            int(data["energy"]),
                                            bool(post_process_results["satisfiable"]),
                                            int(anneal_time),
                                            int(energy_reach),
                                            int(sample_size)))
        
    edb.commit()
    edb_cursor.close()
    edb.close()

def __evaluate_wmis_2_sample(sat, sample):
    post_process_results = {}
    
    assignments = {}
    vars = set()
    
    for node, energy in sample.items():
        if node[0] == "x":
            lit = int(node[1:])
            
            vars.add(abs(lit))
            
            assignments[lit] = energy
                
    model = [True for i in range(len(vars))]
    
    for var in vars:
        if var in assignments:
            model[var - 1] = True if assignments[var] == 1 else False
        elif -var in assignments:
            model[var - 1] = True if assignments[-var] == 0 else False
            
    post_process_results["satisfiable"] = sat.checkAssignment(model)
    
    return post_process_results

def extract_minisat_results():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor();
    
    idb = script.connect_to_instance_pool()
    
    scope = input("scope: ")
    table_name = input("table name: ")
    
    runs = queries.Minisat_run_scope_query_raw(idb)
    runs.query(scope, "minisat_runs")
    
    insert_row = ("INSERT INTO {} "
                  "(run_id, "
                  " instance_id, "
                  " satisfiable) "
                  "VALUES (%s, %s, %s) ").format(table_name)
    
    for run in tqdm(runs):
        data = script.analyde_minisat_run(run)
        
        edb_cursor.execute(insert_row, (str(run["_id"]), 
                                        str(run["instance"]),
                                        int(data["satisfiable"])))
    
    edb.commit()
    edb_cursor.close()
    edb.close()
    
def extract_instance_parameters():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    scope = input("scope: ")
    table_name = input("table name: ")
    
    instances = queries.Instance_scope_query(idb)
    instances.query(scope)
    
    insert_row = ("INSERT INTO {} "
                  "(instance_id, "
                  " number_of_clauses, "
                  " number_of_variables) "
                  "VALUES (%s, %s, %s)").format(table_name)
                  
    for instance, instance_id in tqdm(instances):
        edb_cursor.execute(insert_row, (str(instance_id),
                                        int(instance.getNumberOfClauses()),
                                        int(instance.getNumberOfVariables())))
    
    edb.commit()
    edb_cursor.close()
    edb.close()

def extract_embedding_data():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()

    idb = script.connect_to_instance_pool()

    scope = input("scope: ")
    table_name = input("table name: ")

    qubo_collection = input("qubo collection: ")
    solver_graph_id = ObjectId(input("solver graph id: "))

    solver_input = queries.WMIS_solver_input_scope_query(idb, qubo_collection)
    solver_input.query(scope, solver_graph_id)

    results = pd.DataFrame(columns=["instance_id",
                                    "embedding_list_id",
                                    "emb_list_index",
                                    "num_qubits",
                                    "avgr_chain_length",
                                    "median_chain_length"]);
    
    for soin in solver_input:
        emb_index = 0;
        for emb in soin["embeddings"]:
            chain_lengths = []
            qubits = set()

            for node, chain in emb.items():
                
                chain_lengths.append(len(chain))

                qubits = qubits.union(set(chain))

            results = results.append({"instance_id": soin["instance_id"],
                                      "embedding_list_id": soin["embeddings_id"],
                                      "emb_list_index": emb_index,
                                      "num_qubits": len(qubits),
                                      "avgr_chain_length": np.mean(chain_lengths),
                                      "median_chain_length": np.median(chain_lengths)},
                                     ignore_index=True)

    insert_row = ('''INSERT INTO {} 
                     (instance_id, 
                      embedding_list_id, 
                      emb_list_index, 
                      num_qubits, 
                      avgr_chain_length, 
                      median_chain_length) 
                     VALUES(%s, %s, %s, %s, %s, %s)''').format(table_name)

    for index, row in results.iterrows():
        edb_cursor.execute(insert_row, (str(row["instance_id"]),
                                        str(row["embedding_list_id"]),
                                        int(row["emb_list_index"]),
                                        int(row["num_qubits"]),
                                        float(row["avgr_chain_length"]),
                                        float(row["median_chain_length"])))
    
    edb.commit()
    edb_cursor.close()
    edb.close()