import dimod
import neal
import numpy as np
from dwave.system import DWaveSampler, EmbeddingComposite, DWaveCliqueSampler

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)

exactSolver = dimod.ExactSolver()
sim = neal.SimulatedAnnealingSampler()
num_reads = 2500
exactSet = exactSolver.sample(bqm)
simSet = sim.sample(bqm, num_reads = num_reads)

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

solutions_exact = testSamples(exactSet, bits)
for x, energy in solutions_exact :
    print("Found solution {} with Energy Level {}".format(x,energy))
        

solutions_neal = testSamples(simSet, bits)
solutions_unique = { str(x) : energy for x, energy in solutions_neal}
solutions_count = {x:0 for x in solutions_unique.keys()}
for x,energy in solutions_neal:
    solutions_count[str(x)]+=1
for x in solutions_unique.keys() :
    print("Found solution {} {} times with Energy Level {}".format(x,solutions_count[x],solutions_unique[x]))

samplerComposite = EmbeddingComposite(DWaveSampler())
sampleset = samplerComposite.sample(bqm,num_reads = num_reads) # a7cb9757-6989-43d2-9e06-914c9e947727

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

samplerClique = DWaveCliqueSampler()
sampleset = samplerClique.sample(bqm,num_reads = num_reads) # ecebd686-0b4e-47c2-9087-b5e234d79d89

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

