import copy
import math
import random
import sys

from random import randint, shuffle

import numpy as np

def printBoard(board):
    for i in range(len(board)):
        if (i % 3 == 0 and i != 0):
            print("------+------+------")
        for j in range(len(board[i])):
            if (j % 3 == 0 and j != 0):
                sys.stdout.write("|")
            if(isWrong(board,i,j)):
                sys.stdout.write(printRed(board[i][j]) + " ")
            else:
                sys.stdout.write(str(board[i][j]) + " ")

        print("")
    print("\n")

def isWrong(board, i, j):
    x = board[i][j]
    for y in range(9):
        if y == i:
            continue
        if x == board[y][j]:
            return True
    return False



def fitness(board):
    fitness = 0
    for i in range(0, 9, 3):
        for j in range(0, 9, 3):
            square = np.array([board[i][j], board[i+1][j], board[i+2][j], board[i][j+1], board[i][j+2], board[i+1][j+1], board[i+1][j+2], board[i+2][j+1], board[i+2][j+2]])
            fitness -= np.sum([len(set(square)) != 9])

    #for i in range(9):
      #  print(board[i, :])
    return  fitness + -np.sum([len(set(board[:, i])) != 9 for i in range(9)])



def initSodoku():
    row = np.arange(1,10)

    board = np.tile(row, (9, 1))

    return board

def shuffleBoard(board):
    for i in range(9):
        shuffle(board[i])

def step(board):

    x = randint(0, 8)

    y1 = randint(0,8)
    y2 = randint(0,8)
    while y2 == y1:
        y2 = randint(0, 8)

    board[x][y1], board[x][y2] = board[x][y2], board[x][y1]

def printRed(text):
    return format(f"\033[1;31m{text}\033[0;0m")


def __main__():
    board = initSodoku()
    iteration = 0
    initialTemp = np.float64(100)
    temp = initialTemp

   # shuffleBoard(board)
    newBoard = copy.deepcopy(board)
    lastFitness = fitness(board)
    print(f"Fitness: {lastFitness}")
    printBoard(board)
    while lastFitness < 0 and temp > 0:
        step(newBoard)

        newFitness = fitness(newBoard)
        if(newFitness == 0):
            break
        if(newFitness >= lastFitness):
            lastFitness = newFitness
            board = copy.deepcopy(newBoard)
        else:
            p = math.exp(-(abs(newFitness - lastFitness) / temp))
            print(f"Probability: {p*100}%")

            if(p > random.random()):
                # print(f"Returned to less fit state: {lastFitness}")
                lastFitness = newFitness
                board = copy.deepcopy(newBoard)
            else:
                newBoard = copy.deepcopy(board)
        iteration += 1
        temp *= 0.9999


    print(f"Iteration {iteration}")
    printBoard(newBoard)
    print(fitness(newBoard))


__main__()