SameGame

SameGame.scala

object SameGame {
  final case class Position(col: Int, row: Int)

  sealed trait Color
  case object Green extends Color
  case object Blue  extends Color
  case object Red   extends Color
  case object Brown extends Color
  case object Gray  extends Color

  sealed trait CellState
  final case class Filled(color: Color) extends CellState
  case object Empty                     extends CellState

  final case class Cell(position: Position, state: CellState)
  final case class Group(color: Color, positions: Set[Position])
  final case class Column(cells: List[CellState]) extends AnyVal
  final case class Board(columns: List[Column])   extends AnyVal

  sealed trait SameGameState
  final case class InProgress(board: Board, score: Int) extends SameGameState
  final case class Finished(board: Board, score: Int)   extends SameGameState

  object Color {
    def apply(n: Int): Color = {
      n % 5 match {
        case 0 => Green
        case 1 => Blue
        case 2 => Red
        case 3 => Brown
        case 4 => Gray
      }
    }
  }

  object Position {
    def left(pos: Position): Position = {
      Position(pos.col - 1, pos.row)
    }

    def right(pos: Position): Position = {
      Position(pos.col + 1, pos.row)
    }

    def up(pos: Position): Position = {
      Position(pos.col, pos.row + 1)
    }

    def down(pos: Position): Position = {
      Position(pos.col, pos.row - 1)
    }
  }

  object Column {
    implicit class CellMapper(column: Column) {
      def map(f: (CellState, Int) => CellState): Column = {
        Column(column.cells.zipWithIndex.map { case (cs, i) => f(cs, i) })
      }

      def shiftDown: Column = {
        val nonEmptyCells = column.cells
          .filter(!CellState.isEmpty(_))

        val diff = column.cells.length - nonEmptyCells.length

        Column(nonEmptyCells ++ List.fill(diff)(Empty))
      }
    }

    def empty(height: Int): Column = Column((1 to height).map(_ => Empty).toList)
  }

  object CellState {
    def isEmpty(cellState: CellState): Boolean = {
      cellState match {
        case Empty => true
        case _     => false
      }
    }
  }

  object Board {
    implicit class ColumnMapper(board: Board) {
      def map(f: (Column, Int) => Column): Board = {
        Board(board.columns.zipWithIndex.map { case (c, i) => f(c, i) })
      }

      def shiftLeft: Board = {
        val nonEmptyColumns = board.columns
          .filter(column => !CellState.isEmpty(column.cells.head))

        val diff   = board.columns.length - nonEmptyColumns.length
        val height = board.columns.head.cells.length
        Board(nonEmptyColumns ++ List.fill(diff)(Column.empty(height)))
      }
    }
  }

  private val bonus = 1000

  private def sqr(x: Int): Int = x * x

  private def calcScore(group: Group): Int = sqr(group.positions.size - 2)

  private def penalty(numberOfFilledCells: Int): Int = -sqr(numberOfFilledCells - 2)

  private def getCellState(board: Board, position: Position): CellState = {
    val width  = board.columns.length
    val height = board.columns.head.cells.length
    if (position.col >= 0 && position.col < width && position.row >= 0 && position.row < height) {
      board.columns(position.col).cells(position.row)
    } else {
      Empty
    }
  }

  private def findAdjacentWithSameColor(board: Board, position: Position): Set[Position] = {
    getCellState(board, position) match {
      case Filled(color) =>
        Set(
          Position.up(position),
          Position.right(position),
          Position.down(position),
          Position.left(position)
        ).map(p => (getCellState(board, p), p))
          .filter {
            case (Filled(c), _) => c == color
            case _              => false
          }
          .map(_._2)

      case Empty => Set()
    }
  }

  private def hasValidMoves(board: Board): Boolean = {
    board.columns.zipWithIndex
      .exists {
        case (column, colIndex) =>
          column.cells.zipWithIndex
            .exists {
              case (_, rowIndex) =>
                findAdjacentWithSameColor(board, Position(colIndex, rowIndex)).nonEmpty
            }
      }
  }

  private def filledCells(board: Board): Int = {
    board.columns
      .foldLeft(0)((total, column) =>
        column.cells.foldLeft(total)((count, cell) =>
          cell match {
            case Filled(_) => count + 1
            case Empty     => count
        }))
  }

  private def findGroup(board: Board, position: Position): Option[Group] = {
    def find(toSearch: Set[Position], group: Set[Position]): Set[Position] =
      toSearch.headOption.fold(group) { head =>
        val cellsWithSameColor = findAdjacentWithSameColor(board, head)
        val cellsFoundSoFar    = group + head
        val stillToSearch      = (cellsWithSameColor ++ toSearch.tail) -- cellsFoundSoFar
        find(stillToSearch, cellsFoundSoFar)
      }

    getCellState(board, position) match {
      case Filled(color) =>
        val positions = find(Set(position), Set.empty)
        if (positions.size > 1) {
          Some(Group(color, positions))
        } else {
          None
        }
      case _ => None
    }
  }

  private def removeGroup(board: Board, group: Group): Board = {
    board.map {
      case (column, colIndex) =>
        column.map {
          case (cell, rowIndex) =>
            if (group.positions.contains(Position(colIndex, rowIndex))) {
              Empty
            } else {
              cell
            }
        }.shiftDown
    }.shiftLeft
  }

  private def play(board: Board, position: Position): Option[(Board, Int)] = {
    findGroup(board, position)
      .map(g => (removeGroup(board, g), calcScore(g)))
  }

  def evaluateGameState(board: Board, score: Int): SameGameState = {
    def isEmpty(board: Board): Boolean = filledCells(board) == 0

    if (hasValidMoves(board)) {
      InProgress(board, score)
    } else if (isEmpty(board)) {
      Finished(board, score + bonus)
    } else {
      Finished(board, score + penalty(filledCells(board)))
    }
  }

  def applyMove(position: Position, game: SameGameState): SameGameState =
    game match {
      case InProgress(board, score) =>
        play(board, position)
          .map { case (b, s) => evaluateGameState(b, s + score) }
          .getOrElse(game)
      case Finished(_, _) => game
    }

  def legalMoves(game: SameGameState): List[Position] =
    game match {
      case InProgress(board, _) =>
        board.columns.zipWithIndex
          .flatMap {
            case (col, colIndex) =>
              col.cells.zipWithIndex.map { case (_, rowIndex) => Position(colIndex, rowIndex) }
          }
          .flatMap(pos => findGroup(board, pos).toList)
          .distinct
          .flatMap(g => g.positions.headOption.toList)
      case Finished(_, _) => Nil
    }

  def score(game: SameGameState): Int =
    game match {
      case InProgress(_, score) => score
      case Finished(_, score)   => score
    }

  def apply(board: List[List[Int]]): SameGameState =
    SameGame.evaluateGameState(Board(board.map(c => Column(c.map(s => Filled(Color(s)))))), 0)
}