(require (lib "1" "srfi"))
(require (lib "25" "srfi"))
(load "c4.scm")

; Support function for lambda-memo
; Convert a value to something that can be a hash key
(define hash-key
  (lambda (val) (cond
                  ((array? val) (array->vector val))
                  (else val))))

; Auto-memorizing functions
(define-macro lambda-memo
  (lambda (args . body)
    (let ((memo (gensym)) (hashed (gensym)) (key (gensym)))
      `(let ((,memo (make-hash-table 'equal)))
         (lambda ,args (let* ((,key (list ,@args)) (,hashed (hash-table-get ,memo ,key #f)))
                         (if ,hashed
                             ,hashed
                             (let ((,hashed (begin ,@body)))
                               (hash-table-put! ,memo ,key ,hashed)
                               ,hashed
                               )
                             )
                         )
           )
         )
      )))

; One implementation of an array copier
(define array-copy1
  (lambda (arr)
    (tabulate-array (shape 0 (array-end arr 0) 0 (array-end arr 1)) (lambda (x y) (array-ref arr x y)))))

; Another ....
(define array-copy2
  (lambda (arr)
    (apply array (shape 0 (array-end arr 0) 0 (array-end arr 1)) (vector->list (array->vector arr)))))

; Another ....
(define array-copy3
  (lambda (arr c r)
    (apply array (shape 0 c 0 r) (vector->list (array->vector arr)))))

; Some quick testing showed this should be faster by a bit
(define array-copy array-copy3)


; Find the bottom 0 cell in a column
(define free-square1
  (lambda-memo (arr c r x)
    (do ((y 0 (+ y 1)))
      ((or (= y r) (> (array-ref arr x y) 0)) (if (= y 0) #f (- y 1))))))

; Another of the above ...
(define free-square2
  (lambda (arr c r x)
    (letrec ((f (lambda (y)
                  (if (or (= y r) (> (array-ref arr x y) 0))
                      (if (= y 0) #f (- y 1))
                      (f (+ y 1))
                      )
                  )))
      (f 0)
      )))

; Using lists instead
(define free-square3
  (lambda (board x)
    (let ((col (list-ref board x)))
      (- (or (list-index (lambda (x) (> x 0)) col) (length col)) 1)
      )))

; Pick the fastest
(define free-square free-square3)

; Make a move using the list-of-lists board format
(define make-move-lol-old
  (let ((memo (make-hash-table 'equal)))
    (lambda (board col mover) (let ((hashed (hash-table-get memo (list board col mover) #f)))
                                (if hashed
                                    hashed
                                    (let* ((curcol (list-ref board col))
                                           (nonzero-i (or (list-index (lambda (x) (> x 0)) curcol) (length curcol)))
                                           )
                                      (if (= nonzero-i 0)
                                          (raise 'invalid-move)
                                          (let-values (((newcolhead newcoltail) (split-at curcol (- nonzero-i 1)))
                                                       ((newboardhead newboardtail) (split-at board col))
                                                       )
                                            (let* ((newcol (append newcolhead (list mover) (cdr newcoltail)))
                                                   (newboard (append newboardhead (list newcol) (cdr newboardtail)))
                                                   )
                                              (hash-table-put! memo (list board col mover) newboard)
                                              newboard
                                              )
                                            )
                                          )
                                      )
                                    )
                                ))))

(define make-move-lol
  (lambda (board col mover) 
    (let* ((curcol (list-ref board col))
           (nonzero-i (or (list-index (lambda (x) (> x 0)) curcol) (length curcol)))
           )
      (if (= nonzero-i 0)
          (raise 'invalid-move)
          (let-values (((newcolhead newcoltail) (split-at curcol (- nonzero-i 1)))
                       ((newboardhead newboardtail) (split-at board col))
                       )
            (let* ((newcol (append newcolhead (list mover) (cdr newcoltail)))
                   (newboard (append newboardhead (list newcol) (cdr newboardtail)))
                   )
              newboard
              )
            )
          )
      )
    )
  )

; Make a move using the array board format
(define make-move-array
  (lambda (arr c r x mover)
    (let ((y (free-square arr c r x)))
      (if y
          (let ((arr2 (array-copy arr c r)))
            (array-set! arr2 x y mover)
            arr2
            )
          (raise 'invalid-move)
          )
      )
    )
  )

; Make a move
(define make-move
  (lambda (board move mover) 
    (let-values (((newcolhead newcoltail) (split-at (list-ref board (car move)) (cadr move)))
                 ((newboardhead newboardtail) (split-at board (car move)))
                 )
      (let* ((newcol (append newcolhead (list mover) (cdr newcoltail)))
             (newboard (append newboardhead (list newcol) (cdr newboardtail)))
             )
        newboard
        )
      )
    )
  )

; Count the number of continuous items
(define count-chain
  (lambda (lst val) 
    (do ((curn 0) (maxn 0 (max maxn curn)))
      ((null? lst) maxn)
      (if (= (car lst) val)
          (set! curn (+ curn 1))
          (set! curn 0)
          )
      (set! lst (cdr lst))
      )
    )
  )

; Convert a board to a 2D array
; This should not memorize because the result will be altered destructively
(define board->array
  (lambda (board)
    (apply array (shape 0 (length board) 0 (length (car board))) (apply append board))))

; Print out an array
(define print-array
  (lambda (arr)
    (let ((c (- (array-end arr 0) 1))(r (- (array-end arr 1) 1)))
      (do ((x 0) (y 0))
        ((> y r) #f)
        (let ((n (array-ref arr x y)))
          (if (= n 0)
              (display ".")
              (display n)
              )
          )
        (if (= x c) (begin (set! x 0) (set! y (+ y 1)) (newline)) (set! x (+ x 1)))
        )
      )
    )
  )



; Count the number of continuous squares that have an open square on at least one end
;(define count-open-chain
;  (lambda-memo (lst val)
;    (do ((curn 0) (maxn 0) (start-clear #f))
;      ((null? lst) (if start-clear (max curn maxn) maxn))
;      (let ((n (car lst)))
;        (cond 
;          ((= n val) (set! curn (+ curn 1)) (if start-clear (set! maxn (max maxn curn))))
;          ((= n 0) (set! maxn (max maxn curn)) (set! curn 0) (set! start-clear #f))
;          (else 

; Subfunctions for evaluate-cell
; All verturn two values: A chain length and who the chain is for
(define evaluate-cell-left
  (lambda (arr c r x y)
    (do ((x (- x 1) (- x 1)) (chain 0 (+ chain 1)) (chain-for #f))
      ((or (< x 0) (and chain-for (not (= (array-ref arr x y) chain-for)))) (values chain chain-for))
      (if (not chain-for) (set! chain-for (array-ref arr x y)))
      )))

(define evaluate-cell-right
  (lambda (arr c r x y)
    (do ((x (+ x 1) (+ x 1)) (chain 0 (+ chain 1)) (chain-for #f))
      ((or (= x c) (and chain-for (not (= (array-ref arr x y) chain-for)))) (values chain chain-for))
      (if (not chain-for) (set! chain-for (array-ref arr x y)))
      )))

(define evaluate-cell-down
  (lambda (arr c r x y)
    (do ((y (+ y 1) (+ y 1)) (chain 0 (+ chain 1)) (chain-for #f))
      ((or (= y r) (and chain-for (not (= (array-ref arr x y) chain-for)))) (values chain chain-for))
      (if (not chain-for) (set! chain-for (array-ref arr x y)))
      )))

(define evaluate-cell-downleft
  (lambda (arr c r x y)
    (do ((x (- x 1) (- x 1)) (y (+ y 1) (+ y 1)) (chain 0 (+ chain 1)) (chain-for #f))
      ((or (< x 0) (= y r) (and chain-for (not (= (array-ref arr x y) chain-for)))) (values chain chain-for))
      (if (not chain-for) (set! chain-for (array-ref arr x y)))
      )))

(define evaluate-cell-downright
  (lambda (arr c r x y)
    (do ((x (+ x 1) (+ x 1)) (y (+ y 1) (+ y 1)) (chain 0 (+ chain 1)) (chain-for #f))
      ((or (= x c) (= y r) (and chain-for (not (= (array-ref arr x y) chain-for)))) (values chain chain-for))
      (if (not chain-for) (set! chain-for (array-ref arr x y)))
      )))

(define evaluate-cell-upleft
  (lambda (arr c r x y)
    (do ((x (- x 1) (- x 1)) (y (- y 1) (- y 1)) (chain 0 (+ chain 1)) (chain-for #f))
      ((or (< x 0) (< y 0) (and chain-for (not (= (array-ref arr x y) chain-for)))) (values chain chain-for))
      (if (not chain-for) (set! chain-for (array-ref arr x y)))
      )))

(define evaluate-cell-upright
  (lambda (arr c r x y)
    (do ((x (+ x 1) (+ x 1)) (y (- y 1) (- y 1)) (chain 0 (+ chain 1)) (chain-for #f))
      ((or (= x c) (< y 0) (and chain-for (not (= (array-ref arr x y) chain-for)))) (values chain chain-for))
      (if (not chain-for) (set! chain-for (array-ref arr x y)))
      )))

(define evaluate-cell-subfn
  (lambda (fn . args)
    (call-with-values (lambda () (apply fn args)) 
                      (lambda (chain chain-for) (cond 
                                                  ((not chain-for) (values 0 0))
                                                  ((= chain-for 1) (values chain 0))
                                                  ((= chain-for 2) (values 0 chain))
                                                  (else (values 0 0))
                                                  )
                        )
                      )))

; Evaluate a square on the board for strategic value
; Returns two values: The chains for 1 and the chains for 2
(define evaluate-cell
  (lambda (board x y)
    (let ((arr (board->array board)) 
          (c (length board)) (r (length (car board))))
      (let-values (
                   ((left1 left2) (evaluate-cell-subfn evaluate-cell-left arr c r x y))
                   ((right1 right2) (evaluate-cell-subfn evaluate-cell-right arr c r x y))
                   ((down1 down2) (evaluate-cell-subfn evaluate-cell-down arr c r x y))
                   ((downleft1 downleft2) (evaluate-cell-subfn evaluate-cell-downleft arr c r x y))
                   ((downright1 downright2) (evaluate-cell-subfn evaluate-cell-downright arr c r x y))
                   ((upleft1 upleft2) (evaluate-cell-subfn evaluate-cell-upleft arr c r x y))
                   ((upright1 upright2) (evaluate-cell-subfn evaluate-cell-upright arr c r x y))
                   )
        (list (list
               left1
               right1
               down1
               downleft1
               downright1
               upleft1
               upright1
               )
              (list
               left2
               right2
               down2
               downleft2
               downright2
               upleft2
               upright2
               )
              )))))
  
; Compile the opposing chains
(define compile-chains
  (lambda (eval-list mover)
    (match-let ((
                 (left right down downleft downright upleft upright) (list-ref eval-list (- mover 1))
                                                                     ))
      (list
       (+ left right)
       down
       (+ downleft upright)
       (+ downright upleft)
       ))))

; All valid moves (list x y) for a board
(define valid-moves
  (lambda (board)
    (remove (lambda (x) (= (cadr x) -1)) (list-tabulate (length board) (lambda (x) (list x (free-square board x)))))
    )
  )

; Compile board statistics
; Returns a dotted pair of (list 3s 2s 1s) for (mover . enemy)
(define evaluate-board
  (lambda (board mover) 
    (letrec ((possible-moves (valid-moves board))
             (all-evals (map (lambda (move) (apply evaluate-cell board move)) possible-moves))
             (chains (map (lambda (eval-data) (list (compile-chains eval-data mover) (compile-chains eval-data (if (= mover 1) 2 1)))) all-evals))
             (lengths (apply map (lambda x (apply append x)) chains))
             (calc-stats (lambda (lst 3s 2s 1s) 
                           (if (null? lst)
                               (list 3s 2s 1s)
                               (let ((n (car lst)))
                                 (cond
                                   ((= n 3) (calc-stats (cdr lst) (+ 3s 1) 2s 1s))
                                   ((= n 2) (calc-stats (cdr lst) 3s (+ 2s 1) 1s))
                                   ((= n 1) (calc-stats (cdr lst) 3s 2s (+ 1s 1)))
                                   (else    (calc-stats (cdr lst) 3s 2s 1s))
                                   )
                                 )
                               )
                           ))
             )
      (list (calc-stats (car lengths) 0 0 0) (calc-stats (cadr lengths) 0 0 0))
      )
    )
  )

; Silly scoring stuff
(define score-board
  (lambda-memo (board mover)
    (let ((stats (evaluate-board board mover)))
      (match-let ((
                   ((me-3s me-2s me-1s) (them-3s them-2s them-1s)) stats
                                                                   ))
        (cond
          ((> me-3s 0) 100) ; FTW!
          ((= them-3s 1) 99) ; "I'll take the dying boy to block"
          ((> them-3s 1) -100) ; It's a trap
          (else (- (+ (* me-2s 10) (* me-1s 2)) (+ (* them-2s 5) them-1s)))
          )
        )
      )
    )
  )

; The following pseudo-code is from http://en.wikipedia.org/wiki/Alpha-beta_pruning (Yes, I really am that lazy)
;function minimax(node, depth)
;    return alphabeta(node, depth, -∞, +∞)
;
;function alphabeta(node, depth, α, β)
;    if node is a terminal node or depth = 0
;        return the heuristic value of node
;    if the adversary is to play at node
;        foreach child of node
;            β := min(β, alphabeta(child, depth-1, α, β))
;            if α≥β
;                return α
;        return β
;    else {we are to play at node}
;        foreach child of node
;            α := max(α, alphabeta(child, depth-1, α, β))
;            if α≥β
;                return β
;        return α
(define minimax
  (lambda (board me mover depth)
    (alphabeta board me mover depth -inf.0 +inf.0))
  )

(define alphabeta
  (lambda (board me mover depth a b)
    (let ((alllines (getall board)))
      (cond
        ((not (null? (filter (lambda (l) (four-in-a-row? l me)) alllines))) 1000) ; I win, yay!
        ((not (null? (filter (lambda (l) (four-in-a-row? l (if (= me 1) 2 1))) alllines))) -1000) ; They win, boo hiss!
        ((> (apply min (map (lambda (x) (apply min x)) board)) 0) -inf.0) ; No 0s in the board (terminal) (not bloddy likely)
        ((= depth 0) (score-board board mover)) ; Bottom of the tree
        ((= mover me) (call/cc 
                       (lambda (return) 
                         (letrec ((fn (lambda (moves a b)
                                        (if (null? moves) (return a))
                                        (let ((newa (max a (alphabeta (make-move board (car moves) mover) me (if (= mover 1) 2 1) (- depth 1) a b))))
                                          (if (> newa b) (return b))
                                          (fn (cdr moves) newa b)
                                          )
                                        )
                                      ))
                           (fn (valid-moves board) a b)
                           )
                         ))
                      )
        (else (call/cc 
               (lambda (return) 
                 (letrec ((fn (lambda (moves a b)
                                (if (null? moves) (return b))
                                (let ((newb (min b (alphabeta (make-move board (car moves) mover) me (if (= mover 1) 2 1) (- depth 1) a b))))
                                  (if (> a newb) (return a))
                                  (fn (cdr moves) a newb)
                                  )
                                )
                              ))
                   (fn (valid-moves board) a b)
                   )
                 ))
              )
        ))))

(define get-move
  (lambda (board mover depth)
    (let* ((bestval -inf.0) 
           (bestmove #f)
           (helper (lambda (move)
                     (let ((v (minimax (make-move board move mover) mover (if (= mover 1) 2 1) depth)))
                       (display (list v move))(newline)
                       (if (> v bestval)
                           (begin (set! bestval v) (set! bestmove move))
                           )
                       )
                     )
                   )
           )
      (for-each helper (valid-moves board))
      (list bestval bestmove)
      )
    )
  )

(define get-move-iter
  (lambda (board mover)
    (letrec ((bestval -inf.0) 
             (bestmove #f)
             (best (make-hash-table))
             (allmoves (valid-moves board))
             (write-lock (make-semaphore 1))
             (helper (lambda (move depth)
                       (let ((x (car move))(v (minimax (make-move board move mover) mover (if (= mover 1) 2 1) depth)))
                         (printf "~a: ~a~n" x v)
                         (semaphore-wait write-lock)
                         (if (> v (hash-table-get best x -inf.0))
                              (hash-table-put! best x v) 
                             )
                         (semaphore-post write-lock)
                         )
                       )
                     )
             (helper2 (lambda (depth) 
                        (if (> depth 10) (raise 'omfg))
                        (display "Depth: ")(display depth)(newline)
                        (for-each (lambda (move) (helper move depth)) allmoves)
                        (helper2 (+ depth 1))
                        )
                      )
             (calc-thread (thread (lambda () (helper2 4))))
           )
      
      (sync/timeout 30 (thread-dead-evt calc-thread))
      (semaphore-wait write-lock)
      (kill-thread calc-thread)
      (semaphore-post write-lock) ; Single-threaded again
      (car (car (sort (hash-table-map best (lambda (k v) (list k v))) (lambda (a b) (> (cadr a) (cadr b))))))
      )
    )
  )

;(define a (board->array test-board))
;(array-set! a 5 3 2)
;(array-set! a 5 2 1)
;(array-set! a 6 1 1)
;(print-array a)
;(evaluate-cell a 7 6 4 3)


(define do-for
  (lambda (thunk n)
    (do ((i 0 (+ i 1)))
      ((= n i))
      (thunk))))

(define testfn
  (lambda (fn . args)
    (lambda ()
      ;(fn a 7 6 0 1)
      (apply fn (append args '(1)))
      (apply fn (append args '(2)))
      (apply fn (append args '(3)))
      (apply fn (append args '(4)))
      (apply fn (append args '(5)))
      (apply fn (append args '(6)))
      )
    )
  )

(define COUNT 50000)
;(time (do-for (testfn make-move a 7 6) COUNT))
;(time (do-for (testfn make-move-lol test-board) COUNT))
;(time (do-for (lambda () (array-ref a 4 4)) COUNT))
;(time (do-for (lambda () (list-ref (list-ref test-board 4) 4)) COUNT))
;(time (do-for (testfn free-square a 7 6) COUNT))
;(time (do-for (testfn free-square3 test-board) COUNT))
;(time (get-move-iter (blank-game 6 7) 1))
(play-connect4 6 7 get-move-iter getmove-human)