from copy import deepcopy
from random import choice, randrange
import pandas as pd
import numpy as np
import sys

class cell():
    def __init__(self, row, col):
        self.possible = set([i for i in range(1,10)])
        self.solved = False
        self.solution = 0
        self.prefilled = False
        self.branch = False
        self.row = row
        self.col = col
    
    def __str__(self):
        retstr = 'solved ' if self.solved else 'unsolved '
        retstr += f'sudoku_cell @ ({self.row},{self.col}) '
        if self.solved:
            retstr += f'with value {self.solution}'
        else:
            retstr += f'with possible values {self.possible}'
        return retstr
    
    def set_value(self, value, branch=False):
        self.possible = {value}
        self.solution = value
        self.solved = True
        self.prefilled = not branch
        self.branch = branch
    
    def cleanup(self):
        self.branch = False
    
    def remove_value(self, value):
        if not self.solved:
            self.possible.discard(value)
            if len(self.possible) == 1:
                self.solution = list(self.possible)[0]
                self.solved = True
        pass # for debugging

    def collapse(self, values):
        if not self.possible:
            return False
        a = self.possible - self.possible.intersection(values)
        if not a:
            return False
        self.possible = a
        if len(self.possible) == 1:
            self.solution = list(self.possible)[0]
            self.solved = True
        return True
    
    def __len__(self):
        return len(self.possible)
    
    @property
    def values(self):
        return list(self.possible)

class sudoku_grid():
    def __init__(self, prefilled_cells):
        self.grid = [[cell(i,j) for j in range(9)] for i in range(9)]
        self.branches = list()
        self.branch_points = list()
        self.last_hash = None
        self.iteration = 0

        # prefilled_cells is a nested list (9x9) of values (1-9), 0 specifies an empty cell
        try:
            ctr = 0
            if len(prefilled_cells) != 9:
                raise ValueError.add_note(f'wrong number of rows')
            for i in range(9):
                if len(prefilled_cells[i]) != 9:
                    raise ValueError.add_note(f'wrong number of cells in row {i}')
                for j in range(9):
                    # if prefilled_cell in valid range: fill in matching cell an mark as solved
                    if 1 <= prefilled_cells[i][j] <= 9:
                        self.grid[i][j].set_value(prefilled_cells[i][j])
                        ctr += 1
                    
            if ctr == 0:
                raise ValueError.add_note(f'prefilled_cells is empty')
        except ValueError as e:
            print(f'{e}')

    def __str__(self):
        # retstr = '                Sudoku\n'
        retstr = ''
        retstr += '    0   1   2   3   4   5   6   7   8\n'
        retstr += '  ╔═══╤═══╤═══╦═══╤═══╤═══╦═══╤═══╤═══╗\n'
        for i in range(9):
            for j in range(9):
                if j == 0:
                    retstr += f'{i} ║ '
                if self.grid[i][j].solved:
                    current_val = self.grid[i][j].solution
                    if self.grid[i][j].prefilled:
                        retstr += f'{current_val}'
                    elif self.grid[i][j].branch:
                        retstr += f'\033[91m{current_val}\033[00m'
                    else:
                        retstr += f'\033[92m{current_val}\033[00m'
                else:
                    retstr += ' '
                if j == 2 or j == 5 or j == 8:
                    retstr += ' ║ '
                else:
                    retstr += ' │ '
            if i == 2 or i == 5:
                retstr += '\n  ╠═══╪═══╪═══╬═══╪═══╪═══╬═══╪═══╪═══╣\n'
            elif i == 8:
                pass
            else: 
                retstr += '\n  ╟───┼───┼───╫───┼───┼───╫───┼───┼───╢\n'
        retstr += '\n  ╚═══╧═══╧═══╩═══╧═══╧═══╩═══╧═══╧═══╝'
        return retstr

    def iterate(self):
        try:
            self.iterate1()
            self.iterate2()
            self.check_branching()
            self.solvable()
            self.iteration += 1
        except:
            print('No solution found')
            print(self)
            sys.exit()

    # remove values based on solved cells
    def iterate1(self):
        # iterate over all cells
        for i in range(9):
            for j in range(9):
                # # remove possible values based on solved cells
                current_value = self.grid[i][j].solution
                if current_value:
                    for k in range(9):
                        current_cell = self.grid[i][k]
                        current_cell.remove_value(current_value) # remove value from current row
                        current_cell = self.grid[k][j]
                        current_cell.remove_value(current_value) # remove value from current column
                    for k in range(3):
                        for l in range(3):
                            current_cell = self.grid[(i//3)*3+(i+k)%3][(j//3)*3+(j+l)%3]
                            current_cell.remove_value(current_value) # remove value from current 3x3 box

    # remove values based on possible values in other cells of row/column/square
    def iterate2(self):
        for i in range(9):
            for j in range(9):        
                current_cell = self.grid[i][j]
                if current_cell.solved:
                    continue

                current_set = set()
                for k in range(8):
                    if len(current_set) == 9:
                        continue
                    current_cell2 = self.grid[i][(j+k+1)%9]
                    current_set = current_set.union(current_cell2.possible)
                current_cell.collapse(current_set)

                current_set = set()
                for k in range(8):
                    if len(current_set) == 9:
                        continue
                    current_cell2 = self.grid[(i+k+1)%9][j]
                    current_set = current_set.union(current_cell2.possible)
                current_cell.collapse(current_set)

                add = 0
                current_set = set()
                for k in range(8):
                    if len(current_set) == 9:
                        continue
                    if (i//3)*3+k//3 == i and (j//3)*3+k%3 == j:
                        add = 1
                    l = k + add
                    row = (i//3)*3+l//3
                    col = (j//3)*3+l%3
                    current_cell2 = self.grid[row][col]
                    current_set = current_set.union(current_cell2.possible)
                current_cell.collapse(current_set)

    def cleanup(self):
        for i in range(9):
            for j in range(9):
                self.grid[i][j].cleanup()

    def check_branching(self):
        if self.invalid_states():
            self.rollback_branch()
        new_hash = hash(str(self))
        if new_hash == self.last_hash:
            [row, col, e, sols] = self.find_lowest_entropy()
            self.branch(row, col, sols)
        self.last_hash = new_hash

    # copy current solution and branch point (branch_point: row, column, selected_solution)
    def save_branch(self, branch_point):
        self.branches.append(deepcopy(self.grid))
        self.branch_points.append(branch_point)
    
    # revert latest branch
    def rollback_branch(self):
        solution = self.branches.pop()
        branch_point = self.branch_points.pop()
        self.grid = solution
        current_cell = self.grid[branch_point[0]][branch_point[1]]
        current_cell.remove_value(branch_point[2])
        print('rollback branch')

    def branch(self, row, col, solutions):
        current_cell = self.grid[row][col]
        selected_value = choice(solutions)
        self.save_branch((row, col, selected_value))
        print(f'Branched: ({row},{col}): {current_cell.possible}, chose {selected_value}')
        current_cell.set_value(selected_value, branch=True)
    
    def find_lowest_entropy(self):
        lowest_i = -1
        lowest_j = -1
        sols = None
        lowest_e = 10
        for i in range(9):
            for j in range(9):
                current_cell = self.grid[i][j]
                if current_cell.solved:
                    continue
                e = len(current_cell)
                if e < lowest_e:
                    sols = current_cell.values
                    lowest_i = i
                    lowest_j = j
                    lowest_e = e
                    if lowest_e == 0:
                        lowest_i = -1
                        lowest_j = -1
                        sols = None
                        lowest_e = 10
                        return (lowest_i, lowest_j, lowest_e, sols)
        return (lowest_i, lowest_j, lowest_e, sols)
    
    def collapse_cell(self, row, col):
        if self.grid[row][col].solved:
            return None
        possible = self.grid[row][col].possible
        if len(possible) == 1:
            self.grid[row][col].solved = list(possible)[0]   
    
    def single_solutions_exist(self):
        for i in range(9):
            for j in range(9):
                if self.grid[i][j].possible:
                    if len(self.grid[i][j].possible) == 1:
                        return True
        return False

    def is_solved(self):
        for i in range(9):
            for j in range(9):
                if not self.grid[i][j].solved:
                    return False
        return True
    
    def solvable(self):
        if len(self.branches) == 0:
            for i in range(9):
                for j in range(9):
                    if len(self.grid[i][j]) == 0:
                        raise Exception
        return True

    def invalid_states(self):
        for i in range(9):
            row_list = list()
            col_list = list()
            for j in range(9):
                row_list.append(self.grid[i][j].solution)
                col_list.append(self.grid[j][i].solution)
                if len(self.grid[i][j]) == 0:
                    return True
            row_list.sort()
            col_list.sort()
            row_list = [x for x in row_list if x != 0]
            col_list = [x for x in col_list if x != 0]
            row_comp = list(set(row_list))
            col_comp = list(set(col_list))
            row_comp.sort()
            col_comp.sort()
            if row_comp != row_list:
                return True
            if col_comp != col_list:
                return True
        return False

if __name__ == '__main__':
    # colorama.init()
    # prefilled = [   [0,4,9,  7,0,5,  0,0,0],
    #                 [0,0,0,  0,0,4,  0,0,3],
    #                 [6,0,1,  2,0,0,  0,7,0],

    #                 [0,0,0,  0,9,1,  0,0,5],
    #                 [0,2,4,  0,6,8,  7,3,1],
    #                 [1,5,8,  0,2,7,  4,9,0],

    #                 [0,0,0,  0,0,2,  6,4,0],
    #                 [0,6,0,  1,0,0,  0,0,0],
    #                 [4,0,5,  0,0,0,  3,0,2]]

    # prefilled = [   [0,0,7,  0,0,5,  0,0,3],
    #                 [0,0,9,  0,6,0,  0,0,0],
    #                 [3,6,0,  0,0,8,  2,0,0],

    #                 [0,0,6,  0,0,0,  0,0,0],
    #                 [5,1,0,  0,8,0,  0,0,9],
    #                 [0,0,0,  0,0,2,  0,4,0],

    #                 [0,0,0,  5,0,0,  9,0,0],
    #                 [8,3,0,  0,1,0,  0,0,5],
    #                 [7,0,0,  0,0,0,  0,0,0]]

    prefilled = [   [2,0,4,  0,6,1,  0,0,9],
                    [0,1,0,  0,0,4,  0,0,0],
                    [0,7,0,  0,0,0,  0,2,0],

                    [0,2,0,  0,0,0,  0,0,0],
                    [8,0,3,  0,0,7,  0,9,0],
                    [0,0,0,  5,0,0,  0,0,6],

                    [4,0,9,  0,0,3,  0,8,0],
                    [0,0,1,  0,0,0,  0,0,0],
                    [0,0,0,  0,7,0,  3,0,0]]

    df = pd.read_csv('sudoku.csv')
    row = randrange(len(df))
    quiz = df.loc[row]['quizzes']
    # print(quiz)
    quiz = list(quiz)
    quiz = [int(x) for x in quiz]
    quiz = np.reshape(quiz, (9, 9)).tolist()

    # sudoku = sudoku_grid(prefilled)
    sudoku = sudoku_grid(quiz)
    print(sudoku)
    while not sudoku.is_solved():
        sudoku.iterate()
        # print(f'Iteration {sudoku.iteration}')
        # print(sudoku)
    sudoku.cleanup()
    print('Solved!')
    print(sudoku)