import dimod
import networkx as nx
import dwave_networkx as dnx
from dwave.embedding import pegasus
import neal, dwave.system
import numpy as np

A_int = np.array([[0,1,2,3,4,1],
              [2,1,4,1,0,2],
              [1,5,2,3,1,1]])
m, n = A_int.shape
vartype = dimod.Vartype.BINARY
c_int = np.array([1,2,1,-3,1,-1])
b = np.array([7,8,5])
print(A_int.dot(np.array([1,0,0,0,1,3])))

def get_binary_form(c,bits):
    c_bin = []
    for c_j in c:
        c_bin += [c_j * 2**i for i in range(0,bits)]
    return np.array(c_bin)

# represent each x_i as 3 bits to the power of 2
bits = 3
# Set Penalty for Constraints aprox solution
P = max(abs(c_int))*3
c = get_binary_form(c_int,bits)
print(c)
A = np.array([get_binary_form(A_int[i], bits) for i in range(0,m)])
print(A)

# Set up linear and quadratic coefficients
L = c.T - P * 2 * np.matmul(b.T,A)
Q = P * np.matmul(A.T,A)
offset = P * b.T.dot(b)
bqm = dimod.as_bqm(L, Q, offset, vartype)

def construct_x(sample, bits):
    x=[]
    for k in range(0,len(sample.keys()),bits):
        x += [sum(sample[k+j]* 2**j for j in range(0,bits))]
    return np.array(x)

def testSamples(sampleset, bits):
    print("Minimum Energy Level found in:",sampleset.first)
    X = []
    for sample, energy,  in sampleset.data(fields=['sample', 'energy']):
        x = construct_x(sample,bits)
        if all(A_int.dot(x)[i] == b[i] for i in range(0,m)) :
            X += [(x,energy)]
    return X

# Define target Pegasus Graph

M = 16 # Defines number of qubits in Pegasus graph : 24M(M-1)
P_M = dnx.pegasus_graph(M, data = True) # allow the use of pegasus indizes

# Define source complete graph

a = len(bqm.variables)
K_a = nx.complete_graph(a)

# Embedding K_a into P_M
embedding = pegasus.find_clique_embedding(K_a, target_graph=P_M)

sampler = dwave.system.FixedEmbeddingComposite(dwave.system.DWaveSampler(), embedding = embedding)

sampler.child.properties['num_qubits']

sampleset = sampler.sample(bqm,num_reads = 2000)

sampleset.to_pandas_dataframe()

solutions_real = testSamples(sampleset, bits)
solutions_unique = { str(x) : energy for x, energy in solutions_real}
for x in solutions_unique.keys() :
    print("Found solution {} with Energy Level {}".format(x,solutions_unique[x]))

a = len(bqm.variables)
embedding = pegasus.find_clique_embedding(a,m = a)

sampler = dwave.system.FixedEmbeddingComposite(dwave.system.DWaveSampler(), embedding = embedding)

