;;; Marching Squares Demo
;;; by Peter Salvi, 2011.11.11.

#lang racket

;;; Libraries
(require racket/gui)
(require rnrs/hashtables-6)

;;; Parameters
(define window-width 500)

;;; Default values
(define origin-x 0.0) (define origin-y 0.0)
(define width 3.0)
(define resolution 10)
(define depth 3)
(define a1 0) (define a2 0) (define a3 0) (define a4 0)
(define b1 1) (define b2 0) (define b3 1)
(define c1 0) (define c2 0)
(define d1 -1)

;;; Variables
(define squares '())
(define segments '())
(define cache (make-hashtable equal-hash equal?))

;;; Basic Maths
(define (v+ . args) (apply map + args))
(define (v* u . args) (map (lambda (x) (apply * x args)) u))
(define (interpolate a x b) (v+ (v* a (- 1 x)) (v* b x)))

;;; Algorithm

(define (evaluate p)
  (let ([cached (hashtable-ref cache p #f)])
    (or cached
        (match-let ([(list x y) p])
          (let ([result (+ (* a1 x x x) (* a2 x x y) (* a3 x y y) (* a4 y y y)
                           (* b1 x x) (* b2 x y) (* b3 y y)
                           (* c1 x) (* c2 y) d1)])
            (hashtable-set! cache p result)
            result)))))

(define (display->world p)
  (match-let ([(list x y) p])
    (list (+ origin-x (/ width -2) (* x (/ width window-width)))
          (+ origin-y (/ width -2) (* y (/ width window-width))))))

(define (world->display p)
  (match-let ([(list x y) p])
    (list (* (- x origin-x (/ width -2)) (/ window-width width))
          (* (- y origin-y (/ width -2)) (/ window-width width)))))

(define (segment->display segment)
  (list (world->display (first segment))
        (world->display (second segment))))

(define (same-sign? a b)
  (or (and (< a 0) (< b 0))
      (and (>= a 0) (>= b 0))))

(define (square-segment corners values)
  (let* ([n (count negative? values)]
         [values (if (= n 1) (map - values) values)]
         [black? (lambda (x) (if (= n 1) (<= x 0) (< x 0)))]
         [first? (lambda (prev i next)
                   (and (black? (list-ref values i))
                        (black? (list-ref values next))
                        (not (black? (list-ref values prev)))))]
         [i (cond [(first? 3 0 1) 0]
                  [(first? 0 1 2) 1]
                  [(first? 1 2 3) 2]
                  [else 3])]
         [point (lambda (i j)
                  (interpolate (list-ref corners i)
                               (/ (list-ref values i)
                                  (- (list-ref values i)
                                     (list-ref values j)))
                               (list-ref corners j)))])
    (if (= n 2)
        (list (point i (modulo (- i 1) 4))
              (point (modulo (+ i 1) 4) (modulo (+ i 2) 4)))
        (list (point i (modulo (- i 1) 4))
              (point (modulo (+ i 2) 4) (modulo (+ i 3) 4))))))

(define (ambiguous-square corners values)
  (let ([i (if (negative? (first values)) 0 1)]
        [point (lambda (i j)
                 (interpolate (list-ref corners i)
                              (/ (list-ref values i)
                                 (- (list-ref values i)
                                    (list-ref values j)))
                              (list-ref corners j)))])
    (if (negative? (evaluate (interpolate (first corners) 1/2 (third corners))))
        (list (list (point i (modulo (+ i 1) 4))
                    (point (modulo (+ i 2) 4) (modulo (+ i 1) 4)))
              (list (point i (modulo (- i 1) 4))
                    (point (modulo (- i 2) 4) (modulo (- i 1) 4))))
        (list (list (point i (modulo (- i 1) 4))
                    (point i (modulo (+ i 1) 4)))
              (list (point (modulo (+ i 2) 4) (modulo (+ i 1) 4))
                    (point (modulo (- i 2) 4) (modulo (- i 1) 4)))))))

(define (compute-square top-left bottom-right depth)
  (let ([top-right (list (first bottom-right) (second top-left))]
        [bottom-left (list (first top-left) (second bottom-right))])
    (when (> depth 0)
      (set! squares
            (append (map segment->display
                         `((,top-left ,top-right)
                           (,top-right ,bottom-right)
                           (,bottom-right ,bottom-left)
                           (,bottom-left ,top-left)))
                    squares))
      (let* ([corners (list top-left top-right bottom-right bottom-left)]
             [values (map evaluate corners)])
        (unless (or (for/and ([x values]) (< x 0))
                    (for/and ([x values]) (>= x 0)))
          (if (> depth 1)
              (let* ([top-half (interpolate top-left 1/2 top-right)]
                     [left-half (interpolate top-left 1/2 bottom-left)]
                     [right-half (interpolate top-right 1/2 bottom-right)]
                     [center (interpolate left-half 1/2 right-half)]
                     [bottom-half (interpolate bottom-left 1/2 bottom-right)])
                (compute-square top-left center (- depth 1))
                (compute-square top-half right-half (- depth 1))
                (compute-square left-half bottom-half (- depth 1))
                (compute-square center bottom-right (- depth 1)))
              (set! segments
                    (if (and (same-sign? (first values) (third values))
                             (same-sign? (second values) (fourth values)))
                        (append (map segment->display
                                     (ambiguous-square corners values))
                                segments)
                        (cons (segment->display (square-segment corners values))
                              segments)))))))))

(define (recompute canvas)
  (set! squares '())
  (set! segments '())
  (hashtable-clear! cache)
  (let ([point (lambda (i j)
                 (display->world (list (* i (/ window-width resolution))
                                       (* j (/ window-width resolution)))))])
    (for/list ([i (in-range resolution)])
      (for/list ([j (in-range resolution)])
        (compute-square (point i j) (point (+ i 1) (+ j 1)) depth))))
  (send canvas refresh))

;;; Graphics

(define (draw-segments segments dc)
  (for-each (lambda (segment)
              (match-let ([(list (list x1 y1) (list x2 y2)) segment])
                (send dc draw-line x1 y1 x2 y2)))
            segments))

(define (draw canvas dc)
  (send dc set-pen "BLACK" 0 'solid)
  (draw-segments squares dc)
  (send dc set-pen "RED" 2 'solid)
  (draw-segments segments dc))

;;; GUI

(let* ([frame (new frame% [label "Marching Squares"])]
       [vbox (new vertical-pane% [parent frame])]
       [canvas (new canvas% [parent vbox]
                    [min-width window-width] [stretchable-width #f]
                    [min-height window-width] [stretchable-height #f]
                    [paint-callback draw])]
       [hbox1 (new horizontal-pane% [parent vbox])]
       [hbox2 (new horizontal-pane% [parent vbox])]
       [hbox3 (new horizontal-pane% [parent vbox])])
  (let-syntax ([add-option
                (syntax-rules ()
                  [(_ var a-label box)
                   (new text-field% [parent box] [label a-label]
                        [init-value (number->string var)]
                        [callback
                         (lambda (t e)
                           (let ([n (string->number (send t get-value))])
                             (when n (set! var n))))])])])
    (add-option origin-x "Origin:" hbox1)
    (add-option origin-y "," hbox1)
    (add-option width "Width:" hbox1)
    (add-option resolution "Resolution:" hbox1)
    (add-option depth "Depth:" hbox1)
    (add-option a1 "Function:" hbox2)
    (add-option a2 "x3+" hbox2)
    (add-option a3 "x2y+" hbox2)
    (add-option a4 "xy2+" hbox2)
    (new message% [parent hbox2] [label "y3+"])
    (add-option b1 #f hbox3)
    (add-option b2 "x2+" hbox3)
    (add-option b3 "xy+" hbox3)
    (add-option c1 "y2+" hbox3)
    (add-option c2 "x+" hbox3)
    (add-option d1 "y+" hbox3)
    (new message% [parent hbox3] [label "=0"])
    (new button% [parent vbox] [label "Redraw"] [stretchable-width #t]
         [callback (lambda (b e) (recompute canvas))])
    (recompute canvas)
    (send frame show #t)))