{- 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 = do
  putStrLn "Welcome to Connect Four"
  game init_state
  
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)