import Data.List
import Data.List.Split

type Vec = (Int, Int)
type Robot = (Vec, Vec)

size :: Vec
size = (101, 103)

mul :: Int -> Vec -> Vec
mul s (x, y) = (s*x, s*y)
add :: Vec -> Vec -> Vec
add (x, y) (z, w) = (x+z, y+w)
clamp :: Vec -> Vec -> Vec
clamp m@(mx, my) v@(x, y)
  | x < 0 = clamp m (x + mx, y)
  | y < 0 = clamp m (x, y + my)
  | x >= mx = clamp m (x - mx, y)
  | y >= my = clamp m (x, y - my)
  | otherwise = v

consume1 :: String -> Robot
consume1 str =
  ((atoi px, atoi py), (atoi vx, atoi vy))
  where
    atoi = (read::String->Int)
    ["p", px, py, "v", vx, vy] = splitOneOf "=, " str

consume :: String -> [Robot]
consume file =
  map consume1 $ lines file

step :: [Robot] -> [Robot]
step robots =
  map stepOne robots
  where
    stepOne (p, v) = (clamp size $ add p v, v)

applyN :: Int -> (a -> a) -> a -> a
applyN n f =
  foldr (.) id (replicate n f)

count :: Vec -> Vec -> [Robot] -> Int
count tl br robots =
  length $ filter (inside tl br) robots
  where
    inside (xs, ys) (xe, ye) ((x, y), _) =
      x >= xs && x < xe && y >= ys && y < ye

counts :: [Robot] -> (Int, Int, Int, Int)
counts robots =
  (count (0, 0)               (tx, ty) robots,
   count ((tx + 1), 0)        (mx, ty) robots,
   count (0, (ty + 1))        (tx, my) robots,
   count ((tx + 1), (ty + 1)) (mx, my) robots)
  where
    (mx, my) = size
    tx = (mx - 1) `div` 2
    ty = (my - 1) `div` 2

safety :: (Int, Int, Int, Int) -> Int
safety (a, b, c, d) =
  a * b * c * d

part1 :: [Robot] -> Int
part1 robots =
  safety $ counts $ applyN 100 step robots

main = do
  file <- readFile "day14.input"
  let robots = consume file
  putStrLn ("Part 1: " ++ (show $ part1 robots))
  --putStrLn ("Part 2: " ++ (show $ part2 robots))