#!/usr/bin/env python3

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

import networkx as nx
import dwave_networkx as dnx
import minorminer
from dwave_qbsolv import QBSolv
import dimod

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


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

def test_kv_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(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.primitiveQUBO_5(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 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)))
    

#test3_3()
test2()
#test_kv_range()