content: -- |
  -- Module      : AOC.Challenge.Day11
  -- License     : BSD3
  --
  -- Stability   : experimental
  -- Portability : non-portable
  --
  -- Day 11.  See "AOC.Solver" for the types used in this module!
  
  module AOC.Challenge.Day11 (
      day11a
    , day11b
    ) where
  
  import           AOC.Common                        (countTrue)
  import           AOC.Common.Point                  (Point, boundingBox', inBoundingBox, fullNeighbs, fullNeighbsSet, parseAsciiMap)
  import           AOC.Solver                        ((:~>)(..))
  import           Control.Monad.Loops               (whileM_)
  import           Control.Monad.Primitive           (PrimMonad, PrimState)
  import           Control.Monad.ST                  (runST)
  import           Control.Monad.State.Strict        (execStateT, StateT(..))
  import           Data.Bit                          (Bit(..))
  import           Data.Bits                         (popCount)
  import           Data.Finite                       (Finite, finites)
  import           Data.Foldable                     (find, toList, traverse_)
  import           Data.Map                          (Map)
  import           Data.Maybe                        (mapMaybe)
  import           Data.Set                          (Set)
  import           GHC.TypeNats                      (KnownNat)
  import           Linear                            (V2(..))
  import qualified Data.Map.Strict                   as M
  import qualified Data.Set                          as S
  import qualified Data.Vector.Generic.Sized         as VG
  import qualified Data.Vector.Sized                 as V
  import qualified Data.Vector.Unboxed.Mutable.Sized as MVU
  import qualified Data.Vector.Unboxed.Sized         as VU
  
  parseSeatMap :: String -> Map Point Bool
  parseSeatMap = parseAsciiMap $ \case
      'L' -> Just False
      '#' -> Just True    -- not in the input, but just for completion's sake
      _   -> Nothing
  
  compile
      :: Map Point (Set Point, Bool)
      -> (forall n. KnownNat n => V.Vector n [Finite n] -> VU.Vector n Bit -> r)
      -> r
  compile mp f = V.withSizedList (toList mp) $ \xs ->
              f (map go . S.toList . fst <$> xs)
                (VG.convert (Bit . snd <$> xs))
    where
      go :: KnownNat n => Point -> Finite n
      go = fromIntegral . (`M.findIndex` mp)
  
  seatRule
      :: forall n m s. (KnownNat n, PrimMonad m, s ~ PrimState m)
      => Int                          -- ^ exit seat threshold
      -> V.Vector n [Finite n]        -- ^ neighbors
      -> MVU.MVector n s Bit          -- ^ source
      -> MVU.MVector n s Bit          -- ^ target
      -> StateT Bool m ()             -- ^ changed
  seatRule thr ns src targ = traverse_ go finites
    where
      go :: Finite n -> StateT Bool m ()
      go i = StateT $ \changed -> do
        Bit x <- MVU.read src i
        n     <- countTrue unBit <$> traverse (MVU.read src) (ns `V.index` i)
        let x' = case x of
              False -> n == 0
              True  -> n < thr
        MVU.write targ i (Bit x')
        pure ((), changed || x /= x')
      {-# INLINE go #-}
  {-# INLINE seatRule #-}
  
  solveWith
      :: KnownNat n
      => Int                      -- ^ exit seat threshold
      -> V.Vector n [Finite n]    -- ^ neighbors
      -> VU.Vector n Bit
      -> Int                      -- ^ equilibrium size
  solveWith thr ns xs = runST $ do
      xs0 <- VU.thaw xs
      xs1 <- MVU.unsafeNew
      (`whileM_` pure ()) . flip execStateT False $ do
        seatRule thr ns xs0 xs1
        seatRule thr ns xs1 xs0
      popCount <$> VU.unsafeFreeze xs0
  
  
  -- | Get a map of points to all of those points' neighbors where there is
  -- a seat. Should only need to be computed once.
  lineOfSights1
      :: Set Point
      -> Map Point (Set Point)
  lineOfSights1 pts = M.fromSet go pts
    where
      go p = fullNeighbsSet p `S.intersection` pts
  
  
  day11a :: Map Point Bool :~> Int
  day11a = MkSol
      { sParse = Just . parseSeatMap
      , sShow  = show
      , sSolve = \mp -> Just $
          let los = lineOfSights1 (M.keysSet mp)
          in  compile (M.intersectionWith (,) los mp) (solveWith 4)
          -- compilation itself takes about 38ms
      }
  
  -- | Get a map of points to all of those points' visible neighbors. Should
  -- only need to be computed once.
  lineOfSights2
      :: V2 Point
      -> Set Point
      -> Map Point (Set Point)
  lineOfSights2 bb pts = M.fromSet go pts
    where
      go p = S.fromList
           . mapMaybe (los p)
           $ fullNeighbs 0
      los p d = find (`S.member` pts)
              . takeWhile (inBoundingBox bb)
              . tail
              $ iterate (+ d) p
  
  day11b :: Map Point Bool :~> Int
  day11b = MkSol
      { sParse = Just . parseSeatMap
      , sShow  = show
      , sSolve = \mp -> do
          bb <- boundingBox' (M.keys mp)
          let los = lineOfSights2 bb (M.keysSet mp)
          pure $ compile (M.intersectionWith (,) los mp) (solveWith 5)
          -- compilation itself takes about 32ms
      }
  