;;; Scheme Recursive Art Contest Entry
;;;
;;; Please do not include your name or personal info in this file.
;;;
;;; Title: Fire in Tears
;;;
;;; Description:
;;;   Our world is burning!
;;;   Red flames reflect in blue tears
;;;   On an ocean of sympathy



;;; convenience functions
(define (cadr s) (car (cdr s)))
(define (caddr s) (car (cdr (cdr s))))
(define (cdddr s) (cdr (cdr (cdr s))))
(define (cadddr s) (car (cdr (cdr (cdr s)))))
(define (cddddr s) (cdr (cdr (cdr (cdr s)))))
(define (sq x) (* x x))
(define % modulo)
(define pi 3.141592653589793238)

(define (range lo hi step)
  (define (range-iter lo hi step acc)
    (if (>= lo hi) acc 
          (range-iter lo (- hi step) step (cons hi acc))))
  (range-iter lo hi step nil))

(define (reverse s)
  (define (reverse-iter s acc)
    (if (null? s) acc
      (reverse-iter (cdr s) (cons (car s) acc))))
  (reverse-iter s nil))

(define (map fn s)
  (define (map-iter fn s acc)
    (if (null? s) acc
      (map-iter fn (cdr s) (cons (fn (car s)) acc))))
  (reverse (map-iter fn s nil)))

(define (reduce fn s start)
  (if (null? s) start
    (reduce fn (cdr s) (fn start (car s)))))

(define (zip2 s t)
  (define (zip2-iter s t acc)
    (if (or (null? s) (null? t)) acc
      (zip2-iter (cdr s) (cdr t) (cons (list (car s) (car t)) acc))))
  (reverse (zip2-iter s t nil)))

(define (partition pred s)
  (define (partition-iter pred s tacc facc)
    (cond
      ((null? s) (list (reverse tacc) (reverse facc)))
      ((pred (car s)) (partition-iter pred (cdr s) (cons (car s) tacc) facc))
      (else (partition-iter pred (cdr s) tacc (cons (car s) facc)))))
  (partition-iter pred s nil nil))


;;; convert between cylindrical and rectangular coordinates
; Rectangular coordinates are (x, y, z), but the coordinate system is
; left-handed (not what you are used to). Also, +x is right, +y is up,
; and +z is INTO the screen (away from the observer).
; Cylindrical coordinates are (r, th, ze) defined as:
; r = sqrt(x**2 + z**2)
; th = arctan(z/x)
; ze = y
; This is the regular cylindrical coordinate system we all know and love.
; To get back to rectangular coordinates, we use:
; x = r*cos(th)
; y = ze
; z = r*sin(th)
(define (rect-to-cyl rect)
  (define x (getx rect))
  (define z (getz rect))
  (make-vec3 (sqrt (+ (sq x) (sq z)))
             (nparctan z x)
             (gety rect)))

(define (cyl-to-rect cyl)
  (define r (getx cyl))
  (define th (gety cyl))
  (make-vec3 (* r (npcos th))
             (getz cyl)
             (* r (npsin th))))


;;; the vector data abstraction that lets us speed up computations using Numpy!
; makes a `vec3` with 3 ndarrays or scalars; type checker getters
(define (make-vec3 x y z) (list 'vec3 x y z))
(define (vec3? x) (and (list? x) (= (length x) 4) (eq? (car x) 'vec3)))
(define (getx v) (cadr v))
(define (gety v) (caddr v))
(define (getz v) (cadddr v))

; math operations on `vec3`s
; maps binary function `fn` onto `v1` and `v2` element-wise
; scalars broadcast onto `vec3`s from the LEFT ONLY
(define-macro (map-vec3 fn v1 v2) `(if (vec3? v1)
                                    (make-vec3 (,fn (getx ,v1) (getx ,v2))
                                               (,fn (gety ,v1) (gety ,v2))
                                               (,fn (getz ,v1) (getz ,v2)))
                                    (make-vec3 (,fn ,v1 (getx ,v2))
                                               (,fn ,v1 (gety ,v2))
                                               (,fn ,v1 (getz ,v2)))))
(define (fmap-vec3 fn v1 v2) (map-vec3 fn v1 v2))

; reduce all the elements of `v` starting with `start` using binary function `fn`
; DOES NOT RETURN A `vec3`, but return type matches the type of whatever is inside `v`
(define-macro (reduce-vec3 fn start v) `(,fn (,fn (,fn ,start (getx ,v))
                                                  (gety ,v)) (getz ,v)))

(define (add v1 v2) (map-vec3 + v1 v2))
(define (sub v1 v2) (map-vec3 - v1 v2))
(define (mul v1 v2) (map-vec3 * v1 v2))
(define (dot v1 v2) (reduce-vec3 + (npzeros-like (getx v1)) (map-vec3 * v1 v2)))
(define (normalize v) (let ((norm (dot v v))) 
                        (mul (npsqrt (/ 1 (npwhere (= norm 0) 1 norm))) v)))

; operations that allow working with ndarrays and conditional masks
; interface for `np.extract`
(define (extract-arr condition arr) (if (nparray? arr)
                                      (npextract condition arr)
                                      arr))
(define (extract condition v) (fmap-vec3 extract-arr condition v))
; interface for `np.place`
(define (place condition v) (let ((r (make-vec3 (npzeros-like condition)
                                                (npzeros-like condition)
                                                (npzeros-like condition))))
                                 (fmap-vec3 (lambda (x y) (npplace x condition y)) r v)
                                 r))


;;; sphere data abstraction
; makes a sphere of radius `r` with center `c` (of type ndarray) of color `color`
(define (make-sphere r c diffuse-color emit-color mirror hardness)
  (list 'sphere r c diffuse-color emit-color mirror hardness))
; is obj a sphere?
(define (sphere? obj) (eq? (car obj) 'sphere))
; returns the radius of `sphere`
(define (sphere-radius sphere) (cadr sphere))
; returns the center of `sphere`
(define (sphere-center sphere) (caddr sphere))
; returns the diffuse color of `sphere`
(define (sphere-diffuse-color sphere) (cadddr sphere))
(define (set-diffuse-color! sphere diff-col) (set-car! (cdddr sphere) diff-col))
; returns the emit color of `sphere`
(define (sphere-emit-color sphere) (car (cddddr sphere)))
(define (set-emit-color! sphere emit-color) (set-car! (cddddr sphere) emit-color))
; returns the mirror effect of the sphere
(define (sphere-mirror sphere) (car (cdr (cddddr sphere))))
(define (set-mirror! sphere mirror) (set-car! (cdr (cddddr sphere)) mirror))
; returns the perceived hardness of the sphere
(define (sphere-hardness sphere) (caddr (cddddr sphere)))
(define (set-hardness! sphere hardness) (set-car! (cdr (cdr (cddddr sphere))) hardness))


(define FARAWAY 1e40)
(define L (make-vec3 5.0 5.0 -10.0))  ; point light position
(define E (make-vec3 0.0 1.0 -7.0))
(define w 2400)
(define h 2400)

; returns the distance `t` between the ray origin `O` and the point of tangency
; `D` is the unit vector in the direction of the ray
(define (sphere-intersect sphere O D)
  (define v (sub O (sphere-center sphere)))
  (define b (* -1. (dot D v)))
  (define c (+ (* -1. (sq (sphere-radius sphere))) (dot v v)))
  (define disc (- (sq b) c))
  (define sqtdisc (npsqrt (npmaximum disc 0)))
  (define t0 (- b sqtdisc))
  (define t1 (+ b sqtdisc))
  (define t (npwhere (> t0 0) t0 t1))
  (define intersection (& (> disc 0) (> t 0)))
  (npwhere intersection t FARAWAY))

; returns the color of a (group of) pixel(s)
; `O` is the ray origin, `D` is its direction, `t` is its length
; `scene` is a list of all the objects in the scene
; `L` is a point light source
; `bounce` controls recursion depth for reflection
; `acc` is an accumulator for tail recursion
(define (sphere-shade sphere O D t scene L bounce)
  (define M (add O (mul t D)))  ; intersection point
  (define N (mul (/ 1 (sphere-radius sphere)) (sub M (sphere-center sphere))))
  (define toL (normalize (sub L M)))  ; unit vector to light source
  (define toO (normalize (sub O M)))  ; unit vector to ray origin
  (define nudged (add M (mul 0.0001 N)))  ; avoid numerical errors
  
  ; make shadows
  (define light-distances (map (lambda (sphere) (sphere-intersect sphere nudged toL)) scene))
  (define seelight (= (sphere-intersect sphere nudged toL) (reduce npminimum light-distances FARAWAY)))

  ;;; Blinn-Phong reflection model
  ; Ambient lighting
  (define color (make-vec3 0.05 0.05 0.05))

  ; Diffuse (Lambertian) reflection
  (define amount (npmaximum (dot N toL) 0))
  (define color (add (mul (* amount seelight) ((sphere-diffuse-color sphere) M)) color))

  ; Specular reflection
  (define phong (dot N (normalize (add toL toO))))
  (define color (add (mul (* (nppower (npclip phong 0 1) (sphere-hardness sphere)) seelight)
                          (make-vec3 1 1 1))
                     color))

  ; Emit light
  (define emit (sphere-emit-color sphere))
  (if (not (eq? emit 'no-emit))
    (define color (add color (emit M))))

  ;;; reflections
  (if (= 0 bounce) color
    (begin
      (define reflected-ray (normalize (sub D (mul (* 2 (dot D N)) N))))
      (add color (mul (sphere-mirror sphere) (raytrace nudged reflected-ray scene (- bounce 1)))))))
      ; (raytrace nudged reflected-ray scene (- bounce 1) (add color acc)))))

; shader functions (allows computation of textures!)
; solid color
(define (solid-color color)
  (define (color-fn M)
    color)
  color-fn)

; makes a checkered pattern with frequency depending on `freq`
(define (checkered-color color freq)
  (define (color-fn M)
    (define checker (= (% (int (* freq (getx M))) 2) (% (int (* freq (getz M))) 2)))
    (mul checker color))
  color-fn)

; blends streaks of `color1` and `color2` on `background` (also a color (vec3))
(define (blended-color color1 color2 background)
  (define (color-fn M)
    (define M (add (* 0.2 (nprandn)) M))
    (define c1-mask (npabs (npsin (* 6 (gety M)))))
    (define c2-mask (nppower (npsin (* 6 (gety M))) 20))
    (add (add (mul c1-mask color1) (mul c2-mask color2)) (mul (- 1 c1-mask) background)))
  color-fn)

; sinusoidal waves
(define (waves-color color)
  (define (color-fn M)
    (let
      ((xsq (sq (getx M))) (ysq (sq (getz M))))
      (define waves-mask (npsin (+ xsq ysq)))
      (mul waves-mask color)))
  color-fn)

;;; ray tracing!
(define (raytrace O D scene bounce)
  ; `distances` is a Scheme list of Numpy ndarrays
  (define distances (map (lambda (sphere) (sphere-intersect sphere O D)) scene))
  (define nearest (reduce npminimum distances FARAWAY))
  (define (color-hit s-t)
    (define s (car s-t))   ; sphere
    (define t (cadr s-t))  ; distance
    (define hit (& (neq nearest FARAWAY) (= nearest t)))
    (if (npany hit)
      (begin
        (define tc (extract-arr hit t))
        (define Oc (extract hit O))
        (define Dc (extract hit D))
        (define cc (sphere-shade s Oc Dc tc scene L bounce))
        (place hit cc))
      (make-vec3 0 0 0)))
  (reduce add (map color-hit (zip2 scene distances)) (make-vec3 0 0 0)))

; projection parameters
(define r (* 2 (/ h w)))
(define x (nptile (nplinspace -2.0 2.0 w) h))
(define y (nprepeat (nplinspace (+ 1.0 r) (- 1.0 r) h) w))
(define Q (make-vec3 x y -3.))  ; the screen/canvas to be painted

; defines a 2D parametric teardrop curve with given `z` value
; takes in a `t` and outputs a vec3 of coordinates
(define (teardrop t z)
  (make-vec3 (* 2.20 (npsin t) (npsin (/ t 2))) (* 2.20 (+ (npcos t) 0.8)) z))

; make a teardrop of spheres
; set the radii and centers
(define ts (range 0 pi (/ pi 8)))
(define radii (map (lambda (t) (+ (/ pi 10) 0.1 (/ (sq (- t pi)) (* -10 pi)))) ts))
(define centers (map rect-to-cyl (map (lambda (t) (teardrop t 0.0)) ts)))

(define (rotate cs reps count acc)
  (cond
    ((null? cs) acc)
    ((or (= 0 count) (<= (getx (car cs)) 1e-8))
      (rotate (cdr cs) reps (- reps 1) acc))
    (else (rotate cs reps (- count 1) (cons (make-vec3 (getx (car cs))
                                                       (* -2 pi (/ count reps))  ; assuming th is 0.0
                                                       (getz (car cs)))
                                            acc)))))

(define (replicate-radii rs cs reps count acc)
  (cond
    ((null? rs) acc)
    ((or (= 0 count) (<= (getx (car cs)) 1e-8))
      (replicate-radii (cdr rs) (cdr cs) reps (- reps 1) acc))
    (else (replicate-radii rs cs reps (- count 1) (cons (car rs) acc)))))

; rotate the existing spheres to get all the spheres to be rendered
(define radii (replicate-radii radii centers 4 3 radii))
(define centers (map cyl-to-rect (rotate centers 4 3 centers)))

; make the scene using the centers and radii we just computed
(define scene (map (lambda (r-c) (make-sphere (car r-c) (cadr r-c) 
                                         (blended-color (make-vec3 0.7 0.0 0.0)
                                                        (make-vec3 0.0 0.7 0.0)
                                                        (make-vec3 0.5 0.1 0.0))  ; blend red and green on a gray background
                                         (solid-color (mul 0.2 (make-vec3 1.0 0.113 0.0)))  ; emit pink by default
                                         0.3 30))
                      (zip2 radii centers)))

(define (alt-blue count scene)
  (if (= 0 count) nil
    (begin (if (= 4 (% count 7)) (begin (set-diffuse-color! (car scene)
                                          (blended-color (make-vec3 0.0 0.0 0.2)
                                                         (make-vec3 0.0 0.1 0.0)
                                                         (make-vec3 0.0 0.0 0.0)))
                                        (set-emit-color! (car scene)
                                          (solid-color (mul 0.6 (make-vec3 0.0 0.894 1.0))))
                                        (set-mirror! (car scene) 0.5)
                                        (set-hardness! (car scene) 200)))
           (alt-blue (- count 1) (cdr scene)))))
(alt-blue (length scene) scene)  ; make some spheres bluish and reflective

(define (alt-mirror count scene)
  (if (= 0 count) nil
    (begin (if (= 5 (% count 6)) (begin (set-diffuse-color! (car scene)
                                          (solid-color (make-vec3 0.1 0.1 0.1)))
                                        (set-emit-color! (car scene) 'no-emit)
                                        (set-mirror! (car scene) 0.9)
                                        (set-hardness! (car scene) 200)))
           (alt-mirror (- count 1) (cdr scene)))))
(alt-mirror (length scene) scene)  ; make some spheres mirror-like

; add a wavy lake under the droplet!
(define scene (cons (make-sphere 99999 (make-vec3 0.0 -1e5 0.0) (waves-color (make-vec3 0.0 0.5 0.803))
                      'no-emit 0.9 200)
                    scene))

; uncomment to render the image using ray tracing
; The last parameter controls the reflection depth:
; `0` means no reflections, `1` means you see reflections,
; `2` means you see reflections of reflections in addition
; to reflections, etc. `2` is the perfect tradeoff between
; aesthetics and rendering time.
;;;; (define color (raytrace E (normalize (sub Q E)) scene 2))

;;;; (define img (reshape (npstack-last (getx color) (gety color) (getz color)) h w 3))  ; uncomment to make an image out of `vec3`
; (imshow img)  ; uncomment to view the generated image using Matplotlib (for debugging)
;;;; (imsave img 'img/td-bs45.png)  ; uncomment to save image (requires PIL/pillow)
; (print 'all-done!)  ; makes you feel good


(define (draw)
  ; YOUR CODE HERE
  ; this function does NOTHING
  (exitonclick))

; Please leave this last line alone.  You may add additional procedures above
; this line.
(draw)