#lang racket/base

; 2 dimensionnal array
; Type inside the array can be anything

(require
  racket/vector
  racket/function
  (prefix-in srfi43: srfi/43)
  "graph.rkt")

(provide
  ; Constructors
  make-2d-array build-2d-array
  ; Accessors
  2d-array-width 2d-array-height
  2d-array-ref 2d-array-set!
  ; Other operations
  2d-array-inside?
  2d-array-fill!
  2d-array-map 2d-array-map!
  2d-array-for-each
  2d-array-map/indexes 2d-array-map!/indexes
  2d-array-for-each/indexes
  2d-array-subarray
  2d-array-clone
  2d-array-copy!
  2d-array-fold
  2d-array-fold/indexes
  2d-array-flood-fill 2d-array-flood-fill!
  2d-array-flood-map 2d-array-flood-map!
  2d-array-flood-map/indexes 2d-array-flood-map!/indexes
  2d-array-path-find
  )

; Structure
(struct 2d-array
  (width ; width
   height ; height
   items)) ; Vector of elements

; Constructors
; Make with the same value everywhere
(define (make-2d-array width height [default-value #f])
  (2d-array
    width
    height
    (make-vector (* width height) default-value)))

; Make with a procedure
; The procedure takes 2 parameters (the indexes x, y), and should return the value to set
(define (build-2d-array width height proc)
  (2d-array
    width
    height
    (build-vector
      (* width height)
      (lambda (i)
        (proc (remainder i width) (quotient i width))))))

; Get the vector index from x,y
(define (2d-array-index arr x y)
  (+ x (* y (2d-array-width arr))))

; Get the value at x,y
(define (2d-array-ref arr x y)
  (vector-ref
    (2d-array-items arr)
    (2d-array-index arr x y)))

; Set the value at x,y
(define (2d-array-set! arr x y val)
  (vector-set!
    (2d-array-items arr)
    (2d-array-index arr x y)
    val))

; Check if a coordinate is inside the bounds of the array
(define (2d-array-inside? arr x y)
  (and
    (>= x 0)
    (>= y 0)
    (< x (2d-array-width arr))
    (< y (2d-array-height arr))))

; Fill the array with a value
(define (2d-array-fill! arr val)
  (vector-fill! (2d-array-items arr) val))

; Apply a procedure to each element of a 2d array and returns a new 2d array
(define (2d-array-map arr proc)
  (2d-array
    (2d-array-width arr)
    (2d-array-height arr)
    (vector-map
      proc
      (2d-array-items arr))))

; Update the array by applying a procedure to each element
(define (2d-array-map! arr proc)
  (vector-map! proc (2d-array-items arr)))

; Apply a procedure for each element of a 2d array and discard the result
(define (2d-array-for-each arr proc)
  (void
    (vector-map
      proc
      (2d-array-items arr))))

; Apply a procedure to each element, and returns a new 2d array
; The procedure also takes the indexes: (lambda (val x y) ...)
(define (2d-array-map/indexes arr proc)
  (define width (2d-array-width arr))
  (2d-array
    width
    (2d-array-height arr)
    (srfi43:vector-map
      (lambda (i val)
        (proc
          val
          (remainder i width)
          (quotient i width)))
      (2d-array-items arr))))

; Update the array by applying a procedure to each element
; The procedure also takes the indexes: (lambda (val x y) ...)
(define (2d-array-map!/indexes arr proc)
  (define width (2d-array-width arr))
  (srfi43:vector-map!
    (lambda (i val)
      (proc
        val
        (remainder i width)
        (quotient i width)))
    (2d-array-items arr)))

; Apply a procedure for each element of a 2d array and discard the result
; The procedure also takes the indexes: (lambda (val x y) ...)
(define (2d-array-for-each/indexes arr proc)
  (define width (2d-array-width arr))
  (srfi43:vector-for-each
    (lambda (i val)
      (proc
        val
        (remainder i width)
        (quotient i width)))
    (2d-array-items arr)))

; Construct an array from another array given a bounding box
(define (2d-array-subarray arr [x 0] [y 0] [width (2d-array-width arr)] [height (2d-array-height arr)])
  (define w (min width (- (2d-array-width arr) x)))
  (define h (min height (- (2d-array-height arr) y)))
  (build-2d-array
    w
    h
    (lambda (i j)
      (2d-array-ref arr (+ i x) (+ j y)))))

; Swallow clone a 2d-array (copy the internal vector into a new vector)
; Euivalent to (2d-array-subarray arr)
(define (2d-array-clone arr)
  (define w (2d-array-width arr))
  (define h (2d-array-height arr))
  (2d-array
    w
    h
    (build-vector
      (* w h)
      (lambda (i) (vector-ref (2d-array-items arr) i)))))

; Copy an array into another
; Copy src in dest at position dest-x,dest-y
(define (2d-array-copy! dest dest-x dest-y src)
  (2d-array-for-each/indexes
    src
    (lambda (cell i j)
      ; If inside dest, copy
      (when (2d-array-inside? dest (+ dest-x i) (+ dest-y j))
        (2d-array-set! dest (+ dest-x i) (+ dest-y j) (2d-array-ref src i j))))))

; Reduce a 2d-array into a single value, version without indexes in proc
; proc is (lambda (val res))
(define (2d-array-fold arr init proc)
  (srfi43:vector-fold
    (lambda (i res val)
      (proc val res))
    init
    (2d-array-items arr)))

; Reduce a 2d-array into a single value, version with indexes in proc
; proc is (lambda (val x y res))
(define (2d-array-fold/indexes arr init proc)
  (define width (2d-array-width arr))
  (srfi43:vector-fold
    (lambda (i res val)
      (proc
        val
        (remainder i width)
        (quotient i width)
        res))
    init
    (2d-array-items arr)))

; Flood fill
; Transform an array by choosing a position, a new value, and a comparator
; Each adjacent cell that compares equaly (according to the comparator) to the value at the given position will be replaced by the new value
; Starts at the given position
; Overwrites the 2d-array
(define (2d-array-flood-fill! arr x y val [comp equal?])
  (2d-array-flood-map! arr x y (lambda (_) val) comp))

; Flood-fill
; Returns a new array
(define (2d-array-flood-fill arr x y val [comp equal?])
  (define narr (2d-array-clone arr))
  (2d-array-flood-fill! narr x y val comp)
  narr)

; Flood map
; Similar to flood fill, but instead of changing to a given value, a procedured is applied
; Version taking a proc (lambda (val))
(define (2d-array-flood-map! arr x y proc [comp equal?])
  (2d-array-flood-map!/indexes arr x y (lambda (val x y) (proc val)) comp))

; Flood map
; Returns a new array
(define (2d-array-flood-map arr x y proc [comp equal?])
  (define narr (2d-array-clone arr))
  (2d-array-flood-map! narr x y proc comp)
  narr)

; Flood map
; Version taking a proc (lambda (val x y) ...)
(define (2d-array-flood-map!/indexes arr x y proc [comp equal?])
  ; Value of the cell
  (define old-value (2d-array-ref arr x y))
  ; Already visited cells
  (define visited (make-2d-array (2d-array-width arr) (2d-array-height arr) #f))
  ; Flood fill function
  (define (flood-fill to-visit)
    ; If there are positions to visit
    (when (not (null? to-visit))
      (define pos (car to-visit))
      (define i (car pos))
      (define j (cdr pos))
      (cond
        ; Not inside the array, skip
        ((not (2d-array-inside? arr i j))
         (flood-fill (cdr to-visit)))
        ; Already visited, skip
        ((2d-array-ref visited i j)
         (flood-fill (cdr to-visit)))
        ; If value at pos if not similar to old-value, skip, and add pos to visited
        ((not (comp old-value (2d-array-ref arr i j)))
         (2d-array-set! visited i j #t)
         (flood-fill (cdr to-visit)))
        ; Else, change the value, add its adjacent cells to the list of cell to visit, and mark it as visited
        (#t
          (2d-array-set!
            arr i j
            (proc (2d-array-ref arr i j) i j))
          (2d-array-set! visited i j #t)
          (flood-fill
            (append
              (list
                (cons (+ i 1) j)
                (cons (- i 1) j)
                (cons i (+ j 1))
                (cons i (- j 1)))
              (cdr to-visit)))))))
  ; Apply the flood-fill function to the first point
  (flood-fill (list (cons x y))))

; Flood map
; Returns a new array
(define (2d-array-flood-map/indexes arr x y proc [comp equal?])
  (define narr (2d-array-clone arr))
  (2d-array-flood-map!/indexes narr x y proc comp)
  narr)

; Find a path between two points of the array
; Return #f if there is no path between the two points
; Else return the list of positions of the path, as points
; from-i from-j : the starting point
; to-i to-j : the ending point
; move-cost : a function (move-cost cell next-cell di dj)
;             taking the array, a current cell, the next cell, and the offset (di dj) to which we are moving,
;             and returning the cost to move from the current cell to the next by using the offset, or #f if the move is not possible
; heuristic-cost : a function (heuristic-cost arr i j to-i to-j)
;                  taking the array, a position and the destination position,
;                  and returning the estimated cost to move from the position to the destination, or #f if the move is not possible
(define (2d-array-path-find arr from-i from-j to-i to-j move-cost heuristic-cost)
  ; Frontier structure used inside this function
  (struct frontier (position cost))
  ; destination point
  (define to (point to-i to-j))
  ; origin point
  (define from (point from-i from-j))
  ; cost of each position
  (define position-cost (make-hash))
  ; from which position do we go to another one ? (to->from)
  (define path (make-hash))
  ; Build the path
  (define (build-path lst current)
    (if current
        (build-path (cons current lst) (hash-ref path current))
        lst))
  ; Return the next frontier from a current position and a direction
  ; Return #f if there is no way to go to the next position (or if it has already been visited with a lower cost)
  (define (next-frontier current-position next-direction)
    (define next-position (point-add current-position next-direction))
    (define i (point-x current-position))
    (define di (point-x next-direction))
    (define j (point-y current-position))
    (define dj (point-y next-direction))
    (define next-i (point-x next-position))
    (define next-j (point-y next-position))
    ; No need to compute a cost if the current or next position are outside the array
    (define move (and (2d-array-inside? arr i j)
                      (2d-array-inside? arr next-i next-j)
                      (move-cost
                        (2d-array-ref arr i j)
                        (2d-array-ref arr next-i next-j)
                        di dj)))
    (define new-cost (and move (+ move (hash-ref position-cost current-position 0))))
    (define next-cost (hash-ref position-cost next-position #f))
    (if (and new-cost (or (not next-cost) (< new-cost next-cost)))
        (begin
          (hash-set! position-cost next-position new-cost)
          (hash-set! path next-position current-position)
          (frontier next-position (+ new-cost (heuristic-cost arr next-i next-j to-i to-j))))
        #f))
  ; Get the neighbours of a given position
  (define (neighbours current-position)
    (filter
      identity
      (map
        (lambda (x) (next-frontier current-position x))
        ; Possible next directions
        (list (point 0 1) (point 0 -1) (point 1 0) (point -1 0)))))
  ; Search the path along the list of frontier points
  (define (search-path frontiers)
    (cond
      ; No element in the frontier => no path possible
      ((null? frontiers) #f)
      ; The first element of the frontier is the destination point => build and return the path
      ((point-equal? (frontier-position (car frontiers)) to)
       (build-path (list) to))
      ; Else, replace the first element of the frontier by its neighbours, sort along the cost and repeat the search
      (#t
       (search-path
          (sort
            (append
              (neighbours (frontier-position (car frontiers)))
              (cdr frontiers))
            (lambda (f1 f2) (< (frontier-cost f1) (frontier-cost f2))))))))
  ; Initialise position-cost and path tables
  (hash-set! position-cost from 0)
  (hash-set! path from #f)
  ; Search the path
  (search-path (list (frontier from 0))))
