#lang curly-fn racket

(require "scripts/aoc.rkt")

;; solution for day 9

(define (cmp a b)
  (cond
    [(< a b) '<]
    [(> a b) '>]
    [else '=]))

(define (move1 hx hy tx ty dir)
  (match* (dir (cmp hx tx) (cmp hy ty))
    [("U" _ (not '>)) (values hx (add1 hy) tx ty)]
    [("D" _ (not '<)) (values hx (sub1 hy) tx ty)]
    [("R" (not '>) _) (values (add1 hx) hy tx ty)]
    [("L" (not '<) _) (values (sub1 hx) hy tx ty)]

    [("U" '= _) (values hx (add1 hy) tx (add1 ty))]
    [("U" '< _) (values hx (add1 hy) (sub1 tx) (add1 ty))]
    [("U" '> _) (values hx (add1 hy) (add1 tx) (add1 ty))]

    [("D" '= _) (values hx (sub1 hy) tx (sub1 ty))]
    [("D" '< _) (values hx (sub1 hy) (sub1 tx) (sub1 ty))]
    [("D" '> _) (values hx (sub1 hy) (add1 tx) (sub1 ty))]

    [("R" _ '=) (values (add1 hx) hy (add1 tx) ty)]
    [("R" _ '<) (values (add1 hx) hy (add1 tx) (sub1 ty))]
    [("R" _ '>) (values (add1 hx) hy (add1 tx) (add1 ty))]

    [("L" _ '=) (values (sub1 hx) hy (sub1 tx) ty)]
    [("L" _ '<) (values (sub1 hx) hy (sub1 tx) (sub1 ty))]
    [("L" _ '>) (values (sub1 hx) hy (sub1 tx) (add1 ty))]))

(define (part1 input)
  (for*/fold ([hx 0] [hy 0] [tx 0] [ty 0] [visited (set (cons 0 0))] #:result (set-count visited))
             ([inst (in-list input)]
              [dir (in-value (first inst))]
              [_ (in-range (second inst))])
    (define-values [hx* hy* tx* ty*] (move1 hx hy tx ty dir))
    ; (printf "DIR: ~a | h: (~a, ~a) t: (~a ~a)\n" dir hx* hy* tx* ty*)
    (values hx* hy* tx* ty* (set-add visited (cons tx* ty*)))))

(define (apply-move dir p)
  (match dir
    ["U" (cons (car p) (add1 (cdr p)))]
    ["D" (cons (car p) (sub1 (cdr p)))]
    ["R" (cons (add1 (car p)) (cdr p))]
    ["L" (cons (sub1 (car p)) (cdr p))]))

(define (catch-up p1 p2)
  (match* ((- (car p1) (car p2)) (- (cdr p1) (cdr p2)))
    [(2 0) (cons (add1 (car p2)) (cdr p2))]
    [(-2 0) (cons (sub1 (car p2)) (cdr p2))]
    [(0 2) (cons (car p2) (add1 (cdr p2)))]
    [(0 -2) (cons (car p2) (sub1 (cdr p2)))]
    [(2 1) (cons (add1 (car p2)) (add1 (cdr p2)))]
    [(2 -1) (cons (add1 (car p2)) (sub1 (cdr p2)))]
    [(-2 1) (cons (sub1 (car p2)) (add1 (cdr p2)))]
    [(-2 -1) (cons (sub1 (car p2)) (sub1 (cdr p2)))]
    [(1 2) (cons (add1 (car p2)) (add1 (cdr p2)))]
    [(-1 2) (cons (sub1 (car p2)) (add1 (cdr p2)))]
    [(1 -2) (cons (add1 (car p2)) (sub1 (cdr p2)))]
    [(-1 -2) (cons (sub1 (car p2)) (sub1 (cdr p2)))]
    [(2 2) (cons (add1 (car p2)) (add1 (cdr p2)))]
    [(2 -2) (cons (add1 (car p2)) (sub1 (cdr p2)))]
    [(-2 2) (cons (sub1 (car p2)) (add1 (cdr p2)))]
    [(-2 -2) (cons (sub1 (car p2)) (sub1 (cdr p2)))]
    [(_ _) p2]))

(define (draw state)
  (define lines
    (for/list ([y (in-range 10 -10 -1)])
      (apply string-append
        (for/list ([x (in-range -10 10)])
          (match (index-of state (cons x y))
            [#f "."]
            [0 "H"]
            [i (~a i)])))))
  (string-join lines "\n"))

(define (part2 input)
  (define initial-state (build-list 10 (λ (_) (cons 0 0))))
  (for*/fold ([rope initial-state] [visited (set (cons 0 0))] #:result (set-count visited))
             ([inst (in-list input)]
              [dir (in-value (first inst))]
              [_ (in-range (second inst))])
    (define new-list (list (apply-move dir (first rope))))

    (for ([i (in-range 1 10)])
      (define h (list-ref new-list (sub1 i)))
      (define t (list-ref rope i))
      (define t* (catch-up h t))
      (set! new-list (append new-list (list t*))))

    ; (printf "~a ~a\n" dir new-list)
    ; (displayln (draw new-list))

    (values new-list (set-add visited (last new-list)))))

(define (parse fname)
  (for/list ([line (in-list (file->lines fname))])
    (match-define (list dir (app string->number num)) (string-split line " "))
    (list dir num)))

; (part2 (parse "inputs/9-test2"))
; (error "x")

(module+ main
  (define input (parse "inputs/9"))
  (answer 9 1 (time (part1 input)))
  (answer 9 2 (time (part2 input)))
  (displayln "meow"))