;;; Chickadee Game Toolkit
;;; Copyright © 2018, 2021 David Thompson <dthompson2@worcester.edu>
;;;
;;; Licensed under the Apache License, Version 2.0 (the "License");
;;; you may not use this file except in compliance with the License.
;;; You may obtain a copy of the License at
;;;
;;;    http://www.apache.org/licenses/LICENSE-2.0
;;;
;;; Unless required by applicable law or agreed to in writing, software
;;; distributed under the License is distributed on an "AS IS" BASIS,
;;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
;;; See the License for the specific language governing permissions and
;;; limitations under the License.

(define-module (chickadee graphics particles)
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-4)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-9 gnu)
  #:use-module (system foreign)
  #:use-module (chickadee math)
  #:use-module (chickadee math matrix)
  #:use-module (chickadee math rect)
  #:use-module (chickadee math vector)
  #:use-module (chickadee graphics blend)
  #:use-module (chickadee graphics buffer)
  #:use-module (chickadee graphics color)
  #:use-module (chickadee graphics engine)
  #:use-module (chickadee graphics shader)
  #:use-module (chickadee graphics texture)
  #:export (make-particle-emitter
            particle-emitter?
            particle-emitter-spawn-area
            particle-emitter-rate
            particle-emitter-life
            particle-emitter-done?
            make-particles
            particles?
            particles-capacity
            particles-size
            particles-texture
            particles-blend-mode
            particles-color
            particles-spawn-area
            add-particle-emitter
            remove-particle-emitter
            update-particles
            draw-particles*
            draw-particles))

(define-record-type <particle-emitter>
  (%make-particle-emitter spawn-area rate life)
  particle-emitter?
  (spawn-area particle-emitter-spawn-area)
  (rate particle-emitter-rate)
  (life particle-emitter-life set-particle-emitter-life!))

(define* (make-particle-emitter spawn-area rate #:optional duration)
  "Return a new particle emitter that spawns RATE particles per frame
within SPAWN-AREA (a rectangle or 2D vector) for DURATION frames.  If
DURATION is not specified, the emitter will spawn particles
indefinitely."
  (%make-particle-emitter spawn-area rate duration))

(define (update-particle-emitter emitter)
  "Advance the lifecycle of EMITTER."
  (let ((life (particle-emitter-life emitter)))
    (when life
      (set-particle-emitter-life! emitter (- life 1)))))

(define (particle-emitter-done? emitter)
  "Return #t if EMITTER has finished emitting particles."
  (let ((life (particle-emitter-life emitter)))
    (and life (<= life 0))))

(define-geometry-type <quad-vertex>
  quad-vertex-ref
  quad-vertex-set!
  quad-vertex-append!
  (position vec2)
  (texture vec2))

(define-geometry-type <particle-vertex>
  particle-vertex-ref
  particle-vertex-set!
  particle-vertex-append!
  (position vec2)
  (velocity vec2)
  (acceleration vec2)
  (life int))

(define-graphics-variable particles-shader
  (strings->shader
   "
#ifdef GLSL330
layout (location = 0) in vec2 position;
layout (location = 1) in vec2 tex;
layout (location = 2) in vec2 offset;
layout (location = 3) in vec2 velocity;
layout (location = 4) in vec2 acceleration;
layout (location = 5) in float life;
#elif defined(GLSL130)
in vec2 position;
in vec2 tex;
in vec2 offset;
in vec2 velocity;
in vec2 acceleration;
in float life;
#elif defined(GLSL120)
attribute vec2 position;
attribute vec2 tex;
attribute vec2 offset;
attribute vec2 velocity;
attribute vec2 acceleration;
attribute float life;
#endif
#ifdef GLSL120
varying vec2 fragTex;
varying float t;
#else
out vec2 fragTex;
out float t;
#endif
uniform mat4 mvp;
uniform int lifetime;
uniform int animationRows;
uniform int animationColumns;

void main(void) {
    t = life / lifetime;
    int numTiles = animationRows * animationColumns;
    int tile = int(numTiles * (1.0 - t));
    float tx = float(mod(tile, animationColumns)) / animationColumns;
    float ty = float(tile / animationColumns) / animationRows;
    float tw = 1.0 / animationColumns;
    float th = 1.0 / animationRows;
    fragTex = vec2(tx, ty) + tex * vec2(tw, th);
    gl_Position = mvp * vec4(position.xy + offset, 0.0, 1.0);
}
"
   "
#ifdef GLSL120
varying vec2 fragTex;
varying float t;
#else
in vec2 fragTex;
in float t;
#endif
#ifdef GLSL330
out vec4 fragColor;
#endif
uniform sampler2D color_texture;
uniform vec4 startColor;
uniform vec4 endColor;

void main (void) {
#ifdef GLSL330
    fragColor = mix(endColor, startColor, t) * texture(color_texture, fragTex);
#else
    gl_FragColor = mix(endColor, startColor, t) * texture2D(color_texture, fragTex);
#endif
}
"))

(define-graphics-variable mvp-matrix (make-null-matrix4))

(define-record-type <particles>
  (%make-particles capacity size bv geometry
                   texture animation-rows animation-columns
                   speed-range acceleration-range direction-range
                   blend-mode start-color end-color lifetime
                   sort emitters)
  particles?
  (capacity particles-capacity)
  (size particles-size set-particles-size!)
  (bv particles-bv)
  (geometry particles-geometry)
  (texture particles-texture set-particles-texture!)
  (animation-rows particles-animation-rows)
  (animation-columns particles-animation-columns)
  (speed-range particles-speed-range set-particles-speed-range!)
  (acceleration-range particles-acceleration-range
                      set-particles-acceleration-range!)
  (direction-range particles-direction-range set-particles-direction-range!)
  (blend-mode particles-blend-mode set-particles-blend-mode!)
  (start-color particles-start-color set-particles-start-color!)
  (end-color particles-end-color set-particles-end-color!)
  (lifetime particles-lifetime set-particles-lifetime!)
  (sort particles-sort set-particles-sort!)
  (emitters particles-emitters set-particles-emitters!))

(define (add-particle-emitter particles emitter)
  "Add EMITTER to PARTICLES."
  (set-particles-emitters! particles
                           (cons emitter (particles-emitters particles))))

(define (remove-particle-emitter particles emitter)
  "Remove EMITTER from PARTICLES."
  (set-particles-emitters! particles
                           (delete emitter (particles-emitters particles))))

(define* (make-particles capacity #:key
                         (blend-mode blend:alpha)
                         (start-color white)
                         (end-color (make-color 0.0 0.0 0.0 0.0))
                         (texture null-texture)
                         (animation-rows 1)
                         (animation-columns 1)
                         (width (if (texture-null? texture)
                                    8.0
                                    (inexact->exact
                                     (floor
                                      (/ (texture-width texture)
                                         animation-columns)))))
                         (height (if (texture-null? texture)
                                     8.0
                                     (inexact->exact
                                      (floor
                                       (/ (texture-height texture)
                                          animation-rows)))))
                         (speed-range (vec2 0.1 1.0))
                         (acceleration-range (vec2 0.0 0.1))
                         (direction-range (vec2 0.0 (* 2 pi)))
                         (lifetime 30)
                         sort)
  "Return a new particle system that may contain up to CAPACITY
particles.  Achieving the desired particle effect involves tweaking
the following keyword arguments as needed:

- BLEND-MODE: Pixel blending mode.  'alpha' by default.

- START-COLOR: The tint color of the particle at the beginning of its
life.  White by default.

- END-COLOR: The tint color of the particle at the end of of its life.
Completely transparent by default for a fade-out effect.  The color in
the middle of a particle's life will be an interpolation of
START-COLOR and END-COLOR.

- TEXTURE: The texture applied to the particles.  The texture may be
subdivided into many animation frames.

- ANIMATION-ROWS: How many animation frame rows there are in the
texture.  Default is 1.

- ANIMATION-COLUMNS: How many animation frame columns there are in the
texture.  Default is 1.

- WIDTH: The width of each particle.  By default, the width of an
animation frame (in pixels) is used.

- HEIGHT: The height of each particle.  By default, the height of an
animation frame (in pixels) is used.

- SPEED-RANGE: A 2D vector containing the min and max particle speed.
Each particle will have a speed chosen at random from this range.  By
default, speed ranges from 0.1 to 1.0.

- ACCELERATION-RANGE: A 2D vector containing the min and max particle
acceleration.  Each particle will have an acceleration chosen at
random from this range.  By default, acceleration ranges from 0.0 to
0.1.

- DIRECTION-RANGE: A 2D vector containing the min and max particle
direction as an angle in radians.  Each particle will have a direction
chosen at random from this range.  By default, the range covers all
possible angles.

- LIFETIME: How long each particle lives, measured in updates. 30 by
default.

- SORT: 'youngest' if youngest particle should be drawn last or
'oldest' for the reverse.  By default, no sorting is applied at all."
  (let ((geometry (make-geometry (list (list <quad-vertex> #:capacity 4)
                                       (list <particle-vertex> #:divisor 1))
                                 capacity
                                 #:index-capacity 6)))
    (with-geometry* geometry (<quad-vertex> 'index)
      (let ((hw (/ width 2.0))
            (hh (/ height 2.0)))
        (quad-vertex-append! geometry
                             ((- hw) (- hh) 0.0 0.0)
                             (hw (- hh) 1.0 0.0)
                             (hw hh 1.0 1.0)
                             ((- hw) hh 0.0 1.0)))
      (geometry-index-append! geometry 0 3 2 0 2 1))
    (%make-particles capacity
                     0
                     ;; 1 extra element as swap space for sorting.
                     (make-bytevector (* (+ capacity 1)
                                         (geometry-type-stride
                                          <particle-vertex>)))
                     geometry
                     texture
                     animation-rows
                     animation-columns
                     speed-range
                     acceleration-range
                     direction-range
                     blend-mode
                     start-color
                     end-color
                     lifetime
                     sort
                     '())))

(define (update-particles particles)
  "Advance the simulation of PARTICLES."
  (let* ((speed-range (particles-speed-range particles))
         (acceleration-range (particles-acceleration-range particles))
         (direction-range (particles-direction-range particles))
         (sort (particles-sort particles))
         (lifetime (particles-lifetime particles))
         (float-ref bytevector-ieee-single-native-ref)
         (float-set! bytevector-ieee-single-native-set!)
         (int-ref bytevector-s32-native-ref)
         (int-set! bytevector-s32-native-set!)
         (y-offset 4)
         (dx-offset 8)
         (dy-offset 12)
         (ddx-offset 16)
         (ddy-offset 20)
         (life-offset 24))
    (let* ((bv (particles-bv particles))
           (stride (geometry-type-stride <particle-vertex>))
           (capacity (particles-capacity particles)))
      ;; Update existing particles, removing dead ones.
      (let loop ((i 0)
                 (size (particles-size particles)))
        (if (< i size)
            (let* ((offset (* i stride))
                   (life (- (int-ref bv (+ offset life-offset)) 1)))
              (if (<= life 0)
                  (let ((new-size (- size 1)))
                    (bytevector-copy! bv (* new-size stride) bv offset stride)
                    (loop i new-size))
                  (let ((x (float-ref bv offset))
                        (y (float-ref bv (+ offset y-offset)))
                        (dx (float-ref bv (+ offset dx-offset)))
                        (dy (float-ref bv (+ offset dy-offset)))
                        (ddx (float-ref bv (+ offset ddx-offset)))
                        (ddy (float-ref bv (+ offset ddy-offset))))
                    (int-set! bv (+ offset life-offset) life)
                    (float-set! bv offset (+ x dx))
                    (float-set! bv (+ offset y-offset) (+ y dy))
                    (float-set! bv (+ offset dx-offset) (+ dx ddx))
                    (float-set! bv (+ offset dy-offset) (+ dy ddy))
                    (loop (+ i 1) size))))
            (set-particles-size! particles size)))
      ;; Add particles from each active emitter and then remove
      ;; emitters that have completed.
      (let ((sx (vec2-x speed-range))
            (sy (vec2-y speed-range))
            (ax (vec2-x acceleration-range))
            (ay (vec2-y acceleration-range))
            (dx (vec2-x direction-range))
            (dy (vec2-y direction-range))
            (emitters (particles-emitters particles)))
        (define (emit emitter any-done?)
          (let* ((spawn-area (particle-emitter-spawn-area emitter))
                 (rate (particle-emitter-rate emitter))
                 (rx (rect-x spawn-area))
                 (ry (rect-y spawn-area))
                 (rw (rect-width spawn-area))
                 (rh (rect-height spawn-area))
                 (start (particles-size particles))
                 (end (min (+ start rate) capacity)))
            (let loop ((i start))
              (if (< i end)
                  (let* ((offset (* i stride))
                         (speed (+ (* (random:uniform) (- sy sx)) sx))
                         (accel (+ (* (random:uniform) (- ay ax)) ax))
                         (dir (+ (* (random:uniform) (- dy dx)) dx))
                         (dir-x (cos dir))
                         (dir-y (sin dir)))
                    (float-set! bv offset (+ rx (* (random:uniform) rw)))
                    (float-set! bv (+ offset y-offset)
                                (+ ry (* (random:uniform) rh)))
                    (float-set! bv (+ offset dx-offset) (* dir-x speed))
                    (float-set! bv (+ offset dy-offset) (* dir-y speed))
                    (float-set! bv (+ offset ddx-offset) (* dir-x accel))
                    (float-set! bv (+ offset ddy-offset) (* dir-y accel))
                    (int-set! bv (+ offset life-offset) lifetime)
                    (loop (+ i 1)))
                  (begin
                    (set-particles-size! particles end)
                    (update-particle-emitter emitter)
                    (or any-done? (particle-emitter-done? emitter)))))))
        (when (fold emit #f emitters)
          (set-particles-emitters! particles
                                   (remove particle-emitter-done? emitters))))
      ;; Sort particles.
      (when sort
        (let ((compare (cond
                        ((eq? sort 'young)
                         (lambda (i j)
                           (< (int-ref bv (+ i life-offset))
                              (int-ref bv (+ j life-offset)))))
                        ((eq? sort 'old)
                         (lambda (i j)
                           (> (int-ref bv (+ i life-offset))
                              (int-ref bv (+ j life-offset)))))
                        (else
                         (error "unknown particle sorting method" sort))))
              (tmp (* (particles-capacity particles) stride)))
          (define (swap i j)
            (bytevector-copy! bv i bv tmp stride)
            (bytevector-copy! bv j bv i stride)
            (bytevector-copy! bv tmp bv j stride))
          ;; In the benchmarks I've done, insertion sort has
          ;; performed much better than quicksort here.  The number
          ;; of comparisons and swaps is much fewer.
          (define (sort start end)
            (let outer ((i (+ start stride)))
              (when (< i end)
                (let inner ((j i))
                  (when (and (> j start)
                             (compare j (- j stride)))
                    (swap (- j stride) j)
                    (inner (- j stride))))
                (outer (+ i stride)))))
          (sort 0 (* (particles-size particles) stride))))
      (with-geometry* (particles-geometry particles) (<particle-vertex>)
        (geometry-import! (particles-geometry particles) <particle-vertex> bv 0
                          (particles-size particles))))))

(define (draw-particles* particles matrix)
  "Render PARTICLES with MATRIX applied."
  (let ((shader (graphics-variable-ref particles-shader))
        (mvp (graphics-variable-ref mvp-matrix))
        (geometry (particles-geometry particles)))
    (with-graphics-state ((g:blend-mode (particles-blend-mode particles))
                          (g:texture-0 (particles-texture particles)))
      (shader-apply/instanced shader
                              (geometry-vertex-array geometry)
                              (particles-size particles)
                              #:mvp (if matrix
                                        (begin
                                          (matrix4-mult! mvp matrix
                                                         (current-projection))
                                          mvp)
                                        (current-projection))
                              #:start-color (particles-start-color particles)
                              #:end-color (particles-end-color particles)
                              #:lifetime (particles-lifetime particles)
                              #:animation-rows
                              (particles-animation-rows particles)
                              #:animation-columns
                              (particles-animation-columns particles)))))

(define (draw-particles particles)
  "Render PARTICLES."
  (draw-particles* particles #f))