import copy
from contextlib import suppress
from heapq import heappush, heappop
import matplotlib.pyplot as plt

# return a board that is like the board b, but has domains for each element of b (always 1-9)
def genDomains(b):
    for row in range(0, 9):
        for cell in range(0, 9):
            if (b[row][cell] == 0):
                b[row][cell] = list(range(1, 10))
    return b
    
# returns True if value is valid
def valid(brd, row, col, val):
    # check row
    if (val in brd[row]):
        return False

    # check column
    for i in range(0, 9):
        if (brd[i][col] == val):
            return False
    
    # check "box"
    rownum = int(row / 3)
    colnum = int(col / 3)
    
    for i in range(rownum * 3, rownum * 3 + 3):
        for j in range(colnum * 3, colnum * 3 + 3):
            if (brd[i][j] == val):
                return False
    
    return True

# naive backtracking solver
def naive(start):
    working = copy.deepcopy(start) # this is only "filled in values" and 0s
    solution = genDomains(start)
    # unassigned will be a list of positions we have to fill
    unassigned = []
    for i in range(0, 9):
        for j in range(0, 9):
            if (isinstance(solution[i][j], list)):
                unassigned.append((i, j))
                
    assumptions = []
                    
    if(len(unassigned) == 0):
        return (working, 0)
        
    # count assignments
    count = 0
    
    # while there are unassigned vars, keep going
    while(len(unassigned)):
        index = unassigned[-1]
        success = False
        
        # iterate over all values in the domain list
        while solution[index[0]][index[1]]:
            i = solution[index[0]][index[1]].pop()
            count += 1
            # took too long
            if (count >= 10000):
                print("took too long")
                return 10000

            # check if this part of the domain(solution) is valid
            if (valid(working, index[0], index[1], i)):
                #count += 1
                #if (count >= 10000):
                #  print("took too long")
                #  return False
                solution[index[0]][index[1]].append(i) # keep in the domain
                working[index[0]][index[1]] = i
                assumptions.append(index)
                unassigned.pop()
                success = True
                break
        
        if (success):
            continue
        else:
            # restore domain to full since we failed
            solution[index[0]][index[1]] = list(range(1, 10))
            working[index[0]][index[1]] = 0
            lastdex = assumptions.pop()
            solution[lastdex[0]][lastdex[1]].remove(working[lastdex[0]][lastdex[1]])
            working[lastdex[0]][lastdex[1]] = 0
            unassigned.append(lastdex)
    
    
    # if we exit without assigning everything, we should have failed
    if (unassigned): return 10000

    return count


# returns a board (domains) where inferences are made for the cell at row, col
def infer(solutions, brd, row, col, val):
    domains = copy.deepcopy(solutions)
    # remove from same row & col
    for i in range(0, 9):
        if (val in domains[row][i] and i != col):
            domains[row][i].remove(val)
        if (val in domains[i][col] and i != row):
            domains[i][col].remove(val)
            
    # remove for "box"
    rownum = int(row / 3)
    colnum = int(col / 3)
    
    for i in range(rownum * 3, rownum * 3 + 3):
        for j in range(colnum * 3, colnum * 3 + 3):
            if (val in domains[i][j] and (i != row and j != col)):
                domains[i][j].remove(val)
                
    return domains


# generates domains in a format supporting forward checking
def gen2Domains(b):
    for row in range(0, 9):
        for cell in range(0, 9):
            if (b[row][cell] == 0):
                b[row][cell] = list(range(1, 10))
            else:
                b[row][cell] = [b[row][cell]]
    return b


# recursive solver for forward-checking
def solve(working, domains, unassigned, count):
    if (not unassigned):
        return (True, count)

    index = unassigned.pop()
    
    # for every value in the domain, check if using it works. if all fail, backtrack.
    for i in domains[index[0]][index[1]]:
        working[index[0]][index[1]] = i
        newdomains = infer(domains, working, index[0], index[1], i)
        domains[index[0]][index[1]].remove(i)
        count += 1
        # took too long
        if (count >= 10000):
          print("took too long")
          return (True, 10000)

        # check for invalidated nodes (empty domains)
        flag = True
        result = False
        for i in range(0, 9):
          for j in range(0, 9):
            if (len(newdomains[i][j]) <= 0):
              flag = False
        
        if (flag): result = solve(working, newdomains, copy.deepcopy(unassigned), count)
        if (not result):
          return (False, count)
        if (result[0]):
            return result
        else:
            #domains[index[0]][index[1]].remove(i)
            count = result[1]
    
    return (False, count)

# forward checking solver
def forward(start):
    working = copy.deepcopy(start) # this is only "filled in values" and 0s
    domains = gen2Domains(working)
    # unassigned will be a list of positions we have to fill
    unassigned = []
    for i in range(0, 9):
        for j in range(0, 9):
            if (len(domains[i][j]) == 9):
                unassigned.append((i, j))
   
    # forward-checking on pre-assigned values
    for i in range(0, 9):
        for j in range(0, 9):
            if (working[i][j] != 0):
                domains = infer(domains, working, i, j, working[i][j])
            
    result = solve(working, domains, unassigned, 0)
    return result[1]
    #if (result[0]): return result[1]
    #else: return 10000


# returns size of domain for a given index
def domsize(domains, index):
    return (len(domains[index[0]][index[1]]))


# returns the # of 0s that are in the same row, col, or box as index
def related(brd, index):
    count = 0
    # count 0s in row + col
    for i in range(0, 9):
        if (brd[index[0]][i] == 0 and i != index[1]):
            ++count
        if (brd[i][index[1]] == 0 and i != index[0]):
            ++count
            
    # count for "box" as well
    rownum = int(index[0] / 3)
    colnum = int(index[1] / 3)
    
    for i in range(rownum * 3, rownum * 3 + 3):
        for j in range(colnum * 3, colnum * 3 + 3):
            if (brd[i][j] == 0 and (i != index[0] and j != index[1])):
                ++count
                
    return count


# returns the # of constraints that will follow from assigning index with val
def lcv(solutions, index, val):
    count = 0
    # count 0s in row + col
    for i in range(0, 9):
        if (val in solutions[index[0]][i] and i != index[1]):
            ++count
        if (val in solutions[i][index[1]] and i != index[0]):
            ++count
            
    # count for "box" as well
    rownum = int(index[0] / 3)
    colnum = int(index[1] / 3)
    
    for i in range(rownum * 3, rownum * 3 + 3):
        for j in range(colnum * 3, colnum * 3 + 3):
            if (val in solutions[i][j] and (i != index[0] and j != index[1])):
                ++count
                
    return count
    
    
# return the correct node + val to try
def genVal(domains, working, unassigned):
    # used to track intermediary values
    heap = []
    superheap = []
    bestrating = 1.0
    
    # get the best indices according to domain size
    for i in unassigned:
        rating = domsize(domains, i) / 9.0
        if (rating < bestrating):
            bestrating = rating
            heap = [i]
        elif (rating == bestrating):
            heap.append(i)
    
    # get the best indices according to degree(related cells)
    bestrating = 1
    for i in heap:
        rating = related(working, i) / 27.0
        if (rating < bestrating):
            bestrating = rating
            superheap = [i]
        elif (rating == bestrating):
            superheap.append(i)
    
    index = superheap[0]
    bestrating = 27
    val = working[index[0]][index[1]]
    
    # get best values according to LCV
    for i in domains[index[0]][index[1]]:
        rating = lcv(domains, index, i)
        if (rating <= bestrating):
            bestrating = rating
            val = i
            
    return (index, val)


# recursive solver that uses heuristics to decide what node to explore
def solveh(working, domains, unassigned, count):
    if (not unassigned):
        return (True, count)
    
    # while there are unassigned values keep trying
    while(unassigned):
        # get next value using heuristics, remove this node from assigned
        nextThing = genVal(domains, working, unassigned)
        index = nextThing[0]
        val = nextThing[1]
        working[index[0]][index[1]] = val
        unassigned.remove(index)
        
        # check for invalidated nodes (empty domain)
        flag = True
        result = False
        newdomains = infer(domains, working, index[0], index[1], val)
        for i in range(0, 9):
            for j in range(0, 9):
                if (not newdomains[i][j]):
                    flag = False
        
        count += 1
        # took too long
        if (count >= 10000):
          print("took too long")
          return (False, count)

        # success! recurse
        if (flag): result = solveh(working, newdomains, copy.deepcopy(unassigned), count)
        if (not result): pass
        elif (result[0]):
            return result
        elif (len(domains[index[0]][index[1]]) > 1): # remove from domain, keep going
            working[index[0]][index[1]] = 0
            domains[index[0]][index[1]].remove(val)
            unassigned.append(index)
            if (flag): count = result[1]
        else: # no values worked :( return false
            if (flag): return (False, result[1])
            return (False, count)


# forward checking solver with heuristics
def heuristic(start):
    working = copy.deepcopy(start) # this is only "filled in values" and 0s
    domains = gen2Domains(start)
    # unassigned will be a list of positions we have to fill
    unassigned = []
    for i in range(0, 9):
        for j in range(0, 9):
            if (len(domains[i][j]) == 9):
                unassigned.append((i, j))
    
    # initial inferences
    for i in range(0, 9):
        for j in range(0, 9):
            if (working[i][j] != 0):
                domains = infer(domains, working, i, j, working[i][j])
            
    result = solveh(working, domains, unassigned, 0)
    if (result[0]): return result[1]
    else: return 10000


def main():
    plt.ioff()
    plt.switch_backend('agg')
    averages = []
    bverages = []
    cverages = []
    
    for i in range(1, 72):
        avgA = 0
        avgB = 0
        avgC = 0
        for j in range(1, 11):
            filepath = "sudoku_problems/" + str(i) + "/" + str(j) + ".sd"
            
            # import board
            with open(filepath) as file:
                board = file.read().splitlines()
                board = board[:-1]

            # convert to list of list of ints
            for l in board:
                board[board.index(l)] = list(map(lambda x: int(x), l.split()))
            
            avgA += naive(copy.deepcopy(board));print(i, j)
            avgB += forward(copy.deepcopy(board));print(i, j)
            avgC += heuristic(copy.deepcopy(board));print(i, j)
            
        averages.append(avgA / 10.0)
        bverages.append(avgB / 10.0)
        cverages.append(avgC / 10.0)
        
    figure, axes = plt.subplots(1, 1, True)
    axes.plot(range(1, 72), averages, label='Naive Algorithm')
    axes.plot(range(1, 72), bverages, label='Forward-Checking Algorithm')
    axes.plot(range(1, 72), cverages, label='Heuristics')
    axes.legend()
    plt.xlabel("Number of Initial Valued Filled In")
    plt.ylabel("Average Number of Variable Assignments in 10 Runs")
    plt.savefig("graph.pdf")
    
main()