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

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()