#lang racket/base

(require "../intcode.rkt"
         racket/list)

(module+ test
  (require rackunit))

(define (part1 program)
  (apply max
         (for/list ([phases (in-permutations (range 5))])
           (for/fold ([signal 0])
                     ([phase (in-list phases)])
             (define vm (execute (start-machine program (list phase signal))))
             (car (machine-outputs vm))))))

(module+ test
  (check-equal? (part1 (parse "3,15,3,16,1002,16,10,16,1,16,15,15,4,15,99,0,0")) 43210)
  (check-equal? (part1 (parse "3,23,3,24,1002,24,10,24,1002,23,-1,23,
101,5,23,23,1,24,23,23,4,23,99,0,0")) 54321)
  (check-equal? (part1 (parse "3,31,3,32,1002,32,10,32,1001,31,-2,31,1007,31,0,33,
1002,33,7,33,1,33,31,31,1,32,31,31,4,31,99,0,0,0")) 65210)
  (check-equal? (part1 (parse-file "input.txt")) 298586))

(define (part2 program)
  (define signals
    (for/list ([phases (in-permutations (range 5 10))])
      (define amplifiers
        (for/list ([phase (in-list phases)])
          (define vm (execute (start-machine program (list phase))))
          vm))
      (let loop ([signal 0]
                 [amplifiers amplifiers])
        (if (empty? amplifiers)
            signal
            (let* ([amp (machine-program (car amplifiers))]
                   [pc (machine-pc (car amplifiers))]
                   [new-vm (execute (machine amp (list signal) pc 0 #f '()))]
                   [new-amplifiers (if (machine-terminated new-vm)
                                       (cdr amplifiers)
                                       (append (cdr amplifiers) (list new-vm)))])
              (loop (car (machine-outputs new-vm)) new-amplifiers))))))
  (apply max signals))

(module+ test
  (check-equal? (part2 (parse "3,26,1001,26,-4,26,3,27,1002,27,2,27,1,27,26,
27,4,27,1001,28,-1,28,1005,28,6,99,0,0,5")) 139629729)
  (check-equal? (part2 (parse "3,52,1001,52,-5,52,3,53,1,52,56,54,1007,54,5,55,1005,55,26,1001,54,
-5,54,1105,1,12,1,53,54,53,1008,54,0,55,1001,55,1,55,2,53,55,53,4,
53,1001,56,-1,56,1005,56,6,99,0,0,0,0,10")) 18216)
  (check-equal? (part2 (parse-file "input.txt")) 9246095))