(load "../common.scm")

(define (read-starting-positions filename)
  (call-with-input-file
    filename
    (lambda (file)
      (let* [(player-1 (begin
			 (get-string-n file (string-length "Player 1 starting position: "))
			 (get-number file)))
	     (player-2 (begin
			 (get-string-n file (string-length "Player 2 starting position: "))
			 (get-number file)))]
	(values player-1 player-2)))))

; returns a lambda that yields increasing numbers from 1 to 100, then repeating back from 1
(define (deterministic-dice)
  (let [(value -1)]
    (lambda ()
      (set! value (mod (+ value 1) 100))
      (+ value 1))))

; play a game of dirac with a magic deterministic dice that always yields increasing values
; returned value is multiplication of dice roll count and final score of opponent
(define (play-deterministic position-1 position-2 win-score)
  (let [(dice (deterministic-dice))]
    ; we subtract one from position so that we can just use "mod 10" to wrap around
    (let repeat [(player-1 (cons (- position-1 1) 0)) ; each player pair is (position . score)
		 (player-2 (cons (- position-2 1) 0)) 
		 (rolls 3)]
      (let* [(new-position (mod (+ (car player-1) (dice) (dice) (dice)) 10))
	     ; add 1 to account for off by 1 positions
	     (new-score (+ new-position (cdr player-1) 1))]
	(if (>= new-score win-score)
	  (* (cdr player-2) rolls)
	  ; swap players in next iteration
	  (repeat player-2 (cons new-position new-score) (+ 3 rolls)))))))

; state of an universe playing dirac, two players pawn position with their scores
(define-record-type state
  (fields
    ; positions are 0..9 instead of input's 1..10
    ; so we can just mod 10 while incrementing position
    (mutable position-1) 
    (mutable score-1)
    (mutable position-2) 
    (mutable score-2)))

(define (state-copy state)
  (make-state
    (state-position-1 state)
    (state-score-1 state)
    (state-position-2 state)
    (state-score-2 state)))

(define (state-equal? l r)
  (and (equal? (state-position-1 l) (state-position-1 r))
       (equal? (state-position-2 l) (state-position-2 r))
       (equal? (state-score-1 l) (state-score-1 r))
       (equal? (state-score-2 l) (state-score-2 r))))

(define (state-hash state)
  (+ (state-position-1 state)
     (* 6700417 (+ (state-position-2 state)
		   (* 6700417 (+ (state-score-1 state)
				 (* 6700417 (state-score-2 state))))))))

(define (make-state-hashtable)
  (make-hashtable state-hash state-equal?))

(define (hashtable-for-each proc ht)
  (let-values [((k v) (hashtable-entries ht))]
    (vector-for-each (lambda (i k v) (proc k v)) k v)))

(define (hashtable-empty? ht)
  (= (hashtable-size ht) 0))

; sum and amount of permutations of 3 3-sided dice
(define dirac-dice
  '((3 . 1)   ; (1 1 1)
    (4 . 3)   ; (1 1 2) (1 2 1) (2 1 1)
    (5 . 6)   ; (1 1 3) (1 2 2) (1 3 1) (2 1 2) (2 2 1) (3 1 1)
    (6 . 7)   ; (1 2 3) (1 3 2) (2 1 3) (2 2 2) (2 3 1) (3 1 2) (3 2 1)
    (7 . 6)   ; (1 3 3) (2 2 3) (2 3 2) (3 1 3) (3 2 2) (3 3 1) 
    (8 . 3)   ; (2 3 3) (3 2 3) (3 3 2)
    (9 . 1))) ; (3 3 3)

; roll dice creating as many necessary universes in the process
; returns amount of wins obtained with this call and new states
(define (dirac-roll win-score states player)
  (let-values [((get-position get-score set-position! set-score!)
		(if (= player 1) 
		  (values state-position-1 state-score-1
			  state-position-1-set! state-score-1-set!)
		  (values state-position-2 state-score-2
			  state-position-2-set! state-score-2-set!)))]
    (let [(additional-wins 0)
	  (new-states (make-state-hashtable))]
      (for-each
	(lambda (outcome)
	  (let [(sum (car outcome))
		(permutations (cdr outcome))]
	    (hashtable-for-each
	      (lambda (state amount)
		(let* [(new-position (mod (+ (get-position state) sum) 10))
		       (new-score (+ (get-score state) new-position 1))
		       (universes (* amount permutations))]
		  (if (>= new-score win-score)
		    (set! additional-wins (+ additional-wins universes))
		    (let [(new-state (state-copy state))]
		      (set-position! new-state new-position)
		      (set-score! new-state new-score)
		      (hashtable-set! new-states new-state 
				      (+ universes (hashtable-ref new-states new-state 0)))))))
	      states)))
	dirac-dice)
      (values additional-wins new-states))))

; play dirac dice starting from given positions up until the given score
; returned value is maximum number of wins for any of the two players
(define (play-dirac position-1 position-2 win-score)
  (let [(init-state (make-state-hashtable))]
    (hashtable-set! init-state (make-state (- position-1 1) 0 (- position-2 1) 0) 1)
    ; "states" tracks the amount of universes currently using a given state
    (let repeat [(states init-state) (player-1-wins 0) (player-2-wins 0) (player 1)]
      (let-values [((additional-wins new-states)
		    (dirac-roll win-score states player))]
	(set! player-1-wins (+ player-1-wins additional-wins))
	(if (hashtable-empty? new-states) ; have all universes finished playing?
	  (max player-1-wins player-2-wins)
	  ; swap players in next iteration
	  (repeat new-states player-2-wins player-1-wins (- 1 player)))))))

(let-values [((position-1 position-2) (read-starting-positions "input"))]
  (printf "part 1:~% Rolls * Opponent Score = ~a~%"
	  (play-deterministic position-1 position-2 1000))
  (printf "part 2:~% Max Winning Universes: ~a~%"
	  (play-dirac position-1 position-2 21)))