// Conway's Game of Life in Baba Yaga

// headAt: return the element at index i from list xs
headAt : xs i -> with (tmp : slice xs i (i + 1)) -> tmp.0;

// get2: 2D index for a grid (list of lists)
get2 : grid r c -> headAt (headAt grid r) c;

// safeGet2: bounds-checked 2D get; returns 0 when out of bounds
safeGet2 : grid r c ->
  with (
    rows : length grid;
    cols : length (headAt grid 0);
    rOk  : (r >= 0) and (r < rows);
    cOk  : (c >= 0) and (c < cols);
  ) ->
    when (rOk and cOk) is
      true then get2 grid r c
      _    then 0;

// range [lo, hi) as a list of Int
range : lo hi ->
  when (lo >= hi) is
    true then []
    _    then prepend lo (range (lo + 1) hi);

// sum a list of numbers
sum : xs -> reduce (acc x -> acc + x) 0 xs;

// deltas for neighborhood offsets
deltas : [-1, 0, 1];

// neighborsValues: list of neighbor cell values for (r,c)
neighborsValues : grid r c ->
  reduce (acc dr ->
    reduce (acc2 dc ->
      when ((dr = 0) and (dc = 0)) is
        true then acc2
        _    then append acc2 (safeGet2 grid (r + dr) (c + dc))
    ) acc deltas
  ) [] deltas;

// count live neighbors at (r,c)
countNeighbors : grid r c -> sum (neighborsValues grid r c);

// nextCell: apply Game of Life rules at (r,c)
nextCell : grid r c ->
  with (
    cell    : get2 grid r c;
    n       : countNeighbors grid r c;
    isAlive : cell = 1;
    born    : (cell = 0) and (n = 3);
    survive : isAlive and ((n = 2) or (n = 3));
  ) ->
    when survive is
      true then 1
      _    then when born is true then 1 _ then 0;

// build next row r given width w
rowNext : grid w r -> map (c -> nextCell grid r c) (range 0 w);

// Single simulation step for the entire grid
step : grid ->
  with (
    h : length grid;
    w : length (headAt grid 0);
  ) -> map (r -> rowNext grid w r) (range 0 h);

// Demo: glider pattern in a 5x5 grid
g0 : [
  [0, 1, 0, 0, 0],
  [0, 0, 1, 0, 0],
  [1, 1, 1, 0, 0],
  [0, 0, 0, 0, 0],
  [0, 0, 0, 0, 0]
];

g1 : step g0;
g2 : step g1;
g3 : step g2;
g4 : step g3;

// Print successive generations
io.out g0;
io.out g1;
io.out g2;
io.out g3;
io.out g4;