import Debug.Trace
import Data.Array

traceOut :: Show a => a -> a
traceOut x = traceShow x x

addFirst :: [[a]] -> a -> [[a]]
addFirst [] x = [[x]]
addFirst (y:ys) x = (x : y) : ys

split :: Eq a => [a] -> [a] -> [[a]]
split _ [] = []
split delims (x:xs)
  | x `elem` delims = [] : split delims xs
  | otherwise       = addFirst (split delims xs) x

enumerate :: [a] -> [(Int, a)]
enumerate = zip [0..]

flipIndex :: (a, [b]) -> [(a, b)]
flipIndex (i, xs) = zip (repeat i) xs

flipTuple3 :: (a, (b, c)) -> ((a, b), c)
flipTuple3 (x, (y, z)) = ((x, y), z)

parseFile :: String -> Array (Int, Int) Bool
parseFile = toArray . map (map (=='@')) . split ['\n']

toArray :: [[x]] -> Array (Int, Int) x
toArray xs = array ((0, 0), (outerLength-1, innerLength-1)) $ toArrayList xs
    where outerLength = length xs
          innerLength = foldl max 0 $ map length xs
          toArrayList = concatMap (map flipTuple3 . enumerate . flipIndex) . enumerate

getSum :: Array (Int, Int) Bool -> (Int, Int) -> Int
getSum arr (x, y) = if arr ! (x, y) then sum [if arr ! (a, b) then 1 else 0 | a <- [max minX (x-1) .. min maxX (x+1)], b <- [max minY (y-1) .. min maxY (y+1)]] else 999
    where ((minX, minY), (maxX, maxY)) = bounds arr

removeRolls :: Array (Int, Int) Bool -> (Int, Array (Int, Int) Bool)
removeRolls arr = (length is, arr // [(i, False) | i <- is])
    where is = filter ((<5) . getSum arr) $ indices arr

takeWhileNeq :: Eq a => [a] -> [a]
takeWhileNeq [] = []
takeWhileNeq [a] = [a]
takeWhileNeq (a:b:bs) = a : (if a == b then [] else takeWhileNeq (b:bs))

solve1 :: Array (Int, Int) Bool -> Int
solve1 = fst . removeRolls

solve2 :: Array (Int, Int) Bool -> Int
solve2 arr = sum $ map fst $ takeWhileNeq $ iterate (removeRolls . snd) (0, arr)

main = do 
    fileinp <- readFile "input.txt"
    let input = parseFile fileinp
    let solved1 = solve1 input
    let solved2 = solve2 input
    print solved1
    print solved2