#lang racket/gui

(require racket/draw)

(struct color (r g b))
(struct vec (x y z))
(struct object (is get-intersect get-normal get-color get-surface-point))

(define light (vec -1 -1 -0.75))
(define ambiant-color (color 15 15 30))
(define bg-color (color 70 70 130))
(define (make-scene)
  (list (new-sphere (vec 0 0 500) 100 (color 0 200 200))
    (new-sphere (vec -100 -50 500) 100 (color 250 200 100))
    (new-sphere (vec 200 150 500) 100 (color 0 0 200))
    (new-sphere (vec -200  200 300) 25 (color 100 100 100))
    (new-sphere (vec -200  150 300) 25 (color 100 100 100))
    (new-sphere (vec -200  100 300) 25 (color 100 100 100))
    (new-sphere (vec -200   50 300) 25 (color 100 100 100))
    (new-sphere (vec -200    0 300) 25 (color 100 100 100))
    (new-sphere (vec -200  -50 300) 25 (color 100 100 100))
    (new-sphere (vec -200 -100 300) 25 (color 100 100 100))
    (new-sphere (vec -200 -150 300) 25 (color 100 100 100))
    (new-sphere (vec -200 -200 300) 25 (color 100 100 100))
    (new-plane  (vec 0 200.01 0) (vec 0 1 0)
            (color 200 0 200) (color 130 0 130))
    ))

(define (color-op col f)
  (color (f (color-r col)) (f (color-g col)) (f (color-b col))))

(define (color-mul col m)
  (color-op col (lambda [x] (* x m))))

(define (color-add c1 c2)
  (color (+ (color-r c1) (color-r c2))
     (+ (color-g c1) (color-g c2))
     (+ (color-b c1) (color-b c2))))

(define (norm pt)
  (sqrt (+ (expt (vec-x pt) 2)
       (expt (vec-y pt) 2)
       (expt (vec-z pt) 2))))

(define (normalized pt)
  (let ([n (norm pt)])
    (vec (/ (vec-x pt) n)
     (/ (vec-y pt) n)
     (/ (vec-z pt) n))))

(define (vec-op op)
  (lambda [p1 p2]
    (vec (op (vec-x p1) (vec-x p2))
     (op (vec-y p1) (vec-y p2))
     (op (vec-z p1) (vec-z p2)))))

(define vec-add (vec-op +))
(define vec-sub (vec-op -))

(define (vec-mul p1 m)
  (vec (* (vec-x p1) m)
       (* (vec-y p1) m)
       (* (vec-z p1) m)))

(define (dot-product v1 v2)
  (+ (* (vec-x v1) (vec-x v2))
     (* (vec-y v1) (vec-y v2))
     (* (vec-z v1) (vec-z v2))))

(define (solve-quadratic a b c)
  (let ([discr (- (* b b) (* 4 a c))])
    (cond ((< discr 0) '())
      ((= discr 0) (list (* -0.5 (/ b a))))
      (#t (let* ([q (if (> b 0)
                (* -0.5 (+ b (sqrt discr)))
                (* -0.5 (- b (sqrt discr))))]
             [x0 (/ q a)]
             [x1 (/ c q)])
        (if (< x0 x1)
            (list x0 x1)
            (list x1 x0)))))))

(define (new-sphere pos radius col)
  (define (get-intersect orig dir)
    (let* ([L (vec-sub orig pos)]
       [a (dot-product dir dir)]
       [b (* 2 (dot-product L dir))]
       [c (- (dot-product L L) (* radius radius))]
       [res (filter (lambda [x] (> x 0)) (solve-quadratic a b c))])
      (cond ((null? res) -1)
        ((null? (cdr res)) (car res))
        ((< (car res) (cadr res)) (car res))
        (#t (cadr res)))))
  (define (get-normal pt) (normalized (vec-sub pt pos)))
  (define (get-color pt) col)
  (define (get-surface-point pt)
            (vec-add (vec-mul (get-normal pt)
                  (* 1.0000001 radius)) pos))

  (object 'sphere get-intersect get-normal get-color get-surface-point))

(define (new-plane pos normal col0 col1)
  (define (get-intersect orig dir)
    (let ([denom (dot-product dir normal)]
      [p (vec-sub pos orig)])
      (if (< denom 1e-6)
      -1
      (exact->inexact (/ (dot-product p normal) denom)))))
  (define (get-normal pt) (normalized (vec-mul normal -1)))
  (define (get-color pt)
    (if (xor (= (modulo (exact-floor (* 0.01 (vec-x pt))) 2) 0)
         (xor (= (modulo (exact-floor (* 0.01 (vec-y pt))) 2) 0)
          (= (modulo (exact-floor (* 0.01 (vec-z pt))) 2) 0)))
    col0
    col1))
  (define (get-surface-point pt)
    (vec-add (vec-mul (normalized normal) 0.0000001) pt))

  (object 'plane get-intersect get-normal get-color get-surface-point))

(define (cast-ray from to scene)
  (foldl (lambda [res acc]
       (cond ((null? acc) res)
         ((< (car res) (car acc)) res)
         (#t acc)))
     '()
     (filter (lambda [res] (> (car res) 0))
         (map (lambda [obj]
            (cons
             ((object-get-intersect obj) from to) obj)) scene))))

(define (blend-colors factor col1 col2)
  (let ([c1 (color-mul col1 factor)]
    [c2 (color-mul col2 (- 1 factor))])
    (color-add c1 c2)))

(define (reflect-color limit from obj scene)
  (let ([col ((object-get-color obj) from)])
    (if (and (> limit 0) (eq? (object-is obj) 'sphere))
    (blend-colors 0.75 ((ray-pixel (- limit 1) scene
                       ((object-get-surface-point obj) from))
                ((object-get-normal obj) from)) col)
    col)))

(define (find-color limit from obj scene)
  (let* ([col (reflect-color limit from obj scene)]
     [spoint ((object-get-surface-point obj) from)]
     [not-occluded (null? (cast-ray spoint light scene))]
     [spec (expt (/ (acos (dot-product
                   (vec-mul ((object-get-normal obj) from) -1)
                   (normalized light)))
            (* 2 pi (/ 1 3))) 5)]
     [final (color-mul col spec)])
    (if not-occluded
    final
    (blend-colors 0.25 final ambiant-color))))

(define (ray-pixel limit scene orig)
  (lambda [to]
    (let ([res (cast-ray orig to scene)])
      (if (null? res)
      bg-color
      (find-color limit (vec-mul to (car res)) (cdr res) scene)))))

(define (coords w h z)
  (for/list ([n (in-range (* w h))])
        (let ([x (- (modulo n w) (quotient w 2))]
          [y (- (quotient n w) (quotient h 2))])
          (vec x y z))))

(define (raytrace w h scene)
  (map (ray-pixel 25 scene (vec 0 0 0)) (coords w h (/ w 3))))

(define (color->list col)
  (let ([c (color-op col (lambda [x] (max 0 (min 255 (exact-round x)))))])
    (list 255 (color-r c) (color-g c) (color-b c))))

(define (to-bytes values)
  (apply bytes (flatten (map color->list values))))

(define (bitmap-raytracer w h scene)
  (let* ([bm (make-bitmap w h)]
     [dc (send bm make-dc)]
     [pixels (raytrace w h scene)])
    (send dc set-argb-pixels 0 0 w h
      (to-bytes pixels))
    bm))

(define (show-image bm)
  (let ([f (new frame% [label "Raytracer"])])
    (new message% [parent f] [label bm])
    (send f show #t)))

(define (save-image bm)
  (send bm save-file "raytracer.png" 'png))

(show-image (bitmap-raytracer 1920 1080 (make-scene)))