// Resit about the Shogun Board Game
//====================================

// Task 1 - 6 see below

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

// how to extract components from pieces
//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("----" * 8)
    for (j <- 1 to 8) {
      val opc = occupied((j,i), b)
      if (opc.isDefined) print(s"|${pp(opc.get)}") 
      else print("|   ")
    }
    println("|")
  } 
  println("----" * 8)
}

// 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   // first right, then possibly up
case object LU extends Move   // first left, then possibly up
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

//======================
// ADD YOUR CODE BELOW
//======================


// Task 1: 
def eval(pc: Piece, m: Move, en: Int, b: Board) : Set[Piece] = {
  val pos = pc.pos
  if (!inside(pos, b)) then Set()
  else if (en == 0) then 
    if is_occupied(pos, b) then   
      val occupant = occupied_by(pos, b)
      if occupant.get != pc.col then
        Set(pc) 
      else Set()
    else Set(pc)
  else if (is_occupied(pos,b)) then 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) 
  }
  }
}


// Task 2: 
def all_moves(pc: Piece, b: Board) : Set[Piece] = {
  val temporary_board = b - pc //not sure if this is
  //what was meant by changes to the board (Perth discussion on keats)
  val all_possible_moves = Set(U, D, R, L, RU, LU, RD, LD, UR, UL, DR, DL)
  all_possible_moves.flatMap(move => eval(pc, move, pc.en, temporary_board))
} 


// Task 3: 
def attacked(c: Colour, b: Board) : Set[Piece] = {
  val my_pieces = b.pces.filter(_.col == c)
  val opponent_pieces = b.pces.filter(_.col != c)
  val attacked_pieces = opponent_pieces.filter(pc => my_pieces.flatMap(p => all_moves(p, b)).exists(_.pos == pc.pos))
  attacked_pieces
}


// Task 4: 
def attackedN(pc: Piece, b: Board) : Int = {
  val opponent_pieces = b.pces.filter(_.col != pc.col)
  val attacking_pieces_num = opponent_pieces.count(p => all_moves(p, b).exists(_.pos == pc.pos))
  attacking_pieces_num
}


// Task 5: 
def protectedN(pc: Piece, b: Board) : Int = {
  val temporary_board = b - pc 
  val my_pieces = b.pces.filter(_.col == pc.col)
  val protecting_pieces_num = my_pieces.count(p => all_moves(p, temporary_board).exists(_.pos == pc.pos))
  protecting_pieces_num
}


// Task 6: 
def legal_moves(pc: Piece, b: Board) : Set[Piece] = {
  if is_king(pc) then 
    val all_king_moves = all_moves(pc, b)
    val valid_king_moves = all_king_moves.filter{m =>
      val temporary_board = Board(b.pces.filter(_.pos != m.pos))//remove potentially captured pieces
      val updated_temp_board = temporary_board + m
      val opponent_pieces = updated_temp_board.pces.filter(_.col != pc.col)
      val opponent_moves = opponent_pieces.flatMap(p => all_moves(p, updated_temp_board))
      !(opponent_moves.exists(_.pos == m.pos))}
    valid_king_moves
  else {all_moves(pc, b)}
}

/*
// more test cases
//=================
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)))
*/

}