import Control.Arrow
import Data.List
import Data.List.Split
import qualified Data.Map as M

main :: IO ()
main = interact $ show . (day14a &&& day14b 1000000000) . lines

day14a :: [String] -> Int
day14a = calc . turn . turn . turn . gravity . turn

turn :: [String] -> [String]
turn = transpose . reverse

gravity :: [String] -> [String]
gravity = map (intercalate "#" . map sort . splitOn "#")

calc :: [String] -> Int
calc = sum . zipWith (*) [1 ..] . map (count 'O') . reverse
  where
    count c = length . filter (== c)

genStates :: [String] -> [[String]]
genStates = map snd . filter f . zip [0 ..] . iterate (gravity . turn)
  where
    f (x, _) = x `mod` 4 == 0

findDupe :: M.Map [String] Int -> Int -> [[String]] -> (Int, Int)
findDupe m n (x:xs) =
  if x `M.member` m
    then (m M.! x, n)
    else findDupe (M.insert x n m) (succ n) xs

day14b :: Int -> [String] -> Int
day14b n input = calc (states !! num)
  where
    num = (n - s1) `mod` (s2 - s1) + s1
    (s1, s2) = findDupe M.empty 0 states
    states = genStates input
