// Main Part 4 about the Shogun Board Game

//=========================================

object M4 {   

type Pos = (Int, Int)    // a position on a chessboard 

// Colours: Red or White

abstract class Colour

case object Red extends Colour

case object Wht extends Colour

// Pieces: Either Pawns or Kings

//===============================

abstract class Piece {

  def pos : Pos       

  def col : Colour    

  def en : Int      // energy for Pawns 1 - 4, for Kings 1 - 2

}

case class Pawn(en: Int, col: Colour, pos: Pos) extends Piece

case class King(en: Int, col: Colour, pos: Pos) extends Piece

//val p = Pawn(4, Wht, (3,2))

//assert(p.pos == (3,2))

//assert(p.col == Wht)

//assert(p.en == 4)  

// checks if a piece is a king

def is_king(pc: Piece) : Boolean = pc match {

  case King(_, _, _) => true

  case _ => false

}

// incrementing and decrementing the position of a piece

def incx(pc: Piece) : Piece = pc match {

  case Pawn(en, c, (x,y)) => Pawn(en, c, (x+1,y))

  case King(en, c, (x,y)) => King(en, c, (x+1,y))

}

def incy(pc: Piece) : Piece = pc match {

  case Pawn(en, c, (x,y)) => Pawn(en, c, (x,y+1))

  case King(en, c, (x,y)) => King(en, c, (x,y+1))

}

def decx(pc: Piece) : Piece = pc match {

  case Pawn(en, c, (x,y)) => Pawn(en, c, (x-1,y))

  case King(en, c, (x,y)) => King(en, c, (x-1,y))

}

def decy(pc: Piece) : Piece = pc match {

  case Pawn(en, c, (x,y)) => Pawn(en, c, (x,y-1))

  case King(en, c, (x,y)) => King(en, c, (x,y-1))

}

//pretty printing colours and pieces

def pp_color(c: Colour) : String = c match {

  case Red => "R"

  case Wht => "W"

}

def pp(pc: Piece) : String = pc match {

  case Pawn(n, c, _) => s"P${pp_color(c)}$n"

  case King(n, c, _) => s"K${pp_color(c)}$n"

}

// Boards are sets of pieces

//===========================

case class Board(pces: Set[Piece]) {

  def +(pc: Piece) : Board = Board(pces + pc)

  def -(pc: Piece) : Board = Board(pces - pc)

}

// checking whether a position is occupied in a board

def occupied(p: Pos, b: Board) : Option[Piece] =  

  b.pces.find(p == _.pos)

def occupied_by(p: Pos, b: Board) : Option[Colour] =

  occupied(p, b).map(_.col)

def is_occupied(p: Pos, b: Board) : Boolean =

  occupied(p, b).isDefined

// is a position inside a board

def inside(p: Pos, b: Board): Boolean = 

  1 <= p._1 && 1 <= p._2 && p._1 <= 8 && p._2 <= 8 

// pretty printing a board

def print_board(b: Board): Unit = {

  println()

  for (i <- 8 to 1 by -1) {

    println("+" ++ "-" * 31 ++ "+")

    for (j <- 1 to 8) {

      val opc = occupied((j,i), b)

      if (opc.isDefined) print(s"|${pp(opc.get)}") 

      else print("|   ")

    }

    println("|")

  } 

  println("+" ++ "-" * 31 ++ "+")

}

// example board: initial board

val b_init = Board(Set(King(2,Wht,(4,1)), King(1,Red,(5,8)),

                  		 Pawn(4,Wht,(1,1)), Pawn(4,Red,(1,8)),

                  		 Pawn(3,Wht,(2,1)), Pawn(2,Red,(2,8)),

                  		 Pawn(2,Wht,(3,1)), Pawn(3,Red,(3,8)),

                  		 Pawn(1,Wht,(5,1)), Pawn(1,Red,(4,8)),

                  		 Pawn(4,Wht,(6,1)), Pawn(3,Red,(6,8)),

                  		 Pawn(3,Wht,(7,1)), Pawn(1,Red,(7,8)),

                  		 Pawn(2,Wht,(8,1)), Pawn(3,Red,(8,8))))

//print_board(b_init)

// --------------------------------

// |PR4|PR2|PR3|PR1|KR1|PR3|PR1|PR3|

// --------------------------------

// |   |   |   |   |   |   |   |   |

// --------------------------------

// |   |   |   |   |   |   |   |   |

// --------------------------------

// |   |   |   |   |   |   |   |   |

// --------------------------------

// |   |   |   |   |   |   |   |   |

// --------------------------------

// |   |   |   |   |   |   |   |   |

// --------------------------------

// |   |   |   |   |   |   |   |   |

// --------------------------------

// |PW4|PW3|PW2|KW2|PW1|PW4|PW3|PW2|

// --------------------------------

// Moves

//=======

abstract class Move

case object U extends Move    // up

case object D extends Move    // down

case object R extends Move    // right

case object L extends Move    // left

case object RU extends Move   // ...

case object LU extends Move

case object RD extends Move

case object LD extends Move

case object UR extends Move

case object UL extends Move

case object DR extends Move

case object DL extends Move

// Task 1: calculates all next possible positions according to a move

def eval(pc: Piece, m: Move, en: Int, b: Board) : Set[Piece] = {

  val p = pc.pos

  val c = pc.col

  if (!inside(p, b)) Set() 

  else if (en == 0 && !is_occupied(p, b)) Set(pc)

  else if (en == 0 && is_occupied(p, b) && c != occupied_by(p, b).get) Set(pc)

  else if (is_occupied(p, b)) Set()  

  else m match {

    case  U => eval(incy(pc), U, en - 1, b) 

    case  D => eval(decy(pc), D, en - 1, b) 

    case  R => eval(incx(pc), R, en - 1, b) 

    case  L => eval(decx(pc), L, en - 1, b) 

    case  RU => eval(incx(pc), RU, en - 1, b) ++ eval(pc, U, en, b)

    case  LU => eval(decx(pc), LU, en - 1, b) ++ eval(pc, U, en, b)

    case  RD => eval(incx(pc), RD, en - 1, b) ++ eval(pc, D, en, b)

    case  LD => eval(decx(pc), LD, en - 1, b) ++ eval(pc, D, en, b)

    case  UR => eval(incy(pc), UR, en - 1, b) ++ eval(pc, R, en, b)

    case  UL => eval(incy(pc), UL, en - 1, b) ++ eval(pc, L, en, b)

    case  DR => eval(decy(pc), DR, en - 1, b) ++ eval(pc, R, en, b)

    case  DL => eval(decy(pc), DL, en - 1, b) ++ eval(pc, L, en, b)

}}

/*

// test cases

val pw_a = Pawn(4, Wht, (4,4))

println(eval(pw_a, U,  4, b_init))  // Set(Pawn(4,Wht,(4,8)))

println(eval(pw_a, U,  3, b_init))  // Set(Pawn(4,Wht,(4,7)))

println(eval(pw_a, RU, 4, b_init))  // Set(Pawn(4,Wht,(6,6)), Pawn(4,Wht,(4,8)),

                                    //     Pawn(4,Wht,(5,7)), Pawn(4,Wht,(7,5)), 

                                    //     Pawn(4,Wht,(8,4)))

val pw_b = Pawn(4, Red, (4,4))

println(eval(pw_b, RU, 4, b_init))  // Set(Pawn(4,Red,(8,4)), Pawn(4,Red,(7,5)), 

                                           Pawn(4,Red,(6,6)), Pawn(4,Red,(5,7)))

*/

// Task 2: calculates all possible moves for a piece

def all_moves(pc: Piece, b: Board) : Set[Piece] = {

  Set(U,D,L,R,RU,LU,RD,LD,UR,UL,DR,DL).flatMap(eval(pc, _, pc.en, b - pc))

}

/*

// test cases

val pw_c = Pawn(2, Wht, (4,4))

val pw_d = Pawn(3, Red, (4,4))

println(all_moves(pw_c, b_init))  

  // Set(Pawn(2,Wht,(3,5)), Pawn(2,Wht,(2,4)), Pawn(2,Wht,(3,3)), Pawn(2,Wht,(5,5)), 

  //     Pawn(2,Wht,(6,4)), Pawn(2,Wht,(4,6)), Pawn(2,Wht,(4,2)), Pawn(2,Wht,(5,3)))

println(all_moves(pw_d, b_init)) 

  // Set(Pawn(3,Red,(4,7)), Pawn(3,Red,(5,2)), Pawn(3,Red,(3,2)), Pawn(3,Red,(1,4)), 

  //     Pawn(3,Red,(6,3)), Pawn(3,Red,(3,6)), Pawn(3,Red,(2,5)), Pawn(3,Red,(2,3)), 

  //     Pawn(3,Red,(4,1)), Pawn(3,Red,(5,6)), Pawn(3,Red,(7,4)), Pawn(3,Red,(6,5)))

*/

// Task 3: calculates all pieces that are attacked by colour

def attacked(c: Colour, b: Board) : Set[Piece] = {

  val (me, opponent) = b.pces.partition(_.col == c)

  val all = me.flatMap(all_moves(_, b))

  opponent.filter(pc => is_occupied(pc.pos, Board(all)))

}

// test cases

val b_checkmate = Board(Set(King(2, Red, (4,2)), King(2, Wht, (7,1)),

                            Pawn(3, Red, (6,1)), Pawn(2, Wht, (8,4)),

                            Pawn(4, Red, (4,4)), Pawn(2, Wht, (4,1)),

                            Pawn(4, Red, (5,3)), Pawn(3, Wht, (8,7)),

                            Pawn(3, Red, (6,5))))

print_board(b_checkmate)

println(attacked(Red, b_checkmate)) // Set(Pawn(2,Wht,(8,4)), King(2,Wht,(7,1)))

println(attacked(Wht, b_checkmate)) // Set(Pawn(3,Red,(6,1)))

println(attacked(Wht, b_init)) // Set()

println(attacked(Red, b_init)) // Set()

// Task 4: calculates the number of pieces that attack a piece

def attackedN(pc: Piece, b: Board) : Int = {

  val (me, opponent) = b.pces.partition(_.col == pc.col)

  val all = opponent.toList.flatMap(all_moves(_, b))

  all.count(_.pos == pc.pos)

}

// test cases

println(attackedN(Pawn(2, Wht, (8,4)), b_checkmate)) // 3

println(attackedN(King(2, Wht, (7,1)), b_checkmate)) // 1

println(attackedN(Pawn(3, Red, (6,1)), b_checkmate)) // 1

// Task 5: calculates the number of pieces that protect a piece

def protectedN(pc: Piece, b: Board) : Int = {

  val (me, opponent) = b.pces.partition(_.col == pc.col)

  val all = (me - pc).toList.flatMap(all_moves(_, (b - pc)))

  all.count(_.pos == pc.pos)

}

println(protectedN(Pawn(2, Wht, (8,4)), b_checkmate)) // 1

println(protectedN(Pawn(4, Red, (5,3)), b_checkmate)) // 3

//

val pw1 = Pawn(4, Wht, (4,6))

val pw2 = Pawn(4, Wht, (2,4))

val pw3 = Pawn(3, Red, (6,8))

val pw4 = Pawn(2, Red, (2,8))

val bt = b_init + pw1 + pw2

print_board(bt)

println(s"Capture Red: ${attacked(Wht, bt)}")

  // Set(Pawn(2,Red,(2,8)), Pawn(3,Red,(6,8)))

println(s"Capture Wht: ${attacked(Red, bt)}")

  // Set(Pawn(4,Wht,(4,6)))

println(s"ProtectedN:  ${protectedN(pw3, bt)}")

  // 2

println(s"AttackedN:   ${attackedN(pw4, bt)}")

  // 2

println(s"all moves:   ${all_moves(pw2, bt)}")

  // Set(Pawn(4,Wht,(4,2)), Pawn(4,Wht,(1,7)), Pawn(4,Wht,(5,3)), Pawn(4,Wht,(5,5)), 

  //     Pawn(4,Wht,(2,8)), Pawn(4,Wht,(3,7)), Pawn(4,Wht,(6,4)))

}