Sudoku.scala

import scala.util.Random

object Sudoku {
  lazy val localSquareLength = 3
  lazy val length = localSquareLength * localSquareLength
  lazy val totalCells = length * length
  lazy val sideIndices = 0 to (length - 1)
  lazy val legalValues = 1 to length
  lazy val emptyCells = for {
    y <- sideIndices
    x <- sideIndices
  } yield {
    Cell(x, y)
  }

  def randomPuzzle(blankCellCount: Int): Option[Puzzle] = {
    randomizedFullGrids.flatMap(grid => removeCellValues(grid, blankCellCount)).nextOption()
  }

  private def removeCellValues(puzzle: Sudoku.Puzzle, removeCellCount: Int): Option[Puzzle] = {
    @scala.annotation.tailrec
    def removeCellValues0(filledValues: Seq[CellValue], permanentValues: Seq[CellValue], removeCellCount: Int): Option[Puzzle] = {
      filledValues match {
        case head :: tail => {
          val updated = Puzzle(tail.++(permanentValues).toSet)
          val solutions = fill(updated)
          if (solutions.take(2).size == 1) {
            if (removeCellCount == 1) {
              Some(updated) //done
            } else {
              removeCellValues0(tail, permanentValues, removeCellCount - 1) //successfully removed an element, need to remove more
            }
          } else {
            removeCellValues0(tail, permanentValues.:+(head), removeCellCount) //couldn't remove the element, keep it in the puzzle and move on
          }
        }
        case Nil => None
      }
    }

    removeCellValues0(Random.shuffle(puzzle.filledCells.toSeq), Seq.empty, removeCellCount)
  }

  lazy val randomizedFullGrids: Iterator[Puzzle] = fill(Puzzle(filledCells = Set.empty))

  //@scala.annotation.tailrec
  def fill(puzzle: Puzzle): Iterator[Puzzle] = {
    puzzle.legalChanges match {
      case Failed => Iterator.empty
      case Success(legalChanges, isDone) => {
        if (isDone) {
          legalChanges
        } else {
          legalChanges.flatMap(changedPuzzle => fill(changedPuzzle))
        }
      }
    }
  }

  case class Cell(x: Int, y: Int) {
    def sharesRowOrColumn(other: Cell) = x == other.x || y == other.y || localSquare == other.localSquare
    lazy val localSquare = (x / localSquareLength, y / localSquareLength)
  }

  case class CellValue(cell: Cell, value: Int) {
    def isCompatible(other: CellValue) = {
      val incompatible = value == other.value && cell.sharesRowOrColumn(other.cell)
      !incompatible
    }
  }

  case class Puzzle(filledCells: Set[CellValue]) {
    lazy val legalChanges: LegalAdditions = {
      emptyCells
        .find(cell => !filledCells.exists(fc => fc.cell == cell))
        .map { nextFreeCell =>
          val candidates =
            Random.shuffle(legalValues.map(value => CellValue(nextFreeCell, value)))
              .iterator
              .flatMap(tryAdd)

          if (candidates.hasNext)
            Success(candidates, isDone = filledCells.size + 1 >= totalCells)
          else
            Failed
        }.getOrElse {
        Failed
      }
    }

    def tryAdd(newCell: CellValue): Option[Puzzle] = {
      if (filledCells.forall(filled => filled.isCompatible(newCell))) {
        Some(Puzzle(filledCells + newCell))
      }
      else {
        None
      }
    }

    override lazy val toString = {
      emptyCells
        .map(cell => filledCells.find(cv => cv.cell == cell).map(cv => cv.value.toString).getOrElse("_"))
        .grouped(length)
        .map(row => row.mkString(" "))
        .mkString("\n")
    }
  }

  trait LegalAdditions
  case class Success(legalChanges: Iterator[Puzzle], isDone: Boolean) extends LegalAdditions
  object Failed extends LegalAdditions
}

object RunSudoku extends App {

  import scala.util.chaining.scalaUtilChainingOps
  import Sudoku._

  Sudoku.randomPuzzle(40).tap(x => println(s"final:\n${x.get}"))

  def parsePuzzle(puzzle: String) =
    puzzle
      .trim
      .split("\n")
      .zipWithIndex
      .flatMap { case (row, rowNum) =>
        row.split(" ", Sudoku.length).zipWithIndex.flatMap { case (value, colNum) =>
          if (value == "_")
            None
          else
            Some(CellValue(Cell(x = colNum + 1, y = rowNum + 1), value = value.toInt))
        }
      }
      .pipe(cellValues => Puzzle(cellValues.toSet))
}