"""
SudokuProcessor - Sudoku puzzle generation, solving, and rendering using SAT solver.
Supports efficient diverse multi-solution generation.
"""
import random
from typing import List, Tuple, Optional
import numpy as np
import cv2

try:
    from pysat.solvers import Solver
    HAS_PYSAT = True
except ImportError:
    HAS_PYSAT = False
    print("Warning: pysat not found, install with: pip install python-sat")


class SudokuProcessor:
    """Handles Sudoku puzzle generation, solving, and image rendering."""
    
    def __init__(self, cell_size: int = 60, font_scale: float = 1.2, thickness: int = 2):
        self.cell_size = cell_size
        self.font_scale = font_scale
        self.thickness = thickness
        self.img_size = cell_size * 9
        
        # Colors (RGB)
        self.bg_color = (255, 255, 255)
        self.line_color = (0, 0, 0)
        self.original_color = (0, 0, 0)
        self.filled_color = (200, 0, 0)
        self.highlight_color = (255, 255, 200)
        
        self._base_clauses_cache = None

    # ==================== SAT Encoding ====================
    
    def _var(self, r: int, c: int, n: int) -> int:
        """Map (row, col, num) to SAT variable (1-indexed)."""
        return r * 81 + c * 9 + n + 1
    
    def _decode_var(self, v: int) -> Tuple[int, int, int]:
        v -= 1
        return v // 81, (v % 81) // 9, v % 9
    
    def _base_clauses(self) -> List[List[int]]:
        """Generate base Sudoku constraint clauses (cached)."""
        if self._base_clauses_cache is not None:
            return self._base_clauses_cache
        
        clauses = []
        for i in range(9):
            for j in range(9):
                clauses.append([self._var(i, j, n) for n in range(9)])
                for n1 in range(9):
                    for n2 in range(n1 + 1, 9):
                        clauses.append([-self._var(i, j, n1), -self._var(i, j, n2)])
        
        for n in range(9):
            for i in range(9):
                clauses.append([self._var(i, j, n) for j in range(9)])
                for j1 in range(9):
                    for j2 in range(j1 + 1, 9):
                        clauses.append([-self._var(i, j1, n), -self._var(i, j2, n)])
                clauses.append([self._var(j, i, n) for j in range(9)])
                for j1 in range(9):
                    for j2 in range(j1 + 1, 9):
                        clauses.append([-self._var(j1, i, n), -self._var(j2, i, n)])
            for br in range(3):
                for bc in range(3):
                    box = [self._var(br*3+di, bc*3+dj, n) for di in range(3) for dj in range(3)]
                    clauses.append(box)
                    for i1 in range(9):
                        for i2 in range(i1 + 1, 9):
                            clauses.append([-box[i1], -box[i2]])
        
        self._base_clauses_cache = clauses
        return clauses
    
    def _grid_clauses(self, grid: List[List[int]]) -> List[List[int]]:
        return [[self._var(i, j, grid[i][j] - 1)] 
                for i in range(9) for j in range(9) if grid[i][j] != 0]
    
    def _model_to_grid(self, model: List[int]) -> List[List[int]]:
        grid = [[0] * 9 for _ in range(9)]
        for v in model:
            if 0 < v <= 729:
                r, c, n = self._decode_var(v)
                grid[r][c] = n + 1
        return grid

    # ==================== Solving ====================
    
    def solve(self, grid: List[List[int]]) -> Optional[List[List[int]]]:
        if HAS_PYSAT:
            with Solver(name='g3') as s:
                for c in self._base_clauses() + self._grid_clauses(grid):
                    s.add_clause(c)
                return self._model_to_grid(s.get_model()) if s.solve() else None
        return self._solve_backtrack(grid)
    
    def _solve_backtrack(self, grid: List[List[int]]) -> Optional[List[List[int]]]:
        board = [row[:] for row in grid]
        return board if self._backtrack(board) else None
    
    def _backtrack(self, board: List[List[int]]) -> bool:
        empty = self._find_empty(board)
        if not empty:
            return True
        r, c = empty
        for num in range(1, 10):
            if self._is_valid(board, r, c, num):
                board[r][c] = num
                if self._backtrack(board):
                    return True
                board[r][c] = 0
        return False
    
    def _find_empty(self, board: List[List[int]]) -> Optional[Tuple[int, int]]:
        for i in range(9):
            for j in range(9):
                if board[i][j] == 0:
                    return (i, j)
        return None
    
    def _is_valid(self, board: List[List[int]], row: int, col: int, num: int) -> bool:
        if num in board[row]:
            return False
        if any(board[i][col] == num for i in range(9)):
            return False
        br, bc = 3 * (row // 3), 3 * (col // 3)
        return all(board[i][j] != num for i in range(br, br+3) for j in range(bc, bc+3))

    def count_solutions(self, grid: List[List[int]], limit: int = 2) -> int:
        if HAS_PYSAT:
            count = 0
            with Solver(name='g3') as s:
                for c in self._base_clauses() + self._grid_clauses(grid):
                    s.add_clause(c)
                while count < limit and s.solve():
                    count += 1
                    s.add_clause([-v for v in s.get_model() if 0 < v <= 729])
            return count
        return self._count_backtrack(grid, limit)
    
    def _count_backtrack(self, grid: List[List[int]], limit: int) -> int:
        board = [row[:] for row in grid]
        self._sol_count, self._sol_limit = 0, limit
        self._count_helper(board)
        return self._sol_count
    
    def _count_helper(self, board: List[List[int]]) -> bool:
        if self._sol_count >= self._sol_limit:
            return True
        empty = self._find_empty(board)
        if not empty:
            self._sol_count += 1
            return self._sol_count >= self._sol_limit
        r, c = empty
        for num in range(1, 10):
            if self._is_valid(board, r, c, num):
                board[r][c] = num
                if self._count_helper(board):
                    return True
                board[r][c] = 0
        return False

    def find_solutions(self, grid: List[List[int]], limit: int = 10) -> List[List[List[int]]]:
        if HAS_PYSAT:
            solutions = []
            with Solver(name='g3') as s:
                for c in self._base_clauses() + self._grid_clauses(grid):
                    s.add_clause(c)
                while len(solutions) < limit and s.solve():
                    model = s.get_model()
                    solutions.append(self._model_to_grid(model))
                    s.add_clause([-v for v in model if 0 < v <= 729])
            return solutions
        return self._find_backtrack(grid, limit)
    
    def _find_backtrack(self, grid: List[List[int]], limit: int) -> List[List[List[int]]]:
        board, solutions = [row[:] for row in grid], []
        self._find_helper(board, solutions, limit)
        return solutions
    
    def _find_helper(self, board: List[List[int]], solutions: List, limit: int) -> bool:
        if len(solutions) >= limit:
            return True
        empty = self._find_empty(board)
        if not empty:
            solutions.append([row[:] for row in board])
            return len(solutions) >= limit
        r, c = empty
        for num in range(1, 10):
            if self._is_valid(board, r, c, num):
                board[r][c] = num
                if self._find_helper(board, solutions, limit):
                    return True
                board[r][c] = 0
        return False

    # ==================== Generation ====================
    
    def generate(self, clues: int = 30, unique: bool = True) -> Tuple[List[List[int]], List[List[int]]]:
        """Generate a Sudoku puzzle with specified number of clues."""
        solution = self._generate_full_grid()
        puzzle = [row[:] for row in solution]
        
        cells = [(i, j) for i in range(9) for j in range(9)]
        random.shuffle(cells)
        
        removed, target = 0, 81 - clues
        for r, c in cells:
            if removed >= target:
                break
            backup = puzzle[r][c]
            puzzle[r][c] = 0
            if unique and self.count_solutions(puzzle, 2) != 1:
                puzzle[r][c] = backup
            else:
                removed += 1
        
        return puzzle, solution
    
    def _generate_full_grid(self) -> List[List[int]]:
        if HAS_PYSAT:
            with Solver(name='g3') as s:
                for c in self._base_clauses():
                    s.add_clause(c)
                cells = [(i, j) for i in range(9) for j in range(9)]
                random.shuffle(cells)
                assumptions = []
                for r, c in cells[:11]:
                    nums = list(range(9))
                    random.shuffle(nums)
                    for n in nums:
                        if s.solve(assumptions=assumptions + [self._var(r, c, n)]):
                            assumptions.append(self._var(r, c, n))
                            break
                s.solve(assumptions=assumptions)
                return self._model_to_grid(s.get_model())
        
        board = [[0] * 9 for _ in range(9)]
        self._fill_grid(board)
        return board
    
    def _fill_grid(self, board: List[List[int]]) -> bool:
        empty = self._find_empty(board)
        if not empty:
            return True
        r, c = empty
        nums = list(range(1, 10))
        random.shuffle(nums)
        for num in nums:
            if self._is_valid(board, r, c, num):
                board[r][c] = num
                if self._fill_grid(board):
                    return True
                board[r][c] = 0
        return False

    # ==================== Diverse Multi-Solution Generation ====================

    @staticmethod
    def _hamming(sol1: List[List[int]], sol2: List[List[int]]) -> int:
        """Count differing cells between two complete grids."""
        return sum(sol1[i][j] != sol2[i][j] for i in range(9) for j in range(9))

    @staticmethod
    def _greedy_diverse_select(
        candidates: List[List[List[int]]],
        target_count: int,
        min_hamming: int,
        _hamming_fn=None,
    ) -> List[List[List[int]]]:
        """
        Greedily select diverse solutions using farthest-point sampling.
        
        1. Start with a random candidate.
        2. Repeatedly add the candidate with maximum min-distance to the selected set.
        3. Stop when enough are selected or no candidate meets min_hamming.
        """
        if _hamming_fn is None:
            _hamming_fn = SudokuProcessor._hamming
        
        if len(candidates) <= 1:
            return list(candidates)
        
        n = len(candidates)
        
        # Pre-compute pairwise distances
        dist = [[0] * n for _ in range(n)]
        for i in range(n):
            for j in range(i + 1, n):
                d = _hamming_fn(candidates[i], candidates[j])
                dist[i][j] = d
                dist[j][i] = d
        
        # Farthest-point sampling
        selected = [random.randint(0, n - 1)]
        remaining = set(range(n)) - {selected[0]}
        
        while len(selected) < target_count and remaining:
            best_idx = -1
            best_min_dist = -1
            
            for r in remaining:
                min_d = min(dist[r][s] for s in selected)
                if min_d > best_min_dist:
                    best_min_dist = min_d
                    best_idx = r
            
            if best_min_dist < min_hamming:
                break
            
            selected.append(best_idx)
            remaining.discard(best_idx)
        
        return [candidates[i] for i in selected]

    def generate_multi_solution(
        self,
        clues: int,
        min_solutions: int = 2,
        max_solutions: int = 5,
        max_attempts: int = 100,
        min_hamming: int = 10
    ) -> Tuple[List[List[int]], List[List[List[int]]]]:
        """
        Generate a puzzle with multiple diverse solutions.
        
        Puzzle-first strategy:
          1. Generate a full grid, randomly remove (81-clues) cells WITHOUT
             uniqueness check → guaranteed to have ≥1 solution, likely many.
          2. Enumerate candidate solutions of this puzzle via SAT.
          3. Greedily select diverse solutions (farthest-point sampling).
          4. If not enough diverse solutions, retry with a new puzzle.
        
        This is correct because all returned solutions are guaranteed valid
        completions of the returned puzzle.
        
        Args:
            clues: Number of given cells.
            min_solutions: Minimum diverse solutions required.
            max_solutions: Maximum to return.
            max_attempts: Outer retry budget.
            min_hamming: Minimum pairwise Hamming distance.
        Returns:
            (puzzle, solutions) — all solutions are valid and pairwise diverse.
        Raises:
            RuntimeError: If unable to find a qualifying puzzle.
        """
        # Adaptive hamming threshold
        adaptive_hamming = min_hamming
        if min_hamming == 10:  # default → auto-adapt
            if clues >= 55:
                adaptive_hamming = 3
            elif clues >= 45:
                adaptive_hamming = 5
            elif clues >= 35:
                adaptive_hamming = 8
            else:
                adaptive_hamming = 12
        
        # Adaptive search depth: more empty cells → more solutions likely exist
        empty_cells = 81 - clues
        if empty_cells <= 15:
            max_search = 30
        elif empty_cells <= 25:
            max_search = 80
        elif empty_cells <= 40:
            max_search = 150
        else:
            max_search = 300
        
        for _ in range(max_attempts):
            # Phase 1: Generate puzzle (random removal, no uniqueness check)
            solution = self._generate_full_grid()
            puzzle = [row[:] for row in solution]
            
            cells = [(i, j) for i in range(9) for j in range(9)]
            random.shuffle(cells)
            for r, c in cells[:81 - clues]:
                puzzle[r][c] = 0
            
            # # Phase 2: Quick feasibility — need at least min_solutions solutions
            # quick_count = self.count_solutions(puzzle, min_solutions + 1)
            # if quick_count < min_solutions:
            #     continue
            
            # Phase 3: Enumerate candidates
            candidates = self.find_solutions(puzzle, max_search)
            if len(candidates) < min_solutions:
                continue

            # Phase 4: Greedy diverse selection
            diverse = self._greedy_diverse_select(
                candidates, max_solutions, adaptive_hamming
            )
            
            if len(diverse) >= min_solutions:
                return puzzle, diverse[:max_solutions]
        
        raise RuntimeError(
            f"Failed to generate puzzle with {min_solutions}-{max_solutions} "
            f"diverse solutions (hamming>={adaptive_hamming}) after {max_attempts} attempts"
        )

    # ==================== Encoding ====================
    
    def encode(self, grid: List[List[int]]) -> str:
        return ''.join(str(grid[i][j]) for i in range(9) for j in range(9))
    
    def decode(self, s: str) -> List[List[int]]:
        return [[int(s[i * 9 + j]) for j in range(9)] for i in range(9)]

    # ==================== Rendering ====================
    
    def render(
        self,
        grid: List[List[int]],
        highlight_new: Optional[Tuple[int, int]] = None,
        original: Optional[List[List[int]]] = None
    ) -> np.ndarray:
        img = np.full((self.img_size, self.img_size, 3), self.bg_color, dtype=np.uint8)
        cs = self.cell_size
        
        if highlight_new:
            r, c = highlight_new
            cv2.rectangle(img, (c * cs, r * cs), ((c+1) * cs, (r+1) * cs), self.highlight_color, -1)
        
        for i in range(10):
            thick = 3 if i % 3 == 0 else 1
            pos = i * cs
            cv2.line(img, (pos, 0), (pos, self.img_size), self.line_color, thick)
            cv2.line(img, (0, pos), (self.img_size, pos), self.line_color, thick)
        
        font = cv2.FONT_HERSHEY_SIMPLEX
        for i in range(9):
            for j in range(9):
                if grid[i][j] == 0:
                    continue
                is_original = original is None or original[i][j] != 0
                color = self.original_color if is_original else self.filled_color
                text = str(grid[i][j])
                (tw, th), _ = cv2.getTextSize(text, font, self.font_scale, self.thickness)
                cv2.putText(img, text, (j*cs + (cs-tw)//2, i*cs + (cs+th)//2), 
                           font, self.font_scale, color, self.thickness)
        
        return img


if __name__ == "__main__":
    proc = SudokuProcessor()
    print(f"Using {'SAT solver' if HAS_PYSAT else 'backtracking'}...")
    
    # Test unique puzzle
    puzzle, solution = proc.generate(clues=25, unique=True)
    print("Puzzle:"); [print(row) for row in puzzle]
    print(f"Clues: {sum(c != 0 for row in puzzle for c in row)}")
    
    cv2.imwrite("test_puzzle.png", cv2.cvtColor(proc.render(puzzle), cv2.COLOR_RGB2BGR))
    cv2.imwrite("test_solution.png", cv2.cvtColor(proc.render(solution, original=puzzle), cv2.COLOR_RGB2BGR))
    print("Saved test images.")
    
    # Test diverse multi-solution at various clue levels
    print("\n=== Testing diverse multi-solution generation ===")
    for clues in [25, 35, 45, 55]:
        print(f"\nClue {clues}:")
        try:
            puzzle_m, solutions_m = proc.generate_multi_solution(
                clues=clues, min_solutions=3, max_solutions=3, min_hamming=10
            )
            print(f"  Generated puzzle with {len(solutions_m)} diverse solutions")
            for i in range(len(solutions_m)):
                for j in range(i + 1, len(solutions_m)):
                    print(f"    Hamming(sol{i}, sol{j}) = {proc._hamming(solutions_m[i], solutions_m[j])}")
        except RuntimeError as e:
            print(f"  {e}")