;;; Chickadee Game Toolkit
;;; Copyright © 2018, 2021, 2024 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 (chickadee data bytestruct)
  #:use-module (chickadee graphics)
  #:use-module (chickadee graphics buffer)
  #:use-module (chickadee graphics color)
  #:use-module (chickadee graphics pipeline)
  #:use-module (chickadee graphics shader)
  #:use-module (chickadee graphics texture)
  #:use-module (chickadee math)
  #:use-module (chickadee math matrix)
  #:use-module (chickadee math rect)
  #:use-module (chickadee math vector)
  #:use-module (ice-9 match)
  #: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)
  #: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-view
            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-bytestruct <particle-vertex>
  (struct (position <vec2>)
          (velocity <vec2>)
          (accel <vec2>)
          (life s32)))

(define-bytestruct <quad-vertex>
  (struct (position <vec2>)
          (uv <vec2>)))

(define-bytestruct <particle-uniforms>
  (struct (matrix <matrix4>)
          (lifetime s32)
          (animation-rows s32)
          (animation-columns s32)
          (start-color <color>)
          (end-color <color>)))

(define-record-type <particle-state>
  (make-particle-state uniforms sampler pipeline bindings matrix)
  particle-state?
  (uniforms particle-state-uniforms)
  (sampler particle-state-sampler)
  (pipeline particle-state-pipeline)
  (bindings particle-state-bindings)
  (matrix particle-state-matrix))

(define-graphics-variable particle-state
  (let ((shader (make-shader
                 (lambda (lang)
                   (values
                    "
#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 float life;
#elif defined(GLSL130)
in vec2 position;
in vec2 tex;
in vec2 offset;
in vec2 velocity;
in float life;
#elif defined(GLSL120)
attribute vec2 position;
attribute vec2 tex;
attribute vec2 offset;
attribute vec2 velocity;
attribute float life;
#endif
#ifdef GLSL120
varying vec2 fragTex;
varying float t;
#else
out vec2 fragTex;
out float t;
#endif

#ifdef GLSL120
uniform mat4 matrix;
uniform int lifetime;
uniform int animationRows;
uniform int animationColumns;
#else
layout (std140) uniform Particles
{
    mat4 matrix;
    int lifetime;
    int animationRows;
    int animationColumns;
    vec4 startColor;
    vec4 endColor;
};
#endif

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 = matrix * vec4(position + 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;
#else
#define fragColor gl_FragColor
#define texture texture2D
#endif

uniform sampler2D colorTexture;
#ifdef GLSL120
uniform vec4 startColor;
uniform vec4 endColor;
#else
layout (std140) uniform Particles
{
    mat4 matrix;
    int lifetime;
    int animationRows;
    int animationColumns;
    vec4 startColor;
    vec4 endColor;
};
#endif

void main (void) {
    fragColor = mix(endColor, startColor, t) * texture(colorTexture, fragTex);
}
")))))
    (make-particle-state
     (make-buffer (bytestruct-sizeof <particle-uniforms>)
                  #:name "Particle uniform buffer"
                  #:usage '(uniform))
     (make-sampler #:name "Particle sampler")
     (make-render-pipeline
      #:name "Particles"
      #:shader shader
      #:color-target (make-color-target #:blend-mode blend:alpha)
      #:vertex-layout
      (vector
       (make-vertex-buffer-layout
        #:stride (* 4 4)
        #:attributes (vector
                      ;; Position
                      (make-vertex-attribute
                       #:format 'float32x2)
                      ;; Texture
                      (make-vertex-attribute
                       #:format 'float32x2
                       #:offset (* 2 4))))
       (make-vertex-buffer-layout
        #:stride (* 7 4)
        #:step-mode 'instance
        #:attributes (vector
                      ;; Position
                      (make-vertex-attribute
                       #:format 'float32x2)
                      ;; Velocity
                      (make-vertex-attribute
                       #:format 'float32x2
                       #:offset (* 2 4))
                      ;; Life
                      (make-vertex-attribute
                       #:format 'sint32
                       #:offset (* 6 4)))))
      #:binding-layout (vector (make-texture-layout)
                               (make-sampler-layout)
                               (make-buffer-layout)))
     (make-vector 3 #f)
     (make-null-matrix4))))

(define-record-type <particles>
  (%make-particles capacity size bv vertex-buffers index-buffer
                   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)
  (vertex-buffers particles-vertex-buffers)
  (index-buffer particles-index-buffer)
  (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-view
                         (animation-rows 1)
                         (animation-columns 1)
                         (width (if texture-view
                                    (inexact->exact
                                     (floor
                                      (/ (texture-view-width texture-view)
                                         animation-columns)))
                                    8.0))
                         (height (if texture-view
                                     (inexact->exact
                                      (floor
                                       (/ (texture-view-height texture-view)
                                          animation-rows)))
                                     8.0))
                         (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-VIEW: The 2D 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* ((hw (/ width 2.0))
         (hh (/ height 2.0))
         (size (bytestruct-sizeof <quad-vertex>))
         (quad-vertex
          (bytevector->buffer
           ;; TODO: Add some syntax for declaratively generating a
           ;; bytevector of packed bytestructs.
           (let ((bv (make-bytevector (* size 4))))
             (bytestruct-pack! <quad-vertex>
                               (((position x) (- hw))
                                ((position y) (- hh))
                                ((uv x) 0.0)
                                ((uv y) 0.0))
                               bv 0)
             (bytestruct-pack! <quad-vertex>
                               (((position x) hw)
                                ((position y) (- hh))
                                ((uv x) 1.0)
                                ((uv y) 0.0))
                               bv size)
             (bytestruct-pack! <quad-vertex>
                               (((position x) hw)
                                ((position y) hh)
                                ((uv x) 1.0)
                                ((uv y) 1.0))
                               bv (* size 2))
             (bytestruct-pack! <quad-vertex>
                               (((position x) (- hw))
                                ((position y) hh)
                                ((uv x) 0.0)
                                ((uv y) 1.0))
                               bv (* size 3))
             bv)
           #:name "Particles quad vertex buffer"))
         (quad-index
          (bytevector->buffer (u32vector 0 2 3 0 1 2)
                              #:name "Particles quad index buffer"
                              #:usage '(index)))
         ;; 1 extra element as swap space for sorting.
         (k (* (+ capacity 1) (bytestruct-sizeof <particle-vertex>)))
         (particles (make-bytevector k))
         (particles-vertex (make-buffer k #:name "Particles vertex buffer")))
    (%make-particles capacity
                     0 ; initial size
                     particles
                     (vector quad-vertex particles-vertex)
                     quad-index
                     texture-view
                     animation-rows
                     animation-columns
                     speed-range
                     acceleration-range
                     direction-range
                     blend-mode
                     start-color
                     end-color
                     lifetime
                     sort
                     ;; No emitters initially.
                     '())))

(define (update-particles particles)
  "Advance the simulation of PARTICLES."
  (match particles
    (($ <particles> capacity (? exact-integer? size) bv #(_ vertex) _ _ _ _
                    speed-range acceleration-range direction-range
                    _ _ _ lifetime sort emitters)
     (let ((stride (bytestruct-sizeof <particle-vertex>)))
       ;; Update existing particles, removing dead ones.
       (let loop ((i 0) (size size))
         (if (< i size)
             (let ((offset (* i stride)))
               (call-with-values (lambda ()
                                   (bytestruct-unpack <particle-vertex>
                                                      ((position x)
                                                       (position y)
                                                       (velocity x)
                                                       (velocity y)
                                                       (accel x)
                                                       (accel y)
                                                       (life))
                                                      bv offset))
                 (lambda (x y dx dy ddx ddy life)
                   (let ((life (- life 1)))
                     (if (<= life 0)
                         (let ((size (- size 1)))
                           (bytevector-copy! bv (* size stride) bv offset stride)
                           (loop i size))
                         (begin
                           (bytestruct-pack! <particle-vertex>
                                             (((position x) (+ x dx))
                                              ((position y) (+ y dy))
                                              ((velocity x) (+ dx ddx))
                                              ((velocity y) (+ dy ddy))
                                              ((life) life))
                                             bv offset)
                           (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)))
                     (bytestruct-pack! <particle-vertex>
                                       (((position x) (+ rx (* (random:uniform) rw)))
                                        ((position y) (+ ry (* (random:uniform) rh)))
                                        ((velocity x) (* dir-x speed))
                                        ((velocity y) (* dir-y speed))
                                        ((accel x) (* dir-x accel))
                                        ((accel y) (* dir-y accel))
                                        ((life) lifetime))
                                       bv offset)
                     (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
                (match sort
                  ('young
                   (lambda (i j)
                     (< (bytestruct-unpack <particle-vertex> ((life)) bv i)
                        (bytestruct-unpack <particle-vertex> ((life)) bv j))))
                  ('old
                   (lambda (i j)
                     (> (bytestruct-unpack <particle-vertex> ((life)) bv i)
                        (bytestruct-unpack <particle-vertex> ((life)) bv j))))
                  (_
                   (error "invalid 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))))
       (let* ((k (* stride (particles-size particles)))
              (dst (map-buffer vertex 'write 0 k)))
         (bytevector-copy! bv 0 dst 0 k)
         (unmap-buffer vertex))))))

;; TODO: Blend mode
(define (draw-particles* particles matrix)
  "Render PARTICLES with MATRIX applied."
  (match particles
    (($ <particles> _ size _ vertex-buffers index-buffer texture-view
                    rows columns _ _ _ blend-mode start-color end-color
                    lifetime)
     (match (graphics-variable-ref particle-state)
       (($ <particle-state> uniforms sampler pipeline bindings mvp)
        (matrix4-mult! mvp matrix (current-projection))
        (let ((bv (map-buffer uniforms 'write 0
                              (bytestruct-sizeof <particle-uniforms>))))
          (bytestruct-pack! <particle-uniforms>
                            (((matrix) mvp)
                             ((lifetime) lifetime)
                             ((animation-rows) rows)
                             ((animation-columns) columns)
                             ((start-color) start-color)
                             ((end-color) end-color))
                            bv 0)
          (unmap-buffer uniforms))
        (vector-set! bindings 0 texture-view)
        (vector-set! bindings 1 sampler)
        (vector-set! bindings 2 uniforms)
        (draw 6
              #:instances size
              #:pipeline pipeline
              #:index-buffer index-buffer
              #:vertex-buffers vertex-buffers
              #:bindings bindings))))))

(define %default-matrix (make-identity-matrix4))
(define (draw-particles particles)
  "Render PARTICLES."
  (draw-particles* particles %default-matrix))