{-# options_ghc -Wall #-}
{-# language BlockArguments, LambdaCase #-}
import Data.List
import System.Environment
import qualified Data.Map as M
import Data.Maybe
import qualified Data.Array as A


data Block = Rock | Space | Wall deriving (Eq, Ord, Show)


data Rotation = U | L | D | R deriving (Eq, Ord, Show, Bounded, Enum)

left, right, opp :: Rotation -> Rotation
left  = \case U -> L; L -> D; D -> R; R -> U
right = \case U -> R; R -> D; D -> L; L -> U
opp   = \case U -> D; D -> U; L -> R; R -> L

instance Semigroup Rotation where
  U <> r = r
  L <> r = left r
  R <> r = right r
  D <> r = opp r

instance Monoid Rotation where mempty = U

type I2 = (Int, Int)

data RArray e = RA !Rotation !(A.Array I2 e)

rotateI :: Rotation -> (I2, I2) -> I2 -> I2
rotateI U _                        i      = i
rotateI L (_, (xhi, _))   (x, y) = (y, xhi - x - 1)
rotateI D (_, (xhi, yhi)) (x, y) = (xhi - x - 1, yhi - y - 1)
rotateI R (_, (_,   yhi)) (x, y) = (yhi - y - 1, x)

(!) :: RArray e -> I2 -> e
RA r arr ! i = arr A.! rotateI r (A.bounds arr) i

rotate :: Rotation -> RArray e -> RArray e
rotate r1 (RA r2 arr) = RA (r1 <> r2) arr


block :: Char -> Maybe Block
block 'O' = Just Rock
block '.' = Just Space
block '#' = Just Wall
block _   = Nothing

type Board = [[Block]]

board :: String -> Board
board = map (map block') . lines where
  block' :: Char -> Block
  block' c = fromMaybe (error $ "unrecognised " ++ show c) $ block c

chunks :: (a -> Bool) -> [a] -> [[a]]
chunks _ [] = []
chunks want (x : xs)
  | want x    = let (as, bs) = span  want xs in (x : as) : chunks want bs
  | otherwise = let (as, bs) = break want xs in (x : as) : chunks want bs

fall :: Ord a => (a -> Bool) -> [a] -> [a]
fall p = concat . map sort . chunks p

rotateLeft :: [[a]] -> [[a]]
rotateLeft = transpose . map reverse

rotateRight :: [[a]] -> [[a]]
rotateRight = map reverse . transpose

rotate180 :: [[a]] -> [[a]]
rotate180 = reverse . map reverse

fallLeft, fallUp, fallRight, fallDown :: Board -> Board
fallLeft  = map $ fall (/= Wall)
fallUp    = rotateRight . fallLeft . rotateLeft
fallRight = rotate180 . fallLeft . rotate180
fallDown  = rotateLeft . fallLeft . rotateRight

scoreWith :: (a -> Bool) -> [[a]] -> Int
scoreWith p xss =
  let ℓ = length xss; ℓs = [ℓ, ℓ-1 .. 0] in
  sum $ zipWith (\xs y -> length (filter p xs) * y) xss ℓs

score :: Board -> Int
score = scoreWith (== Rock)

part1 :: Board -> Int
part1 = score . fallUp

type Cycle = (Board, Board, Board, Board)

cycle1 :: Board -> Cycle
cycle1 xs =
  let as = fallUp xs
      bs = fallLeft as
      cs = fallDown bs
      ds = fallRight cs in
  (as, bs, cs, ds)

nextCycle :: Cycle -> Cycle
nextCycle (_, _, _, xs) = cycle1 xs

cycles :: Board -> [Cycle]
cycles start = let c = cycle1 start in iterate nextCycle c

fourth :: (a,b,c,d) -> d
fourth (_,_,_,d) = d

findRepeat :: Ord a =>
  M.Map a Int -> M.Map Int a ->
  Int -> [a] -> (Int, Int, M.Map Int a)
findRepeat ai ia i (a:as)
  | Just j <- M.lookup a ai = (j, i, ia)
  | otherwise               = findRepeat ai' ia' (i + 1) as
    where ai' = M.insert a i ai; ia' = M.insert i a ia
findRepeat _ _ _ _ = error "you said it was infinite"

part2 :: Board -> Int
part2 ℓs =
  let loops      = 1_000_000_000
      (m, n, xs) = findRepeat M.empty M.empty 0 $ cycles ℓs
      index      = (loops - n) `mod` (n - m) - 1 in
  score $ fourth $ xs M.! (m + index)

main :: IO ()
main = do
  [part, file] <- getArgs
  ℓs <- board <$> readFile file
  case part of
    "1" -> print $ part1 ℓs
    "2" -> print $ part2 ℓs
    _   -> error "usage: $0 <part> <file>"

-- example1 :: [String]
-- example1 =
--   ["O....#....",
--    "O.OO#....#",
--    ".....##...",
--    "OO.#O....O",
--    ".O.....O#.",
--    "O.#..O.#.#",
--    "..O..#O..O",
--    ".......O..",
--    "#....###..",
--    "#OO..#...."]