let zip2d 'a 'b [h] [w] (a: [h][w]a) (b: [h][w]b) : [h][w](a, b) =
  map2 (map2 (\x y -> (x, y))) a b

let fst 'a 'b (t: (a, b)) : a = let (x, _) = t in x
let snd 'a 'b (t: (a, b)) : b = let (_, y) = t in y

let get 'a [h] [w] (def: a) (field: [h][w]a) (y: i64) (x: i64) : a =
  if x < 0 || y < 0 || x >= w || y >= h
    then def
    else field[y, x]

let stencil 'a 'b [h] [w] (def: a) (field: [h][w]a) (f: a -> (a,a,a,a) -> b) : [h][w]b =
  tabulate h (\y ->
    tabulate w (\x ->
      f (get def field y x)
        (get def field (y-1) x
        ,get def field y (x+1)
        ,get def field (y+1) x
        ,get def field y (x-1))))

let red4 'a 'b (f: a -> b) (g: b -> b -> b) (tup: (a, a, a, a)) : b =
  let (a, b, c, d) = tup
  in g (g (g (f a) (f b)) (f c)) (f d)

let markpits [h] [w] (field: [h][w]u8) : [h][w]u8 =
  stencil 9 field
    (\v env -> if red4 (\d -> v < d) (&&) env
                 then v + 1
                 else 255)

let pits [h] [w] (field: [h][w]u8) : []u8 =
  filter (!= 255) (flatten (markpits field))

let grow [h] [w] (field: [h][w]u8) (basin: [h][w]bool) : [h][w]bool =
  stencil (9, false) (zip2d field basin) (\(d, v) env -> d < 9 && (v || red4 snd (||) env))

let floodfill [h] [w] (field: [h][w]u8) : *[h][w]bool =
  let basin0 = map (map (!= 255)) (markpits field)
  let basin1 = grow field basin0
  let changed a b = reduce_comm (||) false (flatten (map2 (map2 (!=)) a b))
  let (res, _) =
        loop (a, b) = (basin0, basin1) while changed a b do
          (b, grow field b)
  in res

let merge [h] [w] (ids: [h][w]i64) : [h][w]i64 =
  stencil i64.highest ids
    (\n env -> if n == i64.highest
                 then i64.highest
                 else i64.min n (red4 id i64.min env))

let markbasins [h] [w] (basins: [h][w]bool) : *[h][w]i64 =
  let numbers = unflatten h w (iota (h * w))
  let ids0 = map2 (map2 (\n b -> if b then n else i64.highest)) numbers basins
  let ids1 = merge ids0
  let changed a b = reduce_comm (||) false (flatten (map2 (map2 (!=)) a b))
  let (res, _) =
        loop (a, b) = (ids0, ids1) while changed a b do
          (b, merge b)
  in res

let collect_used_ids [h] [w] (ids: [h][w]i64) : []i64 =
  let nums = flatten ids
  let bitmap = scatter (tabulate (h * w) (\_ -> false)) nums (map (\_ -> true) nums)
  in map snd (filter fst (zip bitmap (iota (h * w))))

let ipow (base: i64) (exponent: i64) : i64 =
  loop res = 1 for _i < exponent do res * base

entry main [h] [w] (field: [h][w]u8) : (i32, i64) =
  let ids = map (map (\n -> if n == i64.highest then -1 else n))
                (markbasins (floodfill field))
  let used = collect_used_ids ids
  let sizes = map (\n -> reduce_comm (+) 0 (map ((==n) >-> i64.bool) (flatten ids))) used
  let sizes_map = reduce_by_index (tabulate (h * w) (\_ -> 0))
                                  (+) 0
                                  sizes (map (\_ -> 1) sizes)
  let count_size_cumulcount =
        zip sizes_map (iota (h * w))
        |> filter (\(cnt, _sz) -> cnt > 0)
        |> reverse
        |> (\arr -> zip arr (scan (+) 0 (map fst arr)))
        |> take 3
  let wanted_count =
        map (\((cnt, _sz), cmlcnt) -> i64.max 0 (i64.min cnt (3 - (cmlcnt - cnt))))
            count_size_cumulcount
  let part2 = reduce_comm (*) 1
                (map2 (\wc ((_cnt, sz), _cmlcnt) -> ipow sz wc)
                      wanted_count count_size_cumulcount)
  in (reduce_comm (+) 0 (map i32.u8 (pits field))
     ,part2)