(library (mandelbrot)
  (export make-ppm mandelbrot!
          steps
          make-rgb rgb-r rgb-g rgb-b
          number->rgb)
  (import (chezscheme))

(define-record-type rgb (fields r g b))
(define (number->rgb n)
  (assert (and (integer? n) (<= 0 n)))
  (make-rgb (bitwise-and #xff (bitwise-arithmetic-shift-right n 16))
            (bitwise-and #xff (bitwise-arithmetic-shift-right n 8))
            (bitwise-and #xff (bitwise-arithmetic-shift-right n 0))))
(define rgb+
  (case-lambda
    [() (make-rgb 0 0 0)]
    [(a) a]
    [(a b) (make-rgb (+ (rgb-r a) (rgb-r b))
                     (+ (rgb-g a) (rgb-g b))
                     (+ (rgb-b a) (rgb-b b)))]
    [rgbs (fold-left rgb+ (rgb+) rgbs)]))

(define (rgb-mean . rgbs)
  (if (null? rgbs)
      (rgb+)
      (let ([sum (apply rgb+ rgbs)]
            [len (length rgbs)])
        (make-rgb (fx/ (rgb-r sum) len)
                  (fx/ (rgb-g sum) len)
                  (fx/ (rgb-b sum) len)))))

;; proc : (x y) -> rgb
(define (make-ppm path width height proc)
  (call-with-port (open-file-output-port path (file-options no-fail))
    (lambda (p)
      (define header (format "P6\n~a ~a\n255\n" width height))
      (put-bytevector p (string->utf8 header))
      (do ([x 0 (mod (+ 1 x) width)]
           [y 0 (if (= (+ 1 x) width) (+ 1 y) y)])
          [(= y height)]
        (let ([color (proc x y)])
          (put-u8 p (rgb-r color))
          (put-u8 p (rgb-g color))
          (put-u8 p (rgb-b color))))))
  'done)


(define (lerp i min max)
  (+ min (* i (- max min))))

(define steps 100)
(define power 2)
(define samples-per-pixel 4)

(define tau (* 4 (acos 0)))

(define (rotate-about origin p angle)
  (+ (* (- p origin)
      (make-rectangular (cos angle) (sin angle)))
     origin))

;; let's rotate the sample?
(define (mandelbrot! path width height color xmin xmax ymin ymax)
  (define pixel-width  (exact->inexact (/ (- xmax xmin) width)))
  (define pixel-height (exact->inexact (/ (- ymax ymin) height)))
  (define sample-rotation (/ tau samples-per-pixel))
  (define (sample-vals center)
    (let ([top (+ center (make-rectangular 0 (- (/ pixel-height 2))))])
      (map (lambda (i) (rotate-about center top (* i sample-rotation)))
           (iota samples-per-pixel))))
  (make-ppm path width height
    (lambda (x y)
      (define ix (lerp (/ x width)  xmin xmax))
      (define iy (lerp (/ y height) ymin ymax))
      (define center (exact->inexact (make-rectangular ix iy)))
      (define (sample p)
        (define (step z) (+ (expt z power) p))
        (let loop ([i 0]
                   [z (step 0)])
          (cond
           [(= i steps) (color -1)]
           [(> (magnitude z) 2) (color i)]
           [else (loop (+ 1 i) (step z))])))
      (define samples (map sample (sample-vals center)))
      (apply rgb-mean samples)))))

(top-level-program
 (import (chezscheme)
         (mandelbrot))

(define (zoom path-fmt x y width steps)
  (for-each
   (lambda (i)
     (let ([idx (+ 1 i)]
           [i (+ 1 i)])
       (mandelbrot!
        (format path-fmt idx)
        width width
        (/ (- x 0.5) i) (/ (+ x 0.5) i)
        (/ (- y 0.5) i) (/ (+ y 0.5) i))))
   (iota steps)))

(define (cli-fail)
  (format (standard-error-port 'line (current-transcoder))
          "usage: ./mandelbrot.ss path width height [color:hex*]\n")
  (exit 1))

(define args (command-line-arguments))

(unless (<= 3 (length args))
  (cli-fail))

(define path   (list-ref args 0))
(define width  (list-ref args 1))
(define height (list-ref args 2))
(set! width (string->number width))
(set! height (string->number height))

(unless (and (integer? width) (integer? height) (positive? width) (positive? height))
  (cli-fail))

(define (color i)
  (define c (floor (* 256 (/ (+ 1 i) steps))))
  (make-rgb c c c))

(when (< 3 (length args))
  (let* ([stripes
          (let* ([colors (cdddr args)]
                 [nums (map (lambda (n) (string->number n 16)) colors)]
                 [_ (unless (for-all (lambda (x) (and (integer? x) (<= 0 x))) nums)
                      (cli-fail))]
                 [colors (map number->rgb nums)])
            (list->vector colors))]
         [stripe-count (vector-length stripes)])
    (set! color (lambda (i) (vector-ref stripes (mod i stripe-count))))))

(mandelbrot! path width height color -2.2 1.1 -1 1))