import dwavebinarycsp
import dimod
import neal
import numpy as np
from dwavebinarycsp.core.constraint import Constraint
from dwave.system import DWaveSampler, EmbeddingComposite

A = np.array([[1,1,1,1,1],
              [2,2,2,2,2],
              [0,1,2,3,4],
              [2,1,4,1,0],
              [0,1,0,1,0]])
b = np.array([2, 4, 3, 5, 1])
solution = {'x0':0, 'x1':1,'x2':1,'x3':0,'x4':0 }
no_solution = {'x0':0, 'x1':1,'x2':1,'x3':1,'x4':1 }

m, n = A.shape
vartype = dimod.Vartype.BINARY
csp = dwavebinarycsp.ConstraintSatisfactionProblem(vartype)
x = tuple(['x' + str(j) for j in range(n)])

def make_f(i): # avoid late binding of i in for loop
    return lambda *z : (sum(A[i,j]*z[j] for j in range(n)) == b[i])

for v in x:
    csp.add_variable(v)
for i in range(m):
    constraint = Constraint.from_func(make_f(i),x,vartype)
    csp.add_constraint(constraint)

def printAndGetTestResults(sampler, bqm):
    sampleset = sampler.sample(bqm)
    print("Sampling by ", sampler)
    for sample, energy in sampleset.data(['sample', 'energy']):
        isValid = csp.check(sample)
        print(sample, isValid, energy)
        if not isValid : return sampleset

bqm = dwavebinarycsp.stitch(csp)
samplers = {"BruteForce" : dimod.ExactSolver() , "Neal" : neal.SimulatedAnnealingSampler()}
samplesets = {s: printAndGetTestResults(samplers[s],bqm) for s in samplers.keys()}

sampler = EmbeddingComposite(DWaveSampler())
sampleset = sampler.sample(bqm,num_reads = 100)

sampleset.to_pandas_dataframe()





