(require pict file/convertible)
(define size 512) 
(define default-max-iterations 32)
(define default-bailout 4)

(define (generate-fractal formula initial-z bailout max-iterations)

;https://iquilezles.org/www/articles/palettes/palettes.htm
(define palette
  (for/vector #:length (add1 max-iterations)
    #:fill '(0 0 0) ([i max-iterations])
    (for/list ([a '(0.5 0.5 0.5)] [b '(0.5 0.5 0.5)] [c '(2.0 1.0 0.0)] [d '(0.50 0.20 0.25)])
      (~> i (/ max-iterations) (* 2) (* c) (+ d) (* pi 2) cos (* b) (+ a) (* 255) floor inexact->exact))))

(define (iterate func bailout z c)
  (define (normsq z) (let ([a (real-part z)] [b (imag-part z)]) (+ (* a a) (* b b))))
  (define (iterate z c i)
    (cond
      [(or (>= i max-iterations) (> (normsq z) bailout) (nan? (real-part z)) (nan? (imag-part z))) i]
      [else (iterate (func z c) c (+ i 1))]))
  (iterate z c 0))

  (define func (eval `(lambda (z c) ,formula)))
  (define init (eval `(lambda (c) ,initial-z)))
  (~> (for*/list ([b size] [a size])
        (~> (make-rectangular a b)
            exact->inexact
            (/ size)
            (* 4)
            (- 2+2i)
            ((lambda (c) (iterate func bailout (exact->inexact (init c)) c)) _)
            (vector-ref palette _)
            (apply bytes 255 _)))
      (apply bytes-append _)
      (argb-pixels->pict size)
      (convert 'png-bytes)
      (make-image)))

(match (read-args)
  [(list formula initial-z bailout max-iterations) (generate-fractal formula initial-z bailout max-iterations)]
  [(list formula initial-z bailout) (generate-fractal formula initial-z bailout default-max-iterations)]
  [(list formula initial-z) (generate-fractal formula initial-z default-bailout default-max-iterations)]
  [(list formula) (generate-fractal formula 0 default-bailout default-max-iterations)]
  [_ "Usage: formula [initial-z = 0] [bailout = 4] [max-iterations = 32]"])