import qualified Data.Set as S
import Debug.Trace (traceShow)

type Coord = (Integer, Integer)
type Rock = [Coord]

rockFix :: Rock -> Rock
rockFix xs = let maxy = maximum (map snd xs) in map (\(x,y) -> (x+2, y-maxy)) xs

rocks :: [Rock]
rocks = map rockFix
  [ [(0,0), (1,0), (2,0), (3,0)]
  , [(1,0), (0,1), (1,1), (2,1), (1,2)]
  , [(2,0), (2,1), (0,2), (1,2), (2,2)]
  , [(0,0), (0,1), (0,2), (0,3)]
  , [(0,0), (0,1), (1,0), (1,1)]
  ]

simulate :: [Rock] -> Maybe Rock -> String -> S.Set Coord -> Integer -> S.Set Coord
simulate [] _ _ _ _ = error "finite stream of rocks"
simulate (r:rs) Nothing jets occupied 0 = occupied
simulate (r:rs) Nothing jets occupied n = simulate rs (Just next) jets occupied (n-1) where
    height = minimum (1 : map snd (S.toList occupied)) - 4
    next = map (\(x,y) -> (x, y+height)) r
simulate (r:rs) (Just rock) (j:js) occupied remaining = if valid dropped
    then simulate (r:rs) (Just dropped) js occupied remaining
    else simulate (r:rs) Nothing js occupied' remaining
    where
        pushF = case j of
            '>' -> \(x,y) -> (x+1, y)
            '<' -> \(x,y) -> (x-1, y)
            _ -> error ("invalid char" ++ show j)
        dropF (x,y) = (x, y+1)
        valid = all (\c@(x,y) -> c `S.notMember` occupied
                              && x >= 0 && x < 7 && y <= 0)
        pushed = if valid (map pushF rock) then map pushF rock else rock
        dropped = map dropF pushed
        occupied' = foldr S.insert occupied pushed


-- simulate2 :: [Rock] -> String -> [Rock] -> Maybe Rock -> String -> S.Set Coord -> Integer -> Integer -> Integer
-- simulate2 origR origJ [] _ _ _ _ _ = error "finite stream of rocks"
-- simulate2 origR origJ (r:rs) Nothing jets occupied 0 iters = undefined
-- simulate2 origR origJ (r:rs) Nothing jets occupied n iters = if iters > 0 && matchesR && matchesJ
--     then let h = minimum (map snd (S.toList occupied)) in
--          let repetitions = n `div` iters in
--          let addedHeight = h * repetitions in addedHeight + simulate2 origR origJ (r:rs) Nothing jets occupied (n `mod` iters) 0
--     else simulate2 origR origJ rs (Just next) jets occupied (n-1) (iters+1) where
--         height = minimum (1 : map snd (S.toList occupied)) - 4
--         next = map (\(x,y) -> (x, y+height)) r
--         matchesR = origR == take (length origR) (r:rs)
--         matchesJ = origJ == take (length origJ) jets
-- simulate2 origR origJ (r:rs) (Just rock) (j:js) occupied remaining iters = if valid dropped
--     then simulate2 origR origJ (r:rs) (Just dropped) js occupied remaining iters
--     else simulate2 origR origJ (r:rs) Nothing js occupied' remaining iters
--     where
--         pushF = case j of
--             '>' -> \(x,y) -> (x+1, y)
--             '<' -> \(x,y) -> (x-1, y)
--             _ -> error ("invalid char" ++ show j)
--         dropF (x,y) = (x, y+1)
--         valid = all (\c@(x,y) -> c `S.notMember` occupied
--                               && x >= 0 && x < 7 && y <= 0)
--         pushed = if valid (map pushF rock) then map pushF rock else rock
--         dropped = map dropF pushed
--         occupied' = foldr S.insert occupied pushed


main :: IO ()
main = do
    jets <- head . lines <$> getContents
    let height set = -(minimum $ map snd (S.toList set)) + 1

    putStr "part 1: "
    let tower = simulate (cycle rocks) Nothing (cycle jets) S.empty 2022
    print (height tower)

    -- part 2
    let repIters = lcm (fromIntegral $ length jets) (fromIntegral $ length rocks) :: Integer
    putStr "repeated: " >> print repIters

    let initial = simulate (cycle rocks) Nothing (cycle jets) S.empty repIters
    let repeated = simulate (cycle rocks) Nothing (cycle jets) S.empty (2*repIters)

    let diff ls = zipWith (-) (tail ls) ls
    print $ diff $ map (\i -> height $ simulate (cycle rocks) Nothing (cycle jets) S.empty (i*repIters)) [1..20]

    let initHeight = height initial
    let repHeight = height repeated - initHeight
    putStr "repeated structure height: " >> print (initHeight, repHeight)
    let remIters = 1000000000000 `mod` repIters
    let remaining = simulate (cycle rocks) Nothing (cycle jets) S.empty remIters
    let remHeight = if remIters > 0 then -(minimum $ map snd (S.toList remaining)) + 1 else 0
    -- let tower = simulate2 rocks jets (cycle rocks) Nothing (cycle jets) S.empty 2022 0

    putStr "part 2: "
    let repNum = 1000000000000 `div` repIters
    print $ (repNum * repHeight) + remHeight