
   0 | module Days.Day12
   1 | 
   2 | import Data.String
   3 | import Data.SortedMap
   4 | import Data.SortedSet
   5 | import Data.List.Lazy
   6 | import System
   7 | import Debug.Trace
   8 | 
   9 | import Util
  10 | import LazyListT
  11 | 
  12 | import Control.Eff
  13 | import Control.Eff.State
  14 | 
  15 | %hide Util.Direction
  16 | %hide Util.directions
  17 | 
  18 | directions : LazyList (Integer, Integer)
  19 | directions =
  20 |   [
  21 |     (0, 0)
  22 |   , (1, 0)
  23 |   , (0, 1)
  24 |   , (-1, 0)
  25 |   , (0, -1)
  26 |   ]
  27 | 
  28 | record Input where
  29 |   constructor MkInput
  30 |   rows, cols : Nat
  31 |   garden : Grid rows cols Char
  32 |   regions : SortedMap Nat (SortedSet (Coord rows cols))
  33 |   coloring : Grid rows cols (Maybe Nat)
  34 | %name Input input
  35 | 
  36 | parseInput : String -> Maybe Input
  37 | parseInput str = do
  38 |   (rows ** cols ** grid) <- stringTo2D str
  39 |   pure $ MkInput rows cols grid empty (const Nothing)
  40 | 
  41 | uncolored : (input : Input) -> LazyList (Coord input.rows input.cols)
  42 | uncolored input =
  43 |   map fst . filter (isNothing . snd ) . flat $ input.coloring
  44 | 
  45 | coloring : (input : Input) -> (loc : Coord input.rows input.cols) -> Maybe Nat
  46 | coloring input loc = loc `index` input.coloring
  47 | 
  48 | nextColoring : (input : Input) -> Nat
  49 | nextColoring input = 1 + foldl max 0 (keys input.regions)
  50 | 
  51 | regionFast : (i : Input) -> (seed : Coord i.rows i.cols)
  52 |            -> Eff [State (SortedSet (Coord i.rows i.cols)), LazyList] (Coord i.rows i.cols)
  53 | regionFast i seed = do
  54 |   cache <- get
  55 |   if contains seed cache
  56 |     then send Data.List.Lazy.Nil
  57 |     else do
  58 |       modify $ insert seed
  59 |       let seed_val = seed `index` i.garden
  60 |       let nexts = mapMaybe (step seed) directions
  61 |       next <- send nexts
  62 |       if next == seed
  63 |         then send (the (LazyList _) [seed])
  64 |         else do
  65 |           let next_val = next `index` i.garden
  66 |           if next_val == seed_val
  67 |             then regionFast i next
  68 |             else send Data.List.Lazy.Nil
  69 | 
  70 | -- Get all the coordinates within a region from a seed value
  71 | region : (input : Input) -> (seed : Coord input.rows input.cols)
  72 |        -> LazyList (Coord input.rows input.cols)
  73 | region input seed =
  74 |   fst . extract . runState empty . runLazyList $ regionFast input seed
  75 | 
  76 | 
  77 | -- Color the region containing the given seed value, returning nothing if the
  78 | -- value is already colored
  79 | colorRegion : (input : Input) -> (seed : Coord input.rows input.cols) -> Maybe Input
  80 | colorRegion input seed =
  81 |   case coloring input seed of
  82 |     Just _ => Nothing
  83 |     Nothing =>
  84 |       let reg = region input seed
  85 |           region_id = nextColoring input
  86 |           MkInput rows cols garden regions coloring = input
  87 |           region_set = foldl (\acc, e => insert e acc) empty reg
  88 |           coloring = foldl (\acc, e => replaceAt e (Just region_id) acc) coloring reg
  89 |           regions = insert region_id region_set regions
  90 |       in Just $ MkInput rows cols garden regions coloring
  91 | 
  92 | -- Color all the regions
  93 | colorAll : (input : Input) -> Input
  94 | colorAll input =
  95 |   case head' $ uncolored input of
  96 |     Nothing => input
  97 |     Just seed =>
  98 |       case colorRegion input seed of
  99 |         Nothing => input
 100 |         Just input => colorAll input
 101 | 
 102 | record Properties where
 103 |   constructor MkP
 104 |   area, perimeter : Nat
 105 | 
 106 | price : Properties -> Nat
 107 | price (MkP area perimeter) = area * perimeter
 108 | 
 109 | -- Get the area and perimeter of a region
 110 | getProps : (input : Input) -> (region : SortedSet (Coord input.rows input.cols)) -> Properties
 111 | getProps input region =
 112 |   let area = length . Data.SortedSet.toList $ region
 113 |       perimeter = sum . map (\x => 4 `minus` neighbors x) . Data.SortedSet.toList $ region
 114 |   in MkP area perimeter
 115 |   where
 116 |   neighbors : Coord input.rows input.cols -> Nat
 117 |   neighbors x =
 118 |     count (\x => contains x region) . mapMaybe (step x) . filter (/= (0,0)) $ directions
 119 | 
 120 | 
 121 | export
 122 | part1 : String -> IO (String, Input)
 123 | part1 str = do
 124 |   Just input <- pure $ parseInput str
 125 |     | _ => do
 126 |       printLn "Failed to parse input"
 127 |       exitFailure
 128 |   let colored = colorAll input
 129 |   let prices = map (price . getProps colored) . values $ colored.regions
 130 |   pure (show . sum $ prices, colored)
 131 | 
 132 | export
 133 | part2 : Input -> String -> IO String
 134 |