{-# LANGUAGE MultiWayIf, BangPatterns #-}
import Prelude hiding (Right, Left)
import Control.Monad
import Data.Array.IO


type Idx = (Int, Int)
data Dir = Up | Right | Down | Left
  deriving (Show, Eq, Enum)

type Chart = IOUArray Idx Int
data State = State !Chart !Idx !Dir

concatMapM :: (Monad m) => (a -> m [b]) -> [a] -> m [b]
concatMapM f l = liftM concat (mapM f l)

printstate :: State -> IO ()
printstate (State chart idx dir) = do
    ((top, left), (bottom, right)) <- getBounds chart
    rows <- mapM (row (left, right)) [top..bottom]
    mapM_ putStrLn rows
  where
    pcell :: Idx -> Int -> String
    pcell idx' c =
        if | idx' == idx -> ['[', c']
           | idx' == add idx Right -> [']', c']
           | otherwise -> [' ', c']
      where c' = ".W#F" !! c

    row :: (Int, Int) -> Int -> IO String
    row (left, right) y = concatMapM (\i -> liftM (pcell i) (readArray chart i)) (range ((y, left), (y, right)))

add :: Idx -> Dir -> Idx
add (y, x) dir = case dir of
    Up -> (y - 1, x)
    Right -> (y, x + 1)
    Down -> (y + 1, x)
    Left -> (y, x - 1)

rotright :: Dir -> Dir
rotright Up = Right
rotright Right = Down
rotright Down = Left
rotright Left = Up

rotleft :: Dir -> Dir
rotleft Up = Left
rotleft Left = Down
rotleft Down = Right
rotleft Right = Up

rot180 :: Dir -> Dir
rot180 = rotright . rotright

type Virus = Int -> Dir -> (Int, Dir)

simulate :: Virus -> Int -> State -> IO (State, Int)
simulate virus numsteps initstate = do
    let foldfunc (st, n) _ = go st >>= \(st', b) -> return (st', n + fromEnum b)
    foldM foldfunc (initstate, 0) [1..numsteps]
  where
    go :: State -> IO (State, Bool)
    go !(State chart idx dir) = do
        bnds <- getBounds chart
        if inRange bnds idx
            then do
                c <- readArray chart idx
                let (c', dir') = virus c dir
                writeArray chart idx c'
                return (State chart (add idx dir') dir', c' == 2)
            else do
                let ((top, left), (bottom, right)) = bnds
                    bnds' = ((2 * top, 2 * left), (2 * bottom, 2 * right))
                chart' <- newArray bnds' 0
                getAssocs chart >>= mapM (uncurry (writeArray chart'))
                go (State chart' idx dir)

virus1 :: Virus
virus1 0 dir = (2, rotleft dir)
virus1 2 dir = (0, rotright dir)

virus2 :: Virus
virus2 0 dir = (1, rotleft dir)
virus2 1 dir = (2, dir)
virus2 2 dir = (3, rotright dir)
virus2 3 dir = (0, rot180 dir)

main :: IO ()
main = do
    input <- (readFile "22.in")
    let w = length (head (lines input))
        h = length (lines input)
        (halfw, halfh) = (w `quot` 2, h `quot` 2)

    chart <- newListArray ((-halfh, -halfw), (halfh, halfw))
                (map (\c -> fromEnum $ if c == '#' then 2 else 0) $ filter (/= '\n') input)

    (state1, numinfs1) <- simulate virus1 10000 (State chart (0, 0) Up)
    -- printstate state1
    print numinfs1

    (state2, numinfs2) <- simulate virus1 1000000 (State chart (0, 0) Up)
    -- printstate state2
    getBounds ((\(State ch _ _) -> ch) state2) >>= print
    print numinfs2