(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))))

;; 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 (mandelbrot! path width height color xmin xmax ymin ymax)
  (make-ppm path width height
    (lambda (x y)
      (define ix (lerp (/ x width)  xmin xmax))
      (define iy (lerp (/ y height) ymin ymax))
      (define c (exact->inexact (make-rectangular ix iy)))
      (define (step z) (+ (expt z power) c))
      (let loop ([i 0]
                 [z (step 0)])
        (cond
         [(= i steps) (color -1)]
         [(> (magnitude z) 2) (color i)]
         [else (loop (+ 1 i) (step z))]))))))

(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))