Advent of Code 2020 Day 11

{-# language DeriveFunctor #-}
{-# language DeriveTraversable #-}
{-# language RecordWildCards #-}

import Control.Applicative (ZipList(ZipList), getZipList)
import Control.Monad ((>=>), (<=<))
import Data.List (tails, unfoldr)
import Data.List.NonEmpty (NonEmpty((:|)), nonEmpty)
import Data.Maybe (mapMaybe)

change1 :: Char -> Int -> Char
change1 'L' 0 = '#'
change1 '#' n | 4 <= n = 'L'
change1 c _ = c

update1 :: [String] -> Maybe Char
update1 ((ul:u:ur:_):(l:c:r:_):(bl:b:br:_):_) = pure . change1 c . length $ filter ('#' ==) [ul,u,ur,l,r,bl,b,br]
update1 ([ul,u]:[l,c]:[bl,b]:_) = pure . change1 c . length $ filter ('#' ==) [ul,u,l,bl,b]
update1 [(ul:u:ur:_),(l:c:r:_)] = pure . change1 c . length $ filter ('#' ==) [ul,u,ur,l,r]
update1 [[ul,u],[l,c]] = pure . change1 c . length $ filter ('#' ==) [ul,u,l]
update1 _ = Nothing

update1All :: [String] -> [String]
update1All = filter (not . null) . map (mapMaybe update1 . horiz) . vert

vert :: [String] -> [[String]]
vert [] = []
vert l@(h:_) = tails $ ('L' <$ h) : l

horiz :: [String] -> [[String]]
horiz [] = []
horiz l = getZipList $ traverse (ZipList . tails . ('L':)) l

fixUpdate1 :: [String] -> [String]
fixUpdate1 l = if u == l then u else fixUpdate1 u
   where u = update1All l

data Zip a
  = MkZip
  { revBefore :: [a]
  , current :: a
  , future :: [a]
  } deriving (Eq, Show, Functor, Foldable, Traversable)

instance Applicative Zip where
  pure x = MkZip { revBefore = repeat x, current = x, future = repeat x }
  f <*> x =
    MkZip
    { revBefore = zipWith ($) (revBefore f) (revBefore x)
    , current = current f (current x)
    , future = zipWith ($) (future f) (future x)
    }

toZip :: NonEmpty a -> Zip a
toZip (x:|xs) = MkZip { revBefore = [], current = x, future = xs }

reverseOnto :: [a] -> [a] -> [a]
reverseOnto [] l = l
reverseOnto (h:t) l = reverseOnto t (h : l)

toList :: Zip a -> [a]
toList MkZip{..} = reverseOnto revBefore (current : future)

moveNext :: Zip a -> Maybe (Zip a)
moveNext MkZip{..} = fmap mn $ nonEmpty future
 where
  mn (x:|xs) =
    MkZip
    { revBefore = current : revBefore
    , current = x
    , future = xs
    }

movePrev :: Zip a -> Maybe (Zip a)
movePrev MkZip{..} = fmap mp $ nonEmpty revBefore
 where
  mp (x:|xs) =
    MkZip
    { revBefore = xs
    , current = x
    , future = current : future
    }

moveD :: Zip (Zip a) -> Maybe (Zip (Zip a))
moveD = moveNext

moveU :: Zip (Zip a) -> Maybe (Zip (Zip a))
moveU = movePrev

moveR :: Zip (Zip a) -> Maybe (Zip (Zip a))
moveR = traverse moveNext

moveL :: Zip (Zip a) -> Maybe (Zip (Zip a))
moveL = traverse movePrev

search :: (Zip (Zip Char) -> Maybe (Zip (Zip Char))) -> Zip (Zip Char) -> Bool
search step = search'
 where
  search' area = case step area of
   Nothing -> False
   Just next -> case current $ current next of
     '#' -> True
     'L' -> False
     _ -> search' next

change2 :: Char -> Int -> Char
change2 'L' 0 = '#'
change2 '#' n | 5 <= n = 'L'
change2 c _ = c

update2 :: Zip (Zip Char) -> Char
update2 target = change2 (current $ current target) . length $ filter (flip search target) dirs
 where
  dirs =
    [ moveU >=> moveL, moveU, moveU >=> moveR
    , moveL, moveR
    , moveD >=> moveL, moveD, moveD >=> moveR
    ]

update2All :: Zip (Zip Char) -> Zip (Zip Char)
update2All = fmap (fmap update2 . traverse dupZ) . dupZ

dupZ :: Zip a -> Zip (Zip a)
dupZ z =
  MkZip
  { revBefore = unfoldr mp z
  , current = z
  , future = unfoldr mn z
  }
 where
  mp = fmap (\x -> (x,x)) . movePrev
  mn = fmap (\x -> (x,x)) . moveNext

fixUpdate2 :: Zip (Zip Char) -> Zip (Zip Char)
fixUpdate2 z = if z == u then z else fixUpdate2 u
 where u = update2All z

interactive :: Show a => (String -> a) -> IO ()
interactive f = print . f =<< getContents

part1 :: IO ()
part1 = interactive (sum . map (length . filter ('#' ==)) . fixUpdate1 . lines)

main :: IO ()
main = interactive (fmap (sum . map (length . filter ('#' ==)) . toList . fmap toList . fixUpdate2 . toZip . fmap toZip) . (traverse nonEmpty <=< nonEmpty) . lines)