{ lib, utils, ... }:
with lib;
with builtins;
with utils;
let
  lines = readLines ./input;
  grid = map stringToCharacters lines;
  cols = map concatStrings (transpose grid);
  construct_diags' = matrix:
    map (filter (invert isNull)) (transpose (map (i:
      (genList (const null) i) ++ (elemAt matrix i)
      ++ (genList (const null) ((length (head matrix)) - i - 1)))
      (listRange matrix)));
  construct_diags = matrix: construct_diags' (map reverseList matrix);
  diags_1 = map concatStrings (construct_diags grid);
  diags_2 = map concatStrings (construct_diags' grid);
in {
  part1 =
    listSumWith (s: (length (matchAll "XMAS" s)) + (length (matchAll "SAMX" s)))
    (lines ++ cols ++ diags_1 ++ diags_2);
  part2 = listSumWith (li:
    let
      l = elemAt grid li;
      t = elemAt grid (li - 1);
      b = elemAt grid (li + 1);
    in count (ri:
      let
        tl = elemAt t (ri - 1);
        tr = elemAt t (ri + 1);
        mm = elemAt l ri;
        bl = elemAt b (ri - 1);
        br = elemAt b (ri + 1);

        tests = map concatStrings [
          [ tl mm br ]
          [ tr mm bl ]
          [ br mm tl ]
          [ bl mm tr ]
        ];

        # Idea: each X-MAS contains exactly 2 matching MAS centered around the A
      in (count (eq "MAS") tests) == 2) (range 1 ((length l) - 2)))
    (range 1 ((length grid) - 2));
}