(define-module (day04)
  #:use-module (ice-9 rdelim)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-26)
  #:use-module (srfi srfi-41))

(define* (read-input #:optional (port (current-input-port)))
  (define (read-lines acc)
    (define line (read-line port))
    (if (and (not (eof-object? line)) (> (string-length line) 0))
        (read-lines (cons line acc))
        acc))
  (define strings (reverse! (read-lines '())))
  (define string->char-codes (compose (cut map char->integer <>) string->list))
  (list->typed-array 'u8 2 (map string->char-codes strings)))

(define xmas (map char->integer (string->list "XMAS")))

(define (start-word? arr i j di dj)
  (let loop ((i i) (j j) (lst xmas))
    (or (null-list? lst)
        (and (= (array-ref arr i j) (first lst))
             (loop (+ i di) (+ j dj) (cdr lst))))))

(define (count-xmas arr)
  (define-values (m n) (apply values (array-dimensions arr)))
  (define (compute-range dx size)
    (stream-range (if (< dx 0) (* -3 dx) 0)
                  (if (> dx 0) (- size (* 3 dx)) size)))
  (define (count-dir di dj)
    (define (count-row i)
      (stream-length (stream-filter (cut start-word? arr i <> di dj)
                                    (compute-range dj n))))
    (stream-fold + 0 (stream-map count-row (compute-range di m))))
  (stream-fold + 0
               (stream-map (cut apply count-dir <>)
                           (list->stream '((-1 -1) (-1 0) (-1 1) (0 -1) (0 1)
                                           (1 -1) (1 0) (1 1))))))

(define (is-x-mas? arr i j)
  (and (= (array-ref arr i j) (char->integer #\A))
       (let loop ((nm 0) (di '(1 1 -1 -1)) (dj '(1 -1 1 -1)))
         (if (null-list? di)
             (= nm 2)
             (let ((cur (array-ref arr (+ i (car di)) (+ j (car dj)))))
               (cond ((= cur (char->integer #\M))
                      (and (< nm 2) (loop (+ nm 1) (cdr di) (cdr dj))))
                     ((= cur (char->integer #\S))
                      (loop nm (cdr di) (cdr dj)))
                     (else #f)))))
       (not (= (array-ref arr (- i 1) (- j 1))
               (array-ref arr (+ i 1) (+ j 1))))))

(define (count-x-mas arr)
  (define-values (m n) (apply values (array-dimensions arr)))
  (define (count-row i)
    (stream-length (stream-filter (cut is-x-mas? arr i <>)
                                  (stream-range 1 (- n 1)))))
  (stream-fold + 0 (stream-map count-row (stream-range 1 (- m 1)))))

(define part1 (compose count-xmas read-input))

(define part2 (compose count-x-mas read-input))