import SparseCheck

data Kind = King | Queen | Rook | Bishop | Knight | Pawn
  deriving (Eq, Show)

data Colour = Black | White
  deriving (Eq, Show)

type Piece = (Colour,Kind)
type Square = (Int,Int)

data Board = Board
		[(Kind,Square)] -- white
		[(Kind,Square)] -- black
  deriving Show

pieceAt :: Board -> Square -> Maybe Piece
pieceAt (Board wkss bkss) sq =
        pieceAtWith White (pieceAtWith Black Nothing bkss) wkss
	where
	pieceAtWith c n [] = n
	pieceAtWith c n ((k,s):xs) = if s==sq then Just (c,k) else pieceAtWith c n xs

emptyAtAll :: Board -> (Square->Bool) -> Bool
emptyAtAll (Board wkss bkss) e =
	emptyAtAllAnd (emptyAtAllAnd True bkss) wkss
	where
	emptyAtAllAnd b []         = b
	emptyAtAllAnd b ((_,s):xs) = not (e s) && emptyAtAllAnd b xs

rmPieceAt White sq (Board wkss bkss) = Board (rPa sq wkss) bkss
rmPieceAt Black sq (Board wkss bkss) = Board wkss (rPa sq bkss)

rPa sq (ks@(k,s):kss) = if s==sq then kss else ks : rPa sq kss

putPieceAt sq (White,k) (Board wkss bkss) = Board ((k,sq):wkss) bkss
putPieceAt sq (Black,k) (Board wkss bkss) = Board wkss ((k,sq):bkss)

kingSquare :: Colour -> Board -> Square
kingSquare White (Board kss _) = kSq kss
kingSquare Black (Board _ kss) = kSq kss

kSq ((King,s):_)   = s
kSq (       _:kss) = kSq kss 

opponent Black = White
opponent White = Black

colourOf :: Piece -> Colour
colourOf (c,_) = c

kindOf :: Piece -> Kind
kindOf (_,k) = k

onboard :: Square -> Bool
onboard (p,q) = 1<=p && p<=8 && 1<=q && q<=8

forcesColoured White (Board kss _) = kss
forcesColoured Black (Board _ kss) = kss

emptyBoard = Board [] []

data Move = Move 
    Square    -- to here
    (Maybe Piece) -- capturing this
    (Maybe Piece) -- gaining promotion to this
    
data MoveInFull = MoveInFull Piece Square Move

tryMove :: Colour -> (Kind,Square) -> Move -> Board -> Maybe (MoveInFull,Board)
tryMove c ksq@(k,sq) m@(Move sq' mcp mpp) bd =
  if not (kingincheck c bd2) then Just (MoveInFull p sq m, bd2)
  else Nothing 
  where
  p   =   (c,k)
  bd1 = rmPieceAt c sq bd
  p'  = maybe p id mpp
  bd2 = maybe (putPieceAt sq' p' bd1)
          (const (putPieceAt sq' p' (rmPieceAt (opponent c) sq' bd1)))
          mcp

moveDetailsFor :: Colour -> Board -> [(MoveInFull,Board)]
moveDetailsFor c bd =
  foldr ( \ksq ms ->
    foldr (\rm ms' -> maybe id (:) (tryMove c ksq rm bd) ms')
                   ms
                   (rawmoves c ksq bd) )
        []
              (forcesColoured c bd)


-- NB raw move = might illegally leave the king in check.
rawmoves :: Colour -> (Kind,Square) -> Board -> [Move]
rawmoves c (k,sq) bd = m c sq bd
	where
        m = case k of
	    King   -> kingmoves
	    Queen  -> queenmoves
	    Rook   -> rookmoves
	    Bishop -> bishopmoves
	    Knight -> knightmoves
	    Pawn   -> pawnmoves

bishopmoves :: Colour -> Square -> Board -> [Move]
bishopmoves c sq bd =
	( moveLine bd c sq (\(x,y) -> (x-1,y+1)) $
	  moveLine bd c sq (\(x,y) -> (x+1,y+1)) $
	  moveLine bd c sq (\(x,y) -> (x-1,y-1)) $
	  moveLine bd c sq (\(x,y) -> (x+1,y-1)) id
        ) []

rookmoves :: Colour -> Square -> Board -> [Move]
rookmoves c sq bd =
	( moveLine bd c sq (\(x,y) -> (x-1,y)) $
	  moveLine bd c sq (\(x,y) -> (x+1,y)) $
	  moveLine bd c sq (\(x,y) -> (x,y-1)) $
	  moveLine bd c sq (\(x,y) -> (x,y+1)) id
        ) []

moveLine :: Board -> Colour -> Square -> (Square->Square) -> ([Move]->a) -> [Move] -> a
moveLine bd c sq inc cont = ml sq
	where
	ml sq ms =
		let sq' = inc sq in
		if onboard sq' then
			case pieceAt bd sq' of
			Nothing -> ml sq' (Move sq' Nothing Nothing : ms)
			Just p' -> if colourOf p' /= c then
					cont (Move sq' (Just p') Nothing : ms)
                                   else cont ms
		else cont ms

kingmoves :: Colour -> Square -> Board -> [Move]
kingmoves c (p,q) bd =
	sift c bd []     [(p-1,q+1), (p,q+1), (p+1,q+1),
	  	 	  (p-1,q),            (p+1,q),
		 	  (p-1,q-1), (p,q-1), (p+1,q-1)]

knightmoves :: Colour -> Square -> Board -> [Move]
knightmoves c (p,q) bd =
	sift c bd [] [	  	 (p-1,q+2),(p+1,q+2),
			  (p-2,q+1),		  (p+2,q+1),
                          (p-2,q-1),		  (p+2,q-1),
		  		 (p-1,q-2),(p+1,q-2) ]

sift :: Colour -> Board -> [Move] -> [Square] -> [Move]
sift _ _  ms [] = ms
sift c bd ms (sq:sqs) =
	if onboard sq then
		case pieceAt bd sq of
                Nothing -> sift c bd (Move sq Nothing Nothing : ms) sqs
		Just p' -> if colourOf p' == c then sift c bd ms sqs
                           else sift c bd (Move sq (Just p') Nothing : ms) sqs
	else sift c bd ms sqs

pawnmoves :: Colour -> Square -> Board -> [Move]
pawnmoves c (p,q) bd = movs ++ caps
	where
	movs =	let on1 = (p,q+fwd)
		    on2 = (p,q+2*fwd) in
		if pieceAt bd on1 == Nothing then
			promote on1 Nothing ++
			if (q==2 && c==White || q==7 && c==Black) &&
			 	pieceAt bd on2 == Nothing then [Move on2 Nothing Nothing] 
			else []
		else []
	caps =	concat [ promote sq mcp
                       | sq <- [(p+1,q+fwd), (p-1,q+fwd)],
                         mcp@(Just p') <- [pieceAt bd sq], colourOf p'/=c ]
	fwd  =	case c of
       		White -> 1
		Black -> -1
	promote sq@(x,y) mcp =  
		if (c==Black && y==1 || c==White && y==8) then
			map (Move sq mcp . Just)
			    [(c,Queen), (c,Rook), (c,Bishop), (c,Knight)]
		else [Move sq mcp Nothing]

queenmoves :: Colour -> Square -> Board -> [Move]
queenmoves c sq bd = bishopmoves c sq bd ++ rookmoves c sq bd

kingincheck :: Colour -> Board -> Bool
kingincheck c bd =
	any givesCheck (forcesColoured (opponent c) bd)
	where
	givesCheck (k,(x,y)) = kthreat k
		where
		kthreat King =
			abs (x-xk) <= 1 && abs (y-yk) <= 1
		kthreat Queen =
			kthreat Rook || kthreat Bishop
		kthreat Rook =
			x==xk &&
                        emptyAtAll bd (\(xe,ye) -> xe==xk && min y yk < ye && ye < max y yk) ||
			y==yk &&
                        emptyAtAll bd (\(xe,ye) -> ye==yk && min x xk < xe && xe < max x xk)
		kthreat	Bishop =
			x+y==xk+yk &&
			emptyAtAll bd (\(xe,ye) -> xe+ye==xk+yk && min x xk < xe && xe < max x xk) ||
			x-y==xk-yk &&
			emptyAtAll bd (\(xe,ye) -> xe-ye==xk-yk && min x xk < xe && xe < max x xk)
		kthreat	Knight =
			abs (x-xk) == 2 && abs (y-yk) == 1 ||
			abs (x-xk) == 1 && abs (y-yk) == 2
		kthreat Pawn =
			abs (x-xk) == 1 &&
			case c of
			Black -> yk == y+1
			White -> yk == y-1
	(xk,yk) = kingSquare c bd

checkmate :: Colour -> Board -> Bool
checkmate col b = null (moveDetailsFor col b) && kingincheck col b

-- Data types for SparseCheck

(king :- queen :- rook :- bishop :- knight :- pawn) =
  datatype $  ctr0 King
           \/ ctr0 Queen
           \/ ctr0 Rook
           \/ ctr0 Bishop
           \/ ctr0 Knight
           \/ ctr0 Pawn

(black :- white) = datatype (ctr0 Black \/ ctr0 White)

board = datatype (ctr2 Board)

-- Conversion functions

instance Convert Kind where
  term King   = king
  term Queen  = queen
  term Rook   = rook
  term Bishop = bishop
  term Knight = knight
  term Pawn   = pawn
  unterm      = converter  $  conv0 King
                          -+- conv0 Queen
                          -+- conv0 Rook
                          -+- conv0 Bishop
                          -+- conv0 Knight
                          -+- conv0 Pawn

instance Convert Colour where
  term Black = black
  term White = white
  unterm     = converter (conv0 Black -+- conv0 White)


instance Convert Board where
  term (Board a b) = board (term a) (term b)
  unterm = converter (conv2 Board)

-- Board generator

allDiff xs =
  caseOf xs alts where
    alts (x :- xs) =
          nil :-> true
      :|: (x |> xs) :-> forall xs (=/= x) & allDiff xs

onBoard sq =
  exists $ \(p, q) ->
      sq === pair p q
    & i 1 |<=| p & p |<=| i 8
    & i 1 |<=| q & q |<=| i 8

exactlyOne a bs =
  caseOf bs alts where
    alts (b, bs) =
          nil :-> false
      :|: (b |> bs) :-> (a === b & forall bs (=/= a))
                      ? (a =/= b & exactlyOne a bs)

kingPosition as x y =
  exists $ \(b, bs) ->
      as === (b |> bs)
    & (b === pair king (pair x y) ? kingPosition bs x y)

kingsDontTouch ws bs =
  exists $ \(wx, wy, bx, by) ->
      kingPosition bs bx by
    & kingPosition ws wx wy
    & ((bx |>| suc wx) ? (wx |>| suc bx)
    ?  (by |>| suc wy) ? (wy |>| suc by))

zipP :: Term [a] -> Term [b] -> Term [(a, b)] -> Pred
zipP as bs cs = caseOf (as, bs, cs) alts where
                  alts (a :- as :- b :- bs :- cs) = 
                        (nil, nil, nil) :-> true
                    :|: (a |> as, b |> bs, pair a b |> cs) :-> zipP as bs cs

validBoard :: Term Board -> Pred
validBoard b =
  exists $ \(ws :- bs :- wps :- bps :- wsqs :- bsqs :- sqs) ->
      b === board ws bs
    & append wsqs bsqs sqs
    & zipP wps wsqs ws
    & zipP bps bsqs bs
    & exactlyOne king wps
    & exactlyOne king bps
    & kingsDontTouch ws bs
    & forall sqs onBoard
    & allDiff sqs

-- Property

a ==> b = a <= b

prop_checkmate b = if   length ws == 2 && Pawn `elem` (map fst ws)
                   then not (checkmate Black b)
                   else True
  where
    ws = forcesColoured White b

main = print $ check 8 validBoard prop_checkmate