#!/usr/bin/env python3

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

from tqdm import tqdm

def main():
    #instance_parameters()
    #wmis_results()
    #wmis_siman_results_alpha_num_of_assignments()
    wmis_qpu_results_alpha_num_of_assignments()
    #primitive_2_siman_results_alpha_num_of_assignments()
    #primitive_5_siman_results_alpha_num_of_assignments()
    #primitive_5_qpu_results_alpha_num_of_assignments()
    #primitive_8_siman_results_alpha_num_of_assignments()
    #wmis_2_qpu_results()
    #minisat_runs()
    
    #wmis_engergy()
    #wmis_2_engergy()

def instance_parameters():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    instances = queries.Instance_scope_query(idb)
    instances.query("c42_vLogistic_6")
    
    insert_row = ("INSERT INTO c42_vLogistic_6_instances "
                  "(instance_id, "
                  " number_of_clauses, "
                  " number_of_variables) "
                  "VALUES (%s, %s, %s)")
                  
    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 wmis_siman_results():
    db = script.connect_to_instance_pool()
    
    q = queries.WMIS_result_scope_query_raw(db)
    q.query("c45_v[5-45]_1", "wmis_siman_results")
    
    for i in range(2):
        result = q.__next__()
        
        sample_set = queries.read_raw_wmis_sample_set(result["data"])
        
        data = script.analyze_wmis_sample(sample_set.first)
        
        print(data)
        
def wmis_siman_results_alpha_num_of_assignments():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query("c42_vLogistic_6", "wmis_siman_results")
    
    insert_row = ("INSERT INTO c42_vLogistic_6_wmis_siman_results "
                  "(result_id, "
                  " run, "
                  " instance_id, "
                  " number_of_found_assignments, "
                  " chain_break_fraction, "
                  " num_occurrences, "
                  " energy, "
                  " satisfiable) "
                  "VALUES (%s, %s, %s, %s, %s, %s, %s, %s) ")
    
    run = int(input("run: "))
    for result in tqdm(q):
        if "run" in result and result["run"] == run:
        #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)
        
            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))
        
    edb.commit()
    edb_cursor.close()
    edb.close()

    
def primitive_2_siman_results_alpha_num_of_assignments():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query("c42_vLogistic_6", "primitive_isings_2_siman_results")
    
    insert_row = ("INSERT INTO c42_vLogistic_6_primitive_2_siman_results "
                  "(result_id, "
                  " run, "
                  " instance_id, "
                  " chain_break_fraction, "
                  " num_occurrences, "
                  " energy, "
                  " satisfiable) "
                  "VALUES (%s, %s, %s, %s, %s, %s, %s) ")
    
        
    for result in tqdm(q):
        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"])
        
        model = queries.extract_primitive_ising_model(sample_set.first.sample)
        
        isSatisfiable = sat.checkAssignment(model)
        
        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"]),
                                        isSatisfiable))
        
    edb.commit()
    edb_cursor.close()
    edb.close()
    
    
def primitive_5_siman_results_alpha_num_of_assignments():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query("c42_vLogistic_6", "primitive_isigns_5_siman")
    
    insert_row = ("INSERT INTO c42_vLogistic_6_primitive_5_siman_results "
                  "(result_id, "
                  " run, "
                  " instance_id, "
                  " chain_break_fraction, "
                  " num_occurrences, "
                  " energy, "
                  " satisfiable, "
                  " num_conflicts, "
                  " monte_carlo_steps) "
                  "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) ")
    
        
    for result in tqdm(q):
        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 = __post_process_prim_5_sample(sat, sample_set.first.sample)
        
        #print(post_process_results)
        
        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(len(post_process_results["conflicts"])),
                                        int(post_process_results["monte_carlo_steps"])))
        
    edb.commit()
    edb_cursor.close()
    edb.close()
    
def primitive_5_qpu_results_alpha_num_of_assignments():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query("c42_vLogistic_6", "primitive_isings_5_qpu")
    
    insert_row = ("INSERT INTO c42_vLogistic_6_primitive_5_qpu_results "
                  "(result_id, "
                  " run, "
                  " instance_id, "
                  " chain_break_fraction, "
                  " num_occurrences, "
                  " energy, "
                  " satisfiable, "
                  " num_conflicts, "
                  " monte_carlo_steps, "
                  " anneal_time) "
                  "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) ")
    
        
    run = int(input("run: "))
    
    for result in tqdm(q):
        if True:
        #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 = __post_process_prim_5_sample(sat,
                                                                sample_set.first.sample)
            
            anneal_time = result["data"]["info"]["timing"]["qpu_anneal_time_per_sample"]
            #print(post_process_results)
            
            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(len(post_process_results["conflicts"])),
                                            int(post_process_results["monte_carlo_steps"]),
                                            int(anneal_time)))
        
    edb.commit()
    edb_cursor.close()
    edb.close()
    
def primitive_8_siman_results_alpha_num_of_assignments():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query("c42_vLogistic_6", "wmis_4_qubos_siman")
    
    insert_row = ("INSERT INTO c42_vLogistic_6_wmis_4_siman_results "
                  "(result_id, "
                  " run, "
                  " instance_id, "
                  " chain_break_fraction, "
                  " num_occurrences, "
                  " energy, "
                  " satisfiable) "
                  "VALUES (%s,  %s, %s, %s, %s, %s, %s) ")
    
        
    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 = __post_process_prim_5_sample(sat,
                                                                sample_set.first.sample)
            
            #print(post_process_results)
            
            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"])))
        
    edb.commit()
    edb_cursor.close()
    edb.close()
    
def wmis_2_qpu_results():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query("c42_vLogistic_6", "wmis_2_qubos_2_qpu")
    
    insert_row = ("INSERT INTO c42_vLogistic_6_wmis_2_vertex_eq_edge_qpu_results "
                  "(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) ")
    
        
    run = int(input("run: "))
    
    for result in tqdm(q):
        if True:
        #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 = __post_process_prim_5_sample(sat,
                                                                sample_set.first.sample)
            
            anneal_time = result["data"]["info"]["timing"]["qpu_anneal_time_per_sample"]
            #print(post_process_results)
            
            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 __post_process_prim_5_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
    
    conflicts = set()
    for var in vars:
        if var in assignments and -var in assignments:
            if assignments[var] == assignments[-var]:
                conflicts.add(var)
                
    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
            
    
    
    #var_list = list(conflicts)
    #
    #monte_carlo_steps = 0
    #if len(conflicts) > 0:
    #    for i in range(1000):
    #        rand_var = random.choice(var_list)
    #    
    #        if sat.checkAssignment(model):
    #            monte_carlo_steps
    #            break
    #    
    #        model[rand_var - 1] = not model[rand_var - 1]
    #    
    #        monte_carlo_steps += 1
        
    post_process_results["conflicts"] = conflicts
    post_process_results["satisfiable"] = sat.checkAssignment(model)
    #post_process_results["monte_carlo_steps"] = monte_carlo_steps
    
    return post_process_results

def wmis_qpu_results_alpha_num_of_assignments():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query("c42_vLogistic_6", "wmis_qpu_results")
    
    insert_row = ("INSERT INTO c42_vLogistic_6_wmis_qpu_results "
                  "(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) ")
    
    run = int(input("run: "))
    for result in tqdm(q):
        #if result["run"] == run:
        if True:
            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 wmis_engergy():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query("c42_vLogistic_6", "wmis_qpu_results")
    
    insert_row = ("INSERT INTO c42_vLogistic_6_wmis_qpu_energy "
                  "(result_id, "
                  " sample_index, "
                  " run, "
                  " instance_id, "
                  " energy, "
                  " anneal_time, "
                  " sample_size) "
                  "VALUES (%s, %s, %s, %s, %s, %s, %s) ")
    
    run = int(input("run: "))
    for result in tqdm(q):
        #if result["run"] == run:
        if True:
            sample_set = queries.read_raw_wmis_sample_set(result["data"])
        
        
            anneal_time = result["data"]["info"]["timing"]["qpu_anneal_time_per_sample"]
            
            sample_size = np.sum(sample_set.record["num_occurrences"])
            
            count = 1
            for energy, in sample_set.data(fields=["energy"], sorted_by=("energy")):
                
                edb_cursor.execute(insert_row, (str(result["_id"]),
                                                int(count),
                                                int(result["run"]),
                                                str(result["instance"]),
                                                int(energy),
                                                int(anneal_time),
                                                int(sample_size)))
                
                count += 1
        
    edb.commit()
    edb_cursor.close()
    edb.close()
    
def wmis_2_engergy():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor()
    
    idb = script.connect_to_instance_pool()
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query("c42_vLogistic_6", "wmis_2_qubos_2_qpu")
    
    insert_row = ("INSERT INTO c42_vLogistic_6_wmis_2_vertex_eq_edge_qpu_energy "
                  "(result_id, "
                  " sample_index, "
                  " run, "
                  " instance_id, "
                  " energy, "
                  " anneal_time, "
                  " sample_size) "
                  "VALUES (%s, %s, %s, %s, %s, %s, %s) ")
    
    run = int(input("run: "))
    for result in tqdm(q):
        if result["run"] == run:
        #if True:
            sample_set = queries.read_raw_wmis_sample_set(result["data"])
        
        
            anneal_time = result["data"]["info"]["timing"]["qpu_anneal_time_per_sample"]
            
            sample_size = np.sum(sample_set.record["num_occurrences"])
            
            count = 1
            for energy, in sample_set.data(fields=["energy"], sorted_by=("energy")):
                
                edb_cursor.execute(insert_row, (str(result["_id"]),
                                                int(count),
                                                int(result["run"]),
                                                str(result["instance"]),
                                                int(energy),
                                                int(anneal_time),
                                                int(sample_size)))
                
                count += 1
        
    edb.commit()
    edb_cursor.close()
    edb.close()
    
def minisat_runs():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor();
    
    idb = script.connect_to_instance_pool()
    
    runs = queries.Minisat_run_scope_query_raw(idb)
    runs.query("c42_vLogistic_6", "minisat_runs")
    
    insert_row = ("INSERT INTO c42_vLogistic_6_minisat_runs "
                  "(run_id, "
                  " instance_id, "
                  " satisfiable) "
                  "VALUES (%s, %s, %s) ")
    
    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 wmis_results():
    edb = script.connect_to_experimetns_db()
    edb_cursor = edb.cursor();
    
    idb = script.connect_to_instance_pool()
    
    q = queries.WMIS_result_scope_query_raw(idb)
    q.query("c42_v[5-84]_1", "wmis_qbsolv_results")
    
    for i in range(501):
        q.__next__()
        
    res = q.__next__()
    
    sample_set = queries.read_raw_wmis_sample_set(res["data"])
    
    
    model = script.majority_vote_sample(sample_set.first.sample)
    
    sat = queries.get_instance_by_id(idb["instances"], res["instance"])
    
    print(model)
    
    print(sat.getNumberOfVariables())
    print(sat.checkAssignment(model))
    
    
if __name__ == "__main__":
    main()