#!/usr/bin/env python3

from util import SAT2QUBO as s2q
from util import randomSAT as rs
from util import queries
from util import graph
from util import script

import networkx as nx
import dwave_networkx as dnx
import minorminer
from dwave_qbsolv import QBSolv
import dimod
from dwave.system.composites import FixedEmbeddingComposite
from neal import SimulatedAnnealingSampler

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import re
from tqdm import tqdm

def frange(start, stop, steps):
    while start < stop:
        yield start
        start += steps

def test_kv_range():
    
    k = 75
    v = 40
    
    data = {
        "p_node_counts" : [],
        "wmis_node_counts" : [],

        "p_conn_counts" : [],
        "wmis_conn_counts" : [],
        
        "p_avrg_connectivity" : [],
        "wmis_avrg_connectivity" : [],
    
        "a_arr" : [],
        
        "k_arr" : []
    }
    
    target = dnx.chimera_graph(25, 25, 4)
    
    print(target.number_of_nodes())
    
    for r in range(2):
        #for k in range(50, 5001, 50):
        for a in frange(0.5, 8, 0.5):
            #v = int(k/4)
            #a = k/v
            
            #v = int(k/a)
            k = int(a*v)
            
            print("k={} v={} k/v={}".format(k, v, k/v))
            
            ksatInstance = rs.generateRandomKSAT(k, v, 3)
            
            p_qubo = s2q.WMISdictQUBO_2(ksatInstance)

            wmis_qubo = s2q.WMISdictQUBO(ksatInstance)
            
            qg = nx.Graph()
        
            for coupler, energy in wmis_qubo.items():
                if coupler[0] != coupler[1]:
                    qg.add_edge(coupler[0], coupler[1], weight=energy)
                    
            
            ig = nx.Graph()
            for coupler, energy in p_qubo.items():
                if coupler[0] != coupler[1]:
                    ig.add_edge(coupler[0], coupler[1], weight=energy)
            
            
            p_node_count = ig.number_of_nodes();
            p_conn_count = ig.number_of_edges();

                    
            wmis_node_count = qg.number_of_nodes();
            wmis_conn_count = qg.number_of_edges();
            
                
            
            print("p_qubo nodes= {}".format(p_node_count));
            print("wwmis qubo nodes= {}".format(wmis_node_count));
            
            print("nodes= {}".format(p_node_count / wmis_node_count));
            
            
            data["p_node_counts"].append(p_node_count)
            data["wmis_node_counts"].append(wmis_node_count)
            
            data["p_conn_counts"].append(p_conn_count)
            data["wmis_conn_counts"].append(wmis_conn_count)
            
            print("calculating connectivity...")
            data["p_avrg_connectivity"].append(nx.edge_connectivity(ig))
            print("calculating connectivity...")
            data["wmis_avrg_connectivity"].append(nx.edge_connectivity(qg))
            
            data["a_arr"].append(k/v)
            data["k_arr"].append(k)
            
            print()

    sns.set()

    ax = sns.scatterplot(x="a_arr", y="p_node_counts", data=data, label="p_qubo")
    ax.set(xlabel='k/v', ylabel='used verticies')
    ax = sns.scatterplot(x="a_arr", y="wmis_node_counts", data=data, ax=ax, label="wmis_qubo")
    plt.show()
    
    ax = sns.scatterplot(x="a_arr", y="p_conn_counts", data=data, label="p_qubo")
    ax.set(xlabel='k/v', ylabel='used edges')
    ax = sns.scatterplot(x="a_arr", y="wmis_conn_counts", data=data, ax=ax, label="wmis_qubo")
    plt.show()
    
    ax = sns.scatterplot(x="a_arr", y="p_avrg_connectivity", data=data, label="p_qubo")
    ax.set(xlabel='k/v', ylabel='avrg connectivity')
    ax = sns.scatterplot(x="a_arr", y="wmis_avrg_connectivity", data=data, ax=ax, label="wmis_qubo")
    plt.show()
    
    
def test_kv_range_emb():
    
    k = 75

    data = {
        "p_node_counts" : [],
        "wmis_node_counts" : [],

        "p_conn_counts" : [],
        "wmis_conn_counts" : [],
    
        "a_arr" : [],
        
        "k_arr" : []
    }
    
    target = dnx.chimera_graph(25, 25, 4)
    
    print(target.number_of_nodes())
    
    for r in range(2):
        #for k in range(50, 5001, 50):
        for a in frange(3.5, 5.5, 0.1):
            #v = int(k/4)
            #a = k/v
            
            v = int(k/a)
            
            print("k={} v={} k/v={}".format(k, v, k/v))
            
            ksatInstance = rs.generateRandomKSAT(k, v, 3)
            
            p_qubo = s2q.WMISdictQUBO_2(ksatInstance)

            wmis_qubo = s2q.WMISdictQUBO(ksatInstance)
            
            qg = nx.Graph()
        
            for coupler, energy in wmis_qubo.items():
                if coupler[0] != coupler[1]:
                    qg.add_edge(coupler[0], coupler[1], weight=energy)
                    
            
            ig = nx.Graph()
            for coupler, energy in p_qubo.items():
                if coupler[0] != coupler[1]:
                    ig.add_edge(coupler[0], coupler[1], weight=energy)
            
            qemb = minorminer.find_embedding(qg.edges(), target.edges(), return_overlap=True, threads=8)
            print("wmis emb. found: {}".format(qemb[1]))
            iemb = minorminer.find_embedding(ig.edges(), target.edges(), return_overlap=True, threads=8)
            print("primitive emb. found: {}".format(iemb[1]))
            
            if qemb[1] == 0 or iemb[1] == 0:
                print()
                continue
            
            p_node_count = 0;
            p_conn_count = 0;
            
            used_nodes = []
            for var, chain in iemb[0].items():
                used_nodes.extend(chain)
            
            p_node_count = len(np.unique(used_nodes))
                    
            wmis_node_count = 0;
            wmis_conn_count = 0;
            
            used_nodes = []
            for var, chain in qemb[0].items():
                used_nodes.extend(chain)
            
            wmis_node_count = len(np.unique(used_nodes))
                
            
            #print("p_qubo:    nodes={} connections={}".format(p_node_count, p_conn_count))
            
            #print("wmis_qubo: nodes={} connections={}".format(wmis_node_count, wmis_conn_count))
            
            print("p_qubo nodes= {}".format(p_node_count));
            print("wwmis qubo nodes= {}".format(wmis_node_count));
            
            print("nodes= {}".format(p_node_count / wmis_node_count));
            #print("conns= {}".format(p_conn_count / wmis_conn_count));
            
            
            data["p_node_counts"].append(p_node_count)
            data["wmis_node_counts"].append(wmis_node_count)
            
            data["p_conn_counts"].append(p_conn_count)
            data["wmis_conn_counts"].append(wmis_conn_count)
            
            data["a_arr"].append(k/v)
            data["k_arr"].append(k)
            
            print()

    sns.set()

    ax = sns.scatterplot(x="a_arr", y="p_node_counts", data=data, label="p_qubo")
    ax.set(xlabel='k/v', ylabel='used verticies after embedding')
    ax = sns.scatterplot(x="a_arr", y="wmis_node_counts", data=data, ax=ax, label="wmis_qubo")
    plt.show()
    
def test_k_range():
    
    k = 75

    data = {
        "p_node_counts" : [],
        "wmis_node_counts" : [],

        "p_conn_counts" : [],
        "wmis_conn_counts" : [],
    
        "a_arr" : [],
        
        "k_arr" : []
    }
    
    target = dnx.chimera_graph(25, 25, 4)
    
    print(target.number_of_nodes())
    
    for r in range(2):
        for k in range(15, 76, 1):
            v = int(k/4)
            
            
            print("k={} v={} k/v={}".format(k, v, k/v))
            
            ksatInstance = rs.generateRandomKSAT(k, v, 3)
            
            p_qubo = s2q.primitiveQUBO(ksatInstance)

            wmis_qubo = s2q.WMISdictQUBO(ksatInstance)
            
            qg = nx.Graph()
        
            for coupler, energy in wmis_qubo.items():
                if coupler[0] != coupler[1]:
                    qg.add_edge(coupler[0], coupler[1], weight=energy)
                    
            
            ig = nx.Graph()
            for coupler, energy in p_qubo.items():
                if coupler[0] != coupler[1]:
                    ig.add_edge(coupler[0], coupler[1], weight=energy)
            
            qemb = minorminer.find_embedding(qg.edges(), target.edges(), return_overlap=True, threads=8)
            print("wmis emb. found: {}".format(qemb[1]))
            iemb = minorminer.find_embedding(ig.edges(), target.edges(), return_overlap=True, threads=8)
            print("primitive emb. found: {}".format(iemb[1]))
            
            if qemb[1] == 0 or iemb[1] == 0:
                print()
                continue
            
            p_node_count = 0;
            p_conn_count = 0;
            
            used_nodes = []
            for var, chain in iemb[0].items():
                used_nodes.extend(chain)
            
            p_node_count = len(np.unique(used_nodes))
                    
            wmis_node_count = 0;
            wmis_conn_count = 0;
            
            used_nodes = []
            for var, chain in qemb[0].items():
                used_nodes.extend(chain)
            
            wmis_node_count = len(np.unique(used_nodes))
                
            
            #print("p_qubo:    nodes={} connections={}".format(p_node_count, p_conn_count))
            
            #print("wmis_qubo: nodes={} connections={}".format(wmis_node_count, wmis_conn_count))
            
            print("p_qubo nodes= {}".format(p_node_count));
            print("wwmis qubo nodes= {}".format(wmis_node_count));
            
            print("nodes= {}".format(p_node_count / wmis_node_count));
            #print("conns= {}".format(p_conn_count / wmis_conn_count));
            
            
            data["p_node_counts"].append(p_node_count)
            data["wmis_node_counts"].append(wmis_node_count)
            
            data["p_conn_counts"].append(p_conn_count)
            data["wmis_conn_counts"].append(wmis_conn_count)
            
            data["a_arr"].append(k/v)
            data["k_arr"].append(k)
            
            print()

    sns.set()

    ax = sns.scatterplot(x="k_arr", y="p_node_counts", data=data, label="p_qubo")
    ax.set(xlabel='k',
           ylabel='used verticies after embedding')
    ax = sns.scatterplot(x="k_arr", y="wmis_node_counts", data=data, ax=ax, label="wmis_qubo")
    plt.show()

def medianChainLength(emb):
    chl = []
    
    for chain in emb.values():
        chl.append(len(chain))
        
    sns.distplot(chl)
    plt.show()
    
    return np.mean(chl)
    
    
def test2():
    sat = rs.generateRandomKSAT(42, 10, 3)
    
    print(sat.toString())
    
    qubo = s2q.WMISdictQUBO(sat)
    #ising = s2q.primitiveQUBO_8(sat)
    ising = s2q.WMISdictQUBO_2(sat)
    
    qg = nx.Graph()
    
    for coupler, energy in qubo.items():
        if coupler[0] != coupler[1]:
            qg.add_edge(coupler[0], coupler[1], weight=energy)
            
    
    ig = nx.Graph()
    for coupler, energy in ising.items():
        if coupler[0] != coupler[1]:
            ig.add_edge(coupler[0], coupler[1], weight=energy)
            
    #plt.ion()
    
    print(qg.number_of_nodes(), qg.number_of_edges())
    print(ig.number_of_nodes(), ig.number_of_edges())
    
    target = dnx.chimera_graph(16, 16, 4)
    #nx.draw_shell(qg, with_labels=True, node_size=50)
    #nx.draw_shell(ig, with_labels=True, node_size=50)
    
    eg = ig
    
    emb = minorminer.find_embedding(eg.edges(), target.edges(), return_overlap=True,
                                    threads=8)
    
    print(emb[1])
    
    for node, chain in emb[0].items():
        print(node, chain)
    
    print("avrg chain length = {}".format(medianChainLength(emb[0])))
    
    dnx.draw_chimera_embedding(G=target, emb=emb[0], embedded_graph=eg, show_labels=True)
    
    plt.show()
    
    
def test3():
    sat = rs.generateRandomKSAT(42, 10, 3)
    
    print(sat.toString())
    
    ising = s2q.primitiveQUBO(sat)
    
    h = {}
    J = {}
    
    for coupler, energy in ising.items():
        if coupler[0] == coupler[1]:
            h[coupler[0]] = energy
        else:
            J[coupler] = energy
    
    res = QBSolv().sample_ising(h, J, find_max=False)
    
    
    sample = list(res.samples())[0]
    
    extracted = {}
    
    r = re.compile("c\d+_l-?\d*")
    
    for label, assignment in sample.items():
        if r.fullmatch(label):
            
            extracted[tuple(re.split(r"\_l", label[1:]))] = assignment
            
    model = [True for i in range(len(extracted))]
    
    assignments = {}
    
    for label, assignment in extracted.items():
        clause = int(label[0])
        lit = int(label[1])
        var = abs(lit)
        
        if lit < 0:
            assignment *= -1
        
        if var in assignments:
            assignments[var].append(assignment)
        else:
            assignments[var] = [assignment]
            

    conflicts = False
    for var, a in assignments.items():
        if abs(np.sum(a)) != len(a):
            conflicts = True
            print("conflicts - no solution found")
            print(var, np.sort(a))
         
    if conflicts:
        return 
    
    model = [True for i in range(sat.getNumberOfVariables())]
    
    for var, assignment in assignments.items():
        model[var - 1] = True if assignment[0] > 0 else False
    
    print(model, sat.checkAssignment(model))

def test3_3():
    sat = rs.generateRandomKSAT(42, 10, 3)
    
    print(sat.toString())
    
    #ising = s2q.primitiveQUBO_8(sat)
    ising = s2q.WMISdictQUBO_2(sat)
    
    #ising = {}
    
    #ising[("x1", "z1")] = +2
    #ising[("x2", "z2")] = +2
    
    #ising[("z1", "z2")] = +2
    
    #ising[("z1", "z1")] = -2
    #ising[("z2", "z2")] = -2
    
    #ising[("x1", "z3")] = -2
    #ising[("x2", "z3")] = -2
    #ising[("x3", "z3")] = +2
    #ising[("z3", "z3")] = +2    

    
    h, J = graph.split_ising(ising)
    
    #res = QBSolv().sample_ising(h, J, find_max=False)
    res = QBSolv().sample_qubo(ising, find_max=False)
    
    #res = dimod.ExactSolver().sample_ising(h, J)
    #res = dimod.ExactSolver().sample_qubo(ising)
    
    sample = res.first.sample
    
    print(res.truncate(50))
    #print(res.truncate(10))
    
    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]:
                print("conflict at var: {}".format(var))
                conflicts.add(var)
                
    #if conflicts:
    #    return 
    
    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
    
    print()
    
    print(model)
    
    print()
    
    print(sat.checkAssignment(model))
    
    print()
    
    degrees = sat.getDegreesOfVariables()
    
    for var in conflicts:
        node_var = "x{}".format(var)
        node_nvar = "x{}".format(-var)
        print("var {}: deg={}, coupler={}, e={}, ne={}"
              .format(var,
                      degrees[var],
                      ising[(node_var, node_nvar)],
                      assignments[var],
                      assignments[-var]))
              
def test3_4():
    sat_count = 0
    
    steps = 200
    
    for i in tqdm(range(steps)):
        sat = rs.generateRandomKSAT(35, 10, 3)
        
        ising = s2q.WMISdictQUBO_2(sat)
        
        
        
        res = QBSolv().sample_qubo(ising, find_max=False)
        
        #res = dimod.ExactSolver().sample_qubo(ising)
        
        sample = res.first.sample
        
        
        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
        
        if sat.checkAssignment(model):
            sat_count += 1
        
    print("sat percentage: {}".format(sat_count / steps))
    
def test_wmis():
    sat_count = 0
    
    #steps = 100
    steps = 1
    
    for i in tqdm(range(steps)):
        sat = rs.generateRandomKSAT(35, 10, 3)
        target = dnx.chimera_graph(16, 16, 4)
        
        qubo = s2q.WMISdictQUBO(sat)
        
        qg = nx.Graph()
    
        for coupler, energy in qubo.items():
            if coupler[0] != coupler[1]:
                qg.add_edge(coupler[0], coupler[1], weight=energy)
            
        qemb = minorminer.find_embedding(qg.edges(), target.edges(), return_overlap=True)
        
        chimera_sampler = dimod.StructureComposite(SimulatedAnnealingSampler(),
                                                   target.nodes(),
                                                   target.edges())
        
        solver = FixedEmbeddingComposite(chimera_sampler, qemb[0])
        
        res = solver.sample_qubo(__negate_qubo(qubo))
        #res = solver.sample_qubo(qubo)
        
        #res = dimod.ExactSolver().sample_qubo(ising)
        
        
        sample = res.first.sample
        
        model = script.majority_vote_sample(sample)
        
        print("energy={}".format(res.first.energy))
        
        if sat.checkAssignment(model):
            sat_count += 1
        
    print("sat percentage: {}".format(sat_count / steps))
    
def test4():
    sat = rs.generateRandomKSAT(50, 13, 3)
    
    print(sat.toString())
    
    qubo = s2q.WMISdictQUBO(sat)
    ising = s2q.primitiveQUBO(sat)
    
    qg = nx.Graph()
    
    for coupler, energy in qubo.items():
        if coupler[0] != coupler[1]:
            qg.add_edge(coupler[0], coupler[1], weight=energy)
            
    
    ig = nx.Graph()
    for coupler, energy in ising.items():
        if coupler[0] != coupler[1]:
            ig.add_edge(coupler[0], coupler[1], weight=energy)
            
    #plt.ion()
    
    print(qg.number_of_nodes(), qg.number_of_edges())
    print(ig.number_of_nodes(), ig.number_of_edges())
    
    target = dnx.chimera_graph(12, 12, 4)
    #nx.draw_shell(qg, with_labels=True, node_size=50)
    #nx.draw_shell(ig, with_labels=True, node_size=50)
   
    print()
    print("wmis:")
    qemb = minorminer.find_embedding(qg.edges(), target.edges(), return_overlap=True)
    print(qemb[1])
    print("median chain length = {}".format(medianChainLength(qemb[0])))
    
    used_nodes = []
    for var, chain in qemb[0].items():
        used_nodes.extend(chain)
    
    used_nodes = np.unique(used_nodes)
    print("used nodes (embedding) = {}".format(len(used_nodes)))
    
    
    print()
    print("primitves:")
    iemb = minorminer.find_embedding(ig.edges(), target.edges(), return_overlap=True)
    print(iemb[1])
    print("median chain length = {}".format(medianChainLength(iemb[0])))
    
    used_nodes = []
    for var, chain in iemb[0].items():
        used_nodes.extend(chain)
        
    used_nodes = np.unique(used_nodes)
    print("used nodes (embedding) = {}".format(len(used_nodes)))

def __negate_qubo(qubo):
    negative_qubo = {}
    
    for coupler, energy in qubo.items():
        negative_qubo[coupler] = -1 * energy
        
    return  negative_qubo

#test3_4()
#test2()
#test_kv_range
test_wmis()