#lang racket

(struct vec [x y z] #:transparent)
(define (vec->list v) (list (vec-x v) (vec-y v) (vec-z v)))

(define (vec+ . vecs)
  (define x (apply + (map vec-x vecs)))
  (define y (apply + (map vec-y vecs)))
  (define z (apply + (map vec-z vecs)))
  (vec x y z))

(define lines (string-split (string-trim (file->string "input.12.txt")) "\n"))
(define (parse-line line)
  (map string->number (regexp-match* #rx"[0-9\\-]+" line)))
(define input (map (compose (curry apply vec) parse-line) lines))

(define (bound val)
  (cond [(negative? val) -1]
        [(positive? val) 1]
        [else 0]))

(define (gravity-axis val1 val2)
  (bound (- val2 val1)))

(define (gravity-axis-all lop pos)
  (apply + (map (curry gravity-axis pos) lop)))

(define (gravity pos1 pos2)
  (apply vec (map gravity-axis (vec->list pos1) (vec->list pos2))))

(define (gravity-all lop pos)
  (apply vec+ (map (curry gravity pos) lop)))

(define (initial-vels lop)
  (map (lambda (x) (vec 0 0 0)) lop))

(define (sim-step lop [lov (initial-vels lop)])
  (define loa (map (curry gravity-all lop) lop))
  (define new-lov (map vec+ lov loa))
  (define new-lop (map vec+ lop new-lov))
  (list new-lop new-lov))

(define (energy v)
  (apply + (map abs (vec->list v))))

(define final-state (for/foldr ([state (list input (initial-vels input))])
                      ([i (in-range 1000)])
                      (apply sim-step state)))
;; Functional heckcode time
;; (Sorry lol)
(apply + (apply map * (map (curry map energy) final-state)))

(define (sim-step-axis lop [lov (map (lambda (x) 0) lop)])
  (define loa (map (curry gravity-axis-all lop) lop))
  (define new-lov (map + lov loa))
  (define new-lop (map + lop new-lov))
  (list new-lop new-lov))

(define (find-repeat-axis initial-lop [lop initial-lop] [lov (map (lambda (x) 0) lop)] [acc 1])
  (define next-state (sim-step-axis lop lov))
  (cond
    [(and (equal? initial-lop (first next-state)) (andmap zero? (second next-state))) acc]
    [else (find-repeat-axis initial-lop (first next-state) (second next-state) (add1 acc))]))

(define repeat-x (find-repeat-axis (map vec-x input)))
(define repeat-y (find-repeat-axis (map vec-y input)))
(define repeat-z (find-repeat-axis (map vec-z input)))
(lcm repeat-x repeat-y repeat-z)