module Tak.Win (
    won,
    territory,
    roadWin
) where

import qualified Data.Foldable as Foldable
import Data.Matrix hiding (trace)

import Tak.Types

won :: Board -> Player -> Player -> Maybe Finished
won board white black = case roadWin board of
    Just win -> Just win
    Nothing -> flatWin board white black

flatWin :: Board -> Player -> Player -> Maybe Finished
flatWin board white black
    | finished && wc > bc = Just $ FlatWin White wc bc
    | finished && wc < bc = Just $ FlatWin Black wc bc
    | finished && wc == bc = Just $ Draw wc bc
    | otherwise = Nothing
    where
        (wc, bc, em) = territory board
        finished = empty white || empty black || em == 0
        empty player = stonesRemaing player == 0 && capsRemaining player == 0

-- | Counts the squares owned by white and black, and empty squares, respectively.
territory :: Board -> (Int, Int, Int)
territory board = Foldable.foldr fn (0, 0, 0) board where
    fn ((Flat, White) : _) (wc, bc, empty) = (wc + 1, bc, empty)
    fn ((Flat, Black) : _) (wc, bc, empty) = (wc, bc + 1, empty)
    fn [] (wc, bc, empty) = (wc, bc, empty + 1)
    fn _ (wc, bc, empty) = (wc, bc, empty)

-- | Lists which colour owns each square, in a road-building sense.
owner :: Board -> Matrix (Maybe Colour)
owner board = fmap squareOwner board where
    squareOwner [] = Nothing
    squareOwner ((Standing, _) : _) = Nothing
    squareOwner ((Flat, colour) : _) = Just colour
    squareOwner ((Cap, colour) : _) = Just colour

roadWin :: Board -> Maybe Finished
roadWin board
    | xwinW || ywinW = Just $ RoadWin White
    | xwinB || ywinB = Just $ RoadWin Black
    | otherwise = Nothing
    where
        (xwinW, _) = foldr (roadFrom White PosX o) (False, unseen) (sides PosX)
        (xwinB, _) = foldr (roadFrom Black PosX o) (False, unseen) (sides PosX)
        (ywinW, _) = foldr (roadFrom White PosY o) (False, unseen) (sides PosY)
        (ywinB, _) = foldr (roadFrom Black PosY o) (False, unseen) (sides PosY)
        sides :: Dir -> [Loc]
        sides PosX = [(1, i) | i <- [1 .. ncols board]]
        sides NegX = sides PosX
        sides PosY = [(i, 1) | i <- [1 .. nrows board]]
        sides NegY = sides PosY
        unseen = fmap (\ _ -> False) board
        o = owner board

roadFrom :: Colour -> Dir -> Matrix (Maybe Colour) -> Loc -> (Bool, Matrix Bool) -> (Bool, Matrix Bool)
roadFrom _ _ _ _ (True, seen) = (True, seen)
roadFrom colour dir own loc@(i, j) (_, seen)
    | i < 1 || j < 1 || i > nrows own || j > ncols own = (False, seen)
    | seen ! loc = (False, seen)
    | own ! loc /= (Just colour) = (False, setElem True loc seen)
    | dir == PosX && i == nrows own = (True, seen)
    | dir == PosY && j == ncols own = (True, seen)
    | otherwise = foldr (roadFrom colour dir own) (False, (setElem True loc seen))
        [(i + 1, j), (i - 1, j), (i, j - 1), (i, j + 1)]