AntPuzzle.hs

module Main where
import Data.ByteString (pack)
import Data.Set (Set, fromList, member)
import Data.List (tails)
import Codec.BMP
import Data.Word

spread :: (a -> a) -> [a] -> [a]
spread f as = as'
  where as' = concat $ as : map (\a -> f a : as) as'

listOfDiffs :: [Int]
listOfDiffs = subtract 9 `spread` replicate 9 1

-- find sums of the digits for all numbers >= 1000
sumOfDigits :: [Int]
sumOfDigits = scanl (+) 1 $ drop 1000 listOfDiffs

main = print $ length points

-- find reachable points starting from (1000, 1000)
(points, queue) = walk Nothing 1 $ ((1000, 1000), sumOfDigits, sumOfDigits) : queue
  where walk prev 0 _ = ([],[])
        walk prev n ((p@(x,y),as'@(a:as),bs'@(b:bs)):q) =
          let -- traverse (and generate) the rest of the queue
              (ps,q') = walk prev' (n - 1 + length q'') q
              -- decide whether to return this point and enqueue its children
              (ps',q'',prev') = if a + b > 25 
                then (ps, [], Nothing)
                else (p:ps, 
                        -- don't go up if the last guy hit that
                        -- point going right
                        (if Just up == prev then id else (up':)) [right']
                        , Just right)
              -- point on the grid
              up      = (x,y+1)
              right   = (x+1,y)
              -- states to push on the queue
              up'     = (up, as', bs)
              right'  = (right, as, bs')
          in (ps',q''++q')

-- image of which points are reachable, valid, or neither
image :: BMP
image = packRGBA32ToBMP 1000 1000 . pack $ bytes
        -- all the points >= (1000,1000) whose coordinates sum to <= 25
  where grid = [ [ sx+sy <= 25 | sx <- sumOfDigits ] | sy <- sumOfDigits ]

        -- all the points >= (1000,1000) whose coordinates sum to <= 25
        -- and are reachable from (1000,1000) (there exists a path on the grid)
        reachable = tail $ scanl meld (True : repeat False) grid
          where meld bs ds = tail . scanl (uncurry . fuse) False $ zip bs ds
                fuse p q d = (p || q) && d

        -- bytestring encoding of BMP image
        bytes = do
          (gs,rs) <- take 1000 $ zip grid reachable
          (g,r)   <- take 1000 $ zip gs rs
          if r then green else if g then cyan else black

        -- colors!
        green     = [0,255,0,0]
        cyan      = [0,255,255,0]
        black     = [0,0,0,0]

-- main = writeBMP "temp.bmp" image