import Data.List
import qualified Data.Map.Strict as Map
import Data.Semigroup

rotateL, rotateR, tiltW :: Endo [[Char]]
rotateL = Endo $ reverse . transpose
rotateR = Endo $ map reverse . transpose
tiltW = Endo $ map tiltRow
  where
    tiltRow xs =
      let (a, b) = break (== '#') xs
          (os, ds) = partition (== 'O') a
          rest = case b of
            ('#' : b') -> '#' : tiltRow b'
            [] -> []
       in os ++ ds ++ rest

load rows = sum $ map rowLoad rows
  where
    rowLoad = sum . map (length rows -) . elemIndices 'O'

lookupCycle xs i =
  let (o, p) = findCycle 0 Map.empty xs
   in xs !! if i < o then i else (i - o) `rem` p + o
  where
    findCycle i seen (x : xs) =
      case seen Map.!? x of
        Just j -> (j, i - j)
        Nothing -> findCycle (i + 1) (Map.insert x i seen) xs

main = do
  input <- lines <$> readFile "input14"
  print . load . appEndo (tiltW <> rotateL) $ input
  print $
    load $
      lookupCycle
        (iterate (appEndo $ stimes 4 (rotateR <> tiltW)) $ appEndo rotateL input)
        1000000000