import pickle


class SumProductPuzzle:
    def __init__(self, lb=2, rb=98):
        self.nodes = [(i, j) for i in range(lb, rb) for j in range(lb, rb) if i < j]
        edgesP = [(tup1, tup2, tup1[0] * tup1[1]) for tup1 in self.nodes for tup2 in self.nodes
                  if (tup1[0] * tup1[1] == tup2[0] * tup2[1]
                      and tup1[0] <= tup2[0])]
        edgesS = [(tup1, tup2, sum(tup1)) for tup1 in self.nodes for tup2 in self.nodes
                  if (tup1[0] + tup1[1] == tup2[0] + tup2[1]
                      and tup1[0] <= tup2[0])]

        self.primes = []
        g = gen_primes()
        p = next(g)
        while p < rb:
            self.primes.append(p)
            p = next(g)

        self.adjListP = {}
        self.adjListS = {}
        createAdjList(edgesP, self.adjListP)
        createAdjList(edgesS, self.adjListS)

    def writeAsPrime(self, num):
        for i in self.primes:
            for j in self.primes:
                if i + j == num or i + i ** 2 == num:
                    return True
        return False

    def cond1(self, u):
        # validates first condition
        for v in self.adjListS[u]:
            if len(self.adjListP[v]) == 1:
                return False
        return True

    def cond2(self, u):
        # validates second condition
        for v in self.adjListP[u]:
            if u != v and not self.writeAsPrime(sum(v)):
                return False
        return True

    def cond3(self, u):
        # validates third condition
        for v in self.adjListS[u]:
            if v != u and self.cond2(v):
                return False
        return True

    def getSolution(self):
        # returns puzzle solution
        return [node for node in self.nodes if self.cond3(node) and self.cond2(node) and self.cond1(node)]


# Sieve of Eratosthenes
# Code by David Eppstein, UC Irvine, 28 Feb 2002
# http://code.activestate.com/recipes/117119/

def gen_primes():
    """ Generate an infinite sequence of prime numbers.
    """
    # Maps composites to primes witnessing their compositeness.
    # This is memory efficient, as the sieve is not "run forward"
    # indefinitely, but only as long as required by the current
    # number being tested.
    #
    D = {}

    # The running integer that's checked for primeness
    q = 2

    while True:
        if q not in D:
            # q is a new prime.
            # Yield it and mark its first multiple that isn't
            # already marked in previous iterations
            #
            yield q
            D[q * q] = [q]
        else:
            # q is composite. D[q] is the list of primes that
            # divide it. Since we've reached q, we no longer
            # need it in the map, but we'll mark the next
            # multiples of its witnesses to prepare for larger
            # numbers
            #
            for p in D[q]:
                D.setdefault(p + q, []).append(p)
            del D[q]

        q += 1


def createAdjList(edgeList, adjList):
    # creates adjacency lists as a dictionary
    for i in edgeList:
        if not adjList.get(i[0]):
            adjList[i[0]] = list([i[1]])
        elif i[1] not in adjList[i[0]]:
            adjList[i[0]].append(i[1])
        if not adjList.get(i[1]):
            adjList[i[1]] = list([i[0]])
        elif i[0] not in adjList[i[1]]:
            adjList[i[1]].append(i[0])


def dumpObject(obj, filename="sumProductObj"):
    # dumps an object in a file
    with open(filename, "wb") as out:
        pickle.dump(obj, out, pickle.HIGHEST_PROTOCOL)


if __name__ == "__main__":
    try:
        with open("sumProductObj", "rb") as obj:
            testInstance = pickle.load(obj)
    except FileNotFoundError as e:
        print("File Not Found! Creating new puzzle instance!")
        testInstance = SumProductPuzzle()
        dumpObject(testInstance)

    print("Puzzle's solution between %d nodes is %s"
          % (len(testInstance.nodes), testInstance.getSolution()))