;; The first three lines of this file were inserted by DrRacket. They record metadata
;; about the language level of this file in a form that our tools can easily process.
#reader(lib "htdp-advanced-reader.ss" "lang")((modname connect-four) (read-case-sensitive #t) (teachpacks ()) (htdp-settings #(#t constructor repeating-decimal #t #t none #f ())))
(require 2htdp/universe)
(require 2htdp/image)

;; A board is a
;;  (make-board string nat nat list-of-columns)
(define-struct board (turn width height columns))

;; A column is a list-of-pieces

;; A piece is a string["red" or "black" or "blank"]
;; (but the "blank" case isn't used for a board)

;; ----------------------------------------
;; Drawing a board

(define RECT-SIDE 30)
(define CIRC-RADIUS 12)

;; make-space : string -> image
(define (make-space color)
  (overlay (circle CIRC-RADIUS 'solid color)
           (rectangle RECT-SIDE RECT-SIDE 'solid "yellow")))

(define empty-space (make-space "white"))
(define red-space (make-space "red"))
(define black-space (make-space "black"))

;; draw-board : board -> image
(define (draw-board b)
  (draw-columns (board-columns b) (board-height b)))

;; draw-columns : list-of-column nat -> image
(define (draw-columns columns height)
  (cond
    [(empty? columns) empty-image]
    [(cons? columns) 
     (beside (draw-column (first columns) height)
             (draw-columns (rest columns) height))]))

;; draw-column : list-of-piece nat -> image
(define (draw-column col height)
  (draw-upsidedown-column (reverse col) height))

;; draw-upsidedown-column : list-of-piece nat -> image
;;  height is the remaining height of column
(define (draw-upsidedown-column col height)
  (cond
    [(empty? col) (stack-copies empty-space height)]
    [(cons? col) 
     (above (draw-upsidedown-column (rest col) (sub1 height))
            (make-space (first col)))]))

;; stack-copies : image nat -> image
(define (stack-copies i n)
  (cond
    [(zero? n) empty-image]
    [else (above i (stack-copies i (sub1 n)))]))

(check-expect (stack-copies red-space 3)
              (above red-space red-space red-space))

(check-expect (draw-column empty 2)
              (above empty-space empty-space))
(check-expect (draw-column (list "red") 2)
              (above empty-space red-space))
(check-expect (draw-column (list "red" "black") 2)
              (above red-space black-space))

(check-expect (draw-board (make-board "red" 2 2 (list empty empty)))
              (above (beside empty-space empty-space)
                     (beside empty-space empty-space)))
(check-expect (draw-board (make-board "red" 2 2 (list (list "black") empty)))
              (above (beside empty-space empty-space)
                     (beside black-space empty-space)))
(check-expect (draw-board (make-board "red" 2 2 (list empty (list "red" "black"))))
              (above (beside empty-space red-space)
                     (beside empty-space black-space)))
(check-expect (draw-board (make-board "red" 2 4 (list empty (list "red" "black"))))
              (above (beside empty-space empty-space)
                     (beside empty-space empty-space)
                     (beside empty-space red-space)
                     (beside empty-space black-space)))

;; ----------------------------------------
;; Adding a piece

;; add-piece : piece board nat -> board
;;  assume that there's space, and that the column exists
(define (add-piece p b pos)
  (make-board (switch-player (board-turn b))
              (board-width b)
              (board-height b)
              (add-piece-to-columns (board-columns b)
                                    p
                                    pos)))

;; switch-player : string -> string
(define (switch-player s)
  (cond
    [(string=? s "red") "black"]
    [else "red"]))

;; add-piece-to-columns : list-of-column piece nat -> list-of-columns
(define (add-piece-to-columns cols p pos)
  (cond
    [(empty? cols) (error "huh?")]
    [(cons? cols)
     (if (zero? pos)
         (cons (add-piece-to-column (first cols) p)
               (rest cols))
         (cons (first cols)
               (add-piece-to-columns (rest cols) p (sub1 pos))))]))

;; add-piece-to-column : list-of-piece p -> list-of-piece
(define (add-piece-to-column col p)
  (cons p col))

(check-expect (add-piece "red"
                         (make-board "red" 2 2 (list empty empty))
                         0)
              (make-board "black" 2 2 (list (list "red") empty)))
(check-expect (add-piece "black"
                         (make-board "black" 3 2 (list empty (list "red") empty))
                         1)
              (make-board "red" 3 2 (list empty (list "black" "red") empty)))

;; ----------------------------------------

;; column-full? : board nat -> boolean
;;   Assume that n is an ok column number
(define (column-full? b n)
  (= (board-height b) 
     (length (list-ref (board-columns b) n))))

(check-expect (column-full? (make-board "red" 2 2 (list empty empty)) 0)
              false)
(check-expect (column-full? (make-board "red" 2 2 (list (list "red" "black") empty)) 0)
              true)
(check-expect (column-full? (make-board "red" 2 2 (list (list "red" "black") empty)) 1)
              false)
(check-expect (column-full? (make-board "red" 2 3 (list (list "red" "black") empty)) 0)
              false)

;; ----------------------------------------

;; wins? : board piece -> boolean
(define (wins? b piece)
  (or (wins-in-column? b piece)
      (wins-in-row? b piece)
      (wins-diagonal? b piece)))

;; wins? : board piece -> boolean
(define (wins-in-column? b piece)
  (ormap (lambda (col)
           (wins-in-seq? col piece 0))
         (board-columns b)))

;; wins-in-seq? : list-of-pieces piece num -> boolean
;;  already-seen is the number of consecutive matching
;;  pieces that we've already seen
(define (wins-in-seq? col piece already-seen)
  (cond
    [(empty? col) (= already-seen 4)]
    [else
     (or (= already-seen 4)
         (wins-in-seq? (rest col) 
                       piece 
                       (if (string=? (first col) piece)
                           (add1 already-seen)
                           0)))]))

(check-expect (wins-in-column? small-empty-board "red")
              false)
(check-expect (wins-in-column? small-empty-board "black")
              false)
(check-expect (wins-in-column? (make-board "black" 2 4 (list (list "red" "red" "red" "red") 
                                                             (list "black" "black" "black")))
                               "red")
              true)
(check-expect (wins-in-column? (make-board "red" 3 4 (list 
                                                      (list "red" "red" "red") 
                                                      (list "black" "black" "black" "black")
                                                      (list "red")))
                               "black")
              true)


(define (wins-in-row? b piece)
  (local [(define upsidedown-columns (map reverse (board-columns b)))]
    (wins-in-a-row? upsidedown-columns piece)))

;; wins-in-a-row? : list-of-column piece -> boolean
(define (wins-in-a-row? cols piece)
  (cond
    [(andmap empty? cols) false]
    [else (or (wins-in-seq? (extract-first-row cols) piece 0)
              (wins-in-a-row? (drop-first-row cols) piece))]))

;; extract-first-row : list-of-columns -> list-of-piece
(define (extract-first-row cols)
  (map (lambda (col)
         (cond
           [(empty? col) "blank"]
           [else (first col)]))
       cols))

(check-expect (extract-first-row (list (list "red") empty))
              (list "red" "blank"))
(check-expect (extract-first-row (list (list "red") (list "black" "red")))
              (list "red" "black"))

;; drop-first-row : list-of-columns -> list-of-columns
(define (drop-first-row cols)
  (map (lambda (col)
         (cond
           [(empty? col) empty]
           [else (rest col)]))
       cols))

(check-expect (drop-first-row (list (list "red") empty))
              (list empty empty))
(check-expect (drop-first-row (list (list "red") (list "black" "red")))
              (list empty (list "red")))

(check-expect (wins-in-row? small-empty-board "red")
              false)
(check-expect (wins-in-row? small-empty-board "black")
              false)
(check-expect (wins-in-row? (make-board "black" 4 2 (list (list "black" "red")
                                                          (list "black" "red")
                                                          (list "black" "red")
                                                          (list         "red")))
                            "red")
              true)
(check-expect (wins-in-row? (make-board "red" 4 3 (list (list "red" "black" "red")
                                                        (list "red" "black" "red")
                                                        (list       "black" "red")
                                                        (list       "black" "black")))
                            "black")
              true)

(define (wins-diagonal? b piece)
  (local [(define uc (map reverse (board-columns b)))]
    (or (wins-in-slash? uc piece)
        (wins-in-slash? (reverse uc) piece))))

;; wins-in-slash? : list-of-columns piece -> bool
(define (wins-in-slash? uc piece)
  (cond
    [(empty? uc) false]
    [(andmap empty? uc) false]
    [else
     (or (wins-in-seq? (extract-bl-diagonal uc) piece 0)
         (wins-in-slash? (rest uc) piece)
         (wins-in-slash? (drop-first-row uc) piece))]))

;; extract-bl-diagonal : list-of-columns -> list-of-piece
(define (extract-bl-diagonal louc)
  (cond
    [(empty? louc) empty]
    [else (cons (get-first-piece (first louc))
                (extract-bl-diagonal (drop-first-row (rest louc))))]))

(check-expect (extract-bl-diagonal (list
                                    empty
                                    (list "red" "black")
                                    (list "black" "red" "red")
                                    empty))
              (list "blank" "black" "red" "blank"))
          
;; get-first-piece : list-of-piece -> piece[maybe "blank"]
(define (get-first-piece lop)
  (cond
    [(empty? lop) "blank"]
    [else (first lop)]))

(check-expect (wins-in-slash? (list)
                              "red")
              false)
(check-expect (wins-in-slash? (list empty)
                              "red")
              false)

(check-expect (wins-in-slash? (list empty empty empty empty)
                              "red")
              false)

(check-expect (wins-diagonal? (make-board "black" 4 4 (list empty empty empty empty))
                              "red")
              false)

(check-expect (wins-diagonal? (make-board "black" 4 4
                                          (list (list                 "red" "red")
                                                (list               "red" "black")
                                                (list       "red" "black" "black")
                                                (list "red" "black" "black" "red")))
                              "red")
              true)
(check-expect (wins-diagonal? (make-board "black" 4 8
                                          (list (list                     "red" "black")
                                                (list              "red"   "red" "black")
                                                (list       "red" "black" "black" "black")
                                                (list "red" "black" "black" "black" "red")))
                              "red")
              true)
(check-expect (wins-diagonal? (make-board "black" 6 8
                                          (list empty
                                                (list                     "red" "black")
                                                (list              "red"   "red" "black")
                                                (list       "red" "black" "black" "black")
                                                (list "red" "black" "black" "black" "red")
                                                empty))
                              "red")
              true)

(define red-wins-board
  (make-board "black" 6 8
              (list empty
                    (list "red" "black" "black" "black" "red")
                    (list       "red" "black" "black" "black")
                    (list              "red"   "red" "black")
                    (list                     "red" "black")
                    empty)))

(check-expect (wins-diagonal? red-wins-board
                              "red")
              true)

;; ----------------------------------------

;; handle-mouse : board num num string -> board
(define (handle-mouse b x y kind)
  (cond
    [(string=? kind "button-down")
     (if (wins? b (switch-player (board-turn b)))
         b
         (handle-click b x y))]
    [else b]))

;; handle-click : board num num -> board
(define (handle-click b x y)
  (local [(define col (quotient x RECT-SIDE))]
    (if (column-full? b col)
        b
        (add-piece (board-turn b) b col))))

(define small-empty-board (make-board "red" 2 2 (list empty empty)))

(check-expect (handle-mouse small-empty-board
                            1 2
                            "button-down")
              (add-piece "red" small-empty-board 0))
(check-expect (handle-mouse small-empty-board
                            (+ RECT-SIDE 1) 25
                            "button-down")
              (add-piece "red" small-empty-board 1))
(check-expect (handle-mouse small-empty-board
                            1 2
                            "move")
              small-empty-board)
(check-expect (handle-mouse (make-board "red" 2 2 (list (list "black" "red") empty))
                            1 2
                            "button-down")
              (make-board "red" 2 2 (list (list "black" "red") empty)))

(check-expect (handle-mouse red-wins-board
                            1 2
                            "button-down")
              red-wins-board)

;; ----------------------------------------
;; Extend `draw-board' to show a winner

(define (draw-board-with-winner b)
  (local [(define image (draw-board b))
          (define (add-winner msg color)
            (overlay (text msg 12 color) image))]
    (cond
      [(wins? b "red") (add-winner "Red Wins!""red")]
      [(wins? b "black") (add-winner "Black Wins!" "black")]
      [else image])))

;; ----------------------------------------

;; We got here by a simple try-to-move that looks for an immediate
;; win, then a try-not-to-lose that looks to avoid for an immediate loss
;; after a move. But this version doesn't quite work, because we've
;; lost the distinction between "don't care" and "definitely don't
;; make this move".

;; try-to-move-as : board nat piece (board piece -> boolean) -> board-or-false
;;   Assume it's p's turn, look for a move starting with
;;   the `pos'th column, keep it if `keep?' says to
(define (try-to-move-as b pos p keep?)
  (cond
    [(= pos (board-width b)) false]
    [else (cond
            [(column-full? b pos)
             (try-to-move-as b (add1 pos) p keep?)]
            [else
             (local [(define new-b (add-piece p b pos))]
               (if (keep? new-b p)
                   new-b
                   (try-to-move-as b (add1 pos) p keep?)))])]))

;; try-to-win-as : board piece nat -> board-or-false
;;   Assume it's p's turn, look for a winning move
(define (try-to-win-as b p depth)
  (try-to-move-as b 0 p (lambda (new-b p)
                          (cond
                            [(wins? new-b p) true]
                            [(zero? depth) false]
                            [else
                             (local [(define new-new-b
                                       (try-to-win-as new-b 
                                                      (switch-player p)
                                                      (sub1 depth)))]
                               (cond
                                 [(not (board? new-new-b)) true]
                                 [(wins? new-new-b (switch-player p))
                                  false]
                                 [else false]))]))))
                                    

;; try-to-win : board -> board-or-false
;;  Look for a winning move up to d steps ahead
(define (try-to-win b)
  (try-to-win-as b "black" 2))

(check-expect (try-to-win-as (make-board "black" 4 4
                                         (list empty
                                               empty
                                               empty
                                               empty))
                             "black"
                             0)
              false)
(check-expect (try-to-win-as (make-board "black" 4 4
                                         (list (list "red")
                                               empty
                                               (list "red" "red" "red")
                                               (list "black" "black" "black")))
                             "black"
                             0)
              (make-board "red" 4 4
                          (list (list "red")
                                empty
                                (list "red" "red" "red")
                                (list "black" "black" "black" "black"))))
(check-expect (try-to-win-as (make-board "black" 4 4
                                         (list (list "black" "black" "black" "red")
                                               (list "red" "red" "red")
                                               empty
                                               empty))
                             "black"
                             0)
              false)
(check-expect (try-to-win-as (make-board "black" 5 4
                                         (list (list "black")
                                               empty
                                               (list "red" "black")
                                               (list "red" "red")
                                               (list "red" "black")))
                             "black"
                             0)
              (make-board "red" 5 4
                          (list (list "black")
                                (list "black")
                                (list "red" "black")
                                (list "red" "red")
                                (list "red" "black"))))

(check-expect (try-to-win (make-board "black" 4 4
                                      (list empty
                                            empty
                                            empty
                                            empty)))
              false)

;; auto-move-left : board num -> board
;;  Assuming that it's black's turn
(define (auto-move-left b try-pos)
  (cond
    [(column-full? b try-pos)
     (auto-move-left b (add1 try-pos))]
    [else (add-piece "black" b try-pos)]))

;; auto-move-random : board -> board
;;  Assuming that it's black's turn
(define (auto-move-random b)
  (local [(define try-pos (random (board-width b)))]
    (cond
      [(column-full? b try-pos)
       (auto-move-random b)]
      [else (add-piece "black" b try-pos)])))

(define (auto-move b)
  (local [(define maybe-b (try-to-win b))]
    (if (board? maybe-b)
        maybe-b
        (auto-move-random b))))

;; ----------------------------------------

;; board-full? : board -> bool
(define (board-full? b)
  (ormap (lambda (col)
           (< (length col) (board-height b)))                          
         (board-columns b)))

;; handle-mouse-and-auto-move : board num num string -> board
;;  Handles the mouse; if the result is black's
;;  turn and red didn't win, then move for black
;;  assuming the board isn't full
(define (handle-mouse-and-auto-move b x y key)
  (local [(define new-b (handle-mouse b x y key))]
    (cond
      [(and (string=? "black" (board-turn new-b))
            (not (wins? new-b "red"))
            (board-full? b))
       (auto-move new-b)]
      [else new-b])))

(check-expect (handle-mouse-and-auto-move red-wins-board 0 0 "move")
              red-wins-board)

;; ----------------------------------------

(big-bang
 (make-board "red" 7 6 (list empty empty empty empty empty empty empty))
 [to-draw draw-board-with-winner]
 [on-mouse handle-mouse-and-auto-move])


;; To do:
;;   - computer player