{- logic of Connect Four -}

type Tile   = Int   -- 0, 1, or 2
type Player = Int   -- 1 and 2
type Move   = Int   -- column number
data State = State Player [[Tile]]  -- list of rows

empty :: Tile
empty = 0

num_rows, num_cols :: Int
num_rows = 6
num_cols = 7

start_player :: Player
start_player = 1

init_state :: State
init_state = State start_player
  (replicate num_rows (replicate num_cols empty))

other_player :: Player -> Player
other_player = (3 -)

drop_tile :: Move -> State -> State
drop_tile col (State player rows) = State
  (other_player player)
  (reverse $ drop_aux $ reverse rows)
    where
      drop_aux (row : rows) =
        case splitAt col row of
         (first, i : last) ->
           if i == empty
             then (first ++ player : last) : rows
             else row : drop_aux rows

valid_moves :: State -> [Move]
valid_moves (State _ rows) =
  map fst . filter ((== empty) . snd) . zip [0..] $ head rows

show_player :: Player -> String
show_player 1 = "X"
show_player 2 = "O"

show_tile :: Tile -> Char
show_tile t = if t == empty then '.' else head $ show_player t

show_state :: State -> String
show_state (State player rows) =
  unlines $ map (head . show) [0 .. num_cols - 1] :
    map (map show_tile) rows
     ++ ["\nPlayer " ++ show_player player ++ " to go"]

transpose_rows ([] : _) = []
transpose_rows xs = map head xs : transpose_rows (map tail xs)

winning_row :: Player -> [Tile] -> Bool
winning_row player [] = False
winning_row player row = take 4 row == replicate 4 player
  || winning_row player (tail row)

winning_player :: State -> Maybe Player
winning_player (State player rows) =
  let oplayer = other_player player
      long_rows = rows ++ transpose_rows rows
    in if any (winning_row oplayer) long_rows
      then Just oplayer
      else Nothing



{- user interface with I/O -}

main :: IO ()
main = do
  putStrLn "Welcome to Connect Four"
  game init_state

game :: State -> IO ()
game state = do
  putStrLn $ show_state state
  case winning_player state of
    Just player ->
      putStrLn $ show_player player ++ " wins!"
    Nothing -> let moves = valid_moves state in
      if null moves then putStrLn "Game ends in draw."
      else do
        putStr $ "Choose one of " ++ show moves ++ ": "
        move_str <- getLine
        let move = read move_str
        game (drop_tile move state)