import math
import numpy as np
from copy import deepcopy

from .csp import Variable, Factor, ConstraintSatisfactionProblem
from .csp.util import aslist

SudokuPuzzle = list[list[int]]  # type alias for a Sudoku puzzle template

SUDOKU_TEMPLATES:dict[str, dict[str, SudokuPuzzle]] = {}
SUDOKU_TEMPLATES['SIMPLE'] = {
    '0': [[0, 3, 4, 0],
          [4, 0, 0, 2],
          [1, 0, 0, 3],
          [0, 2, 1, 0]],
    '1': [[0, 0, 4, 0],
          [4, 0, 3, 0],
          [0, 4, 0, 3],
          [0, 1, 0, 0]],
    '2': [[0, 0, 1, 0],
          [4, 0, 0, 0],
          [0, 0, 0, 2],
          [0, 3, 0, 0]],
    '3': [[2, 0, 0, 0],
          [0, 0, 3, 0],
          [0, 4, 0, 0],
          [0, 0, 0, 1]],
    '4': [[1, 0, 4, 0],
          [0, 0, 0, 0],
          [0, 0, 0, 0],
          [0, 1, 0, 2]],
    '5': [[0, 0, 0, 3],
          [3, 2, 4, 0],
          [0, 4, 3, 2],
          [2, 0, 0, 0]],
    '6': [[3, 4, 1, 0],
          [0, 2, 0, 0],
          [0, 0, 2, 0],
          [0, 1, 4, 3]],
    '7': [[0, 1, 3, 0],
          [2, 0, 0, 0],
          [0, 0, 0, 3],
          [0, 2, 1, 0]],
    '8': [[0, 0, 1, 0],
          [4, 0, 0, 0],
          [0, 0, 0, 2],
          [0, 3, 0, 0]]    
}

SUDOKU_TEMPLATES['EASY'] = {
      '0': [[3, 0, 5, 0, 0, 9, 0, 0, 2],
            [7, 0, 0, 8, 0, 5, 1, 9, 0],
            [0, 1, 9, 4, 7, 0, 0, 3, 0],
          [1, 0, 6, 0, 2, 4, 0, 0, 3],
          [0, 0, 8, 3, 5, 7, 0, 1, 9],
          [9, 5, 3, 0, 0, 0, 2, 7, 0],
          [0, 9, 1, 2, 4, 0, 3, 0, 0],
          [0, 3, 0, 7, 0, 6, 9, 0, 5],
          [2, 6, 0, 0, 9, 0, 4, 8, 0]]

}

SUDOKU_TEMPLATES['HARD'] = {
      '0': [[0, 0, 1, 8, 0, 0, 0, 0, 0],
            [6, 7, 0, 2, 0, 3, 0, 1, 0],
            [2, 0, 5, 7, 0, 0, 6, 0, 3],
            [3, 5, 0, 6, 2, 0, 8, 0, 0],
            [7, 6, 2, 0, 5, 8, 3, 0, 0],
            [0, 0, 0, 0, 0, 4, 0, 5, 0],
            [0, 9, 0, 5, 8, 6, 0, 0, 0],
            [5, 0, 0, 9, 0, 0, 7, 0, 0],
            [0, 2, 6, 4, 3, 0, 5, 9, 0]],
      '1': [[0, 0, 0, 9, 0, 0, 6, 7, 2],
            [0, 0, 2, 0, 0, 1, 0, 4, 3],
            [0, 3, 0, 0, 2, 6, 0, 0, 0],
            [4, 0, 0, 0, 0, 2, 0, 0, 9],
            [0, 0, 7, 3, 0, 9, 8, 0, 0],
            [5, 0, 0, 6, 0, 0, 0, 0, 1],
            [0, 0, 0, 0, 6, 0, 0, 9, 0],
            [0, 1, 0, 4, 0, 0, 3, 0, 0],
            [0, 7, 8, 0, 0, 5, 0, 0, 0]]
}

SUDOKU_SOLUTIONS = {}
SUDOKU_SOLUTIONS['SIMPLE'] = {}
SUDOKU_SOLUTIONS['EASY'] = {
      '0': [[3, 8, 5, 1, 6, 9, 7, 4, 2],
            [7, 4, 2, 8, 3, 5, 1, 9, 6],
            [6, 1, 9, 4, 7, 2, 5, 3, 8],
            [1, 7, 6, 9, 2, 4, 8, 5, 3],
            [4, 2, 8, 3, 5, 7, 6, 1, 9],
            [9, 5, 3, 6, 8, 1, 2, 7, 4],
            [5, 9, 1, 2, 4, 8, 3, 6, 7],
            [8, 3, 4, 7, 1, 6, 9, 2, 5],
            [2, 6, 7, 5, 9, 3, 4, 8, 1]]
}

SUDOKU_SOLUTIONS['HARD'] = {
    '0': [[9, 3, 1, 8, 6, 5, 4, 2, 7],
          [6, 7, 8, 2, 4, 3, 9, 1, 5],
          [2, 4, 5, 7, 9, 1, 6, 8, 3],
          [3, 5, 4, 6, 2, 9, 8, 7, 1],
          [7, 6, 2, 1, 5, 8, 3, 4, 9],
          [8, 1, 9, 3, 7, 4, 2, 5, 6],
          [4, 9, 7, 5, 8, 6, 1, 3, 2],
          [5, 8, 3, 9, 1, 2, 7, 6, 4],
          [1, 2, 6, 4, 3, 7, 5, 9, 8]],    
    '1': [[8, 5, 1, 9, 4, 3, 6, 7, 2],
          [6, 9, 2, 7, 8, 1, 5, 4, 3],
          [7, 3, 4, 5, 2, 6, 9, 1, 8],
          [4, 6, 3, 8, 1, 2, 7, 5, 9],
          [1, 2, 7, 3, 5, 9, 8, 6, 4],
          [5, 8, 9, 6, 7, 4, 2, 3, 1],
          [3, 4, 5, 2, 6, 8, 1, 9, 7],
          [2, 1, 6, 4, 9, 7, 3, 8, 5],
          [9, 7, 8, 1, 3, 5, 4, 2, 6]]
}



class SudokuBase(ConstraintSatisfactionProblem):
      """ Sudoku Base CSP class. Initialises a Sudoku problem as a constraint satisfaction problem."""
      def __init__(self, init:SudokuPuzzle|None=None):
            """ Initialise the Sudoku CSP from a given template. If no template is given, uses a default easy template.
            
            :param init: A Sudoku puzzle template, represented as a 2D list of integers. 0 represents an empty cell.
            """
            if not init:  # default case
                  init = SUDOKU_TEMPLATES['EASY']['0']
            
            template = self.validate_template(init)

            self.create_variables(template)
            variables = aslist(self.grid)

            constraints = self.create_constraints(self.grid)

            super().__init__(variables, constraints)
      

      def validate_template(self, init:SudokuPuzzle) -> np.matrix:
            """ Validates the given Sudoku template and returns it as a numpy matrix.
            
            :param init: A Sudoku puzzle template, represented as a 2D list of integers. 0 represents an empty cell.
            :return: The validated Sudoku template as a numpy matrix
            """
            template = np.matrix(init)
            n = template.shape[0]
            assert n in [4,9], f"Sudoku template must be 4x4 or 9x9"
            assert template.shape==(n,n), f"Sudoku template invalid: {template}"
            
            self.n, self.m = n, (2 if n == 4 else 3)
            return template
      

      def create_variables(self, template: np.matrix) -> None:
            """ Creates a grid of Variables for the Sudoku CSP, and assigns values to any pre-filled cells in the template.
            
            :param template: The Sudoku puzzle template as a numpy matrix
            """
            domain = set(range(1,self.n+1))
            self.grid = np.matrix([[Variable(domain, name=str((i,j))) 
                                          for j in range(self.n)]
                                                for i in range(self.n)])
            for i in range(self.n):
                  for j in range(self.n):
                        if template[i,j] != 0:
                              self.grid[i,j].value = template[i,j]


      def __repr__(self) -> str:
            """ Print Sudoku CSP as a grid """
            w = int((self.n+self.m)*2+1)
            outstr = "—"*w+"\n"
            for i in range(self.n):
                  outstr += "| "
                  for j in range(self.n):
                        variable = self.grid[i,j]
                        outstr += str(variable.value) if variable.is_assigned else "?"
                        outstr += " " if (j+1) % self.m else " | "
                  outstr += "\n" if (i+1) % self.m else "\n"+"—"*w+"\n"
            return outstr


      def create_constraints(self, grid: np.matrix) -> set[Factor]:
            """ Constraints class for Sudoku CSP. Creates the constraints for the Sudoku problem. Called in the constructor.
            
            input: grid - the grid of Variables for the Sudoku CSP
            output: a set of Factors representing the constraints for the Sudoku problem
            """
            raise NotImplementedError("You need to implement the create_constraints method for the Sudoku class")