(define (calc-xyz index max-index r)
    (let*
        ( [angle (* 6.28312 (/ index (- max-index 1)))]
            [x (* (cos angle) r)]
            [y (* (sin angle) r)] )
        (vector x y 0)))

(define (build-ellipse rmin rmaj num-verts)
    (define p (build-polygons (* 3 num-verts) 'triangle-list))
    (with-primitive p
        (for ([i (in-range 0 (* 3 num-verts) 3) ])
            (pdata-set! "p" i (vector 0 0 0))
            (pdata-set! "n" i (vector 0 0 1))
            (pdata-set! "p" (+ i 1) (calc-xyz (/ i 3) num-verts rmin))
            (pdata-set! "n" (+ i 1) (vnormalise (calc-xyz (/ i 3) num-verts rmin)))
            (pdata-set! "p" (+ i 2) (calc-xyz (+ (/ i 3) 1) num-verts rmin))
            (pdata-set! "n" (+ i 2) (vnormalise (calc-xyz (+ (/ i 3) 1) num-verts rmin))))
        (poly-convert-to-indexed))
    p)

(define-struct stones ((pos-list #:mutable) size-list (root #:mutable) (obj-list #:mutable)))

(define (choose l)
    (list-ref l (random (length l))))

(define (stones-init num area size)
    (make-stones
        (build-list num
            (lambda (_)
                (vmul (vector (vx (srndvec)) (vy (srndvec)) 0) area)))
        (build-list num
            (lambda (_)
                (* size (rndf))))
        0
        '()))

(define (stones-build stones)
    (let* ((root (build-locator))
            (objs (with-state
                    (parent root)
                    (map
                        (lambda (pos size)
                            (if (and (< size 0.4) (zero? (random 3)))
                                (let ((o (with-state
                                                (hint-unlit)
                                                (scale 0.2)
                                                (colour (vector 0.25 0.5 0))
                                                (build-plane))))
                                    (with-primitive o (apply-transform))
                                    o)
                                (with-state
                                    ;(hint-unlit)
                                    (hint-ignore-depth)
                                    (colour (hsv->rgb (vector (+ 0 (* 0.2 (rndf))) 0.5 (+ 0.1 (rndf)))))
                                    (translate pos)
                                    (build-ellipse size size 32))))
                        (stones-pos-list stones)
                        (stones-size-list stones)))))
        (set-stones-obj-list! stones objs)
        (set-stones-root! stones root)
        stones))

(define (stones-relax stones amount)
    (set-stones-pos-list! stones
        (map 
            (lambda (pos size)
                (foldl
                    (lambda (opos osize r)
                        (if (< (vdist pos opos) (+ size osize))
                            (vadd r (vmul (vnormalise (vsub pos opos)) amount))
                            r))
                    pos
                    (stones-pos-list stones)
                    (stones-size-list stones)))
            (stones-pos-list stones)
            (stones-size-list stones))))

(define (stones-update stones)
    (let ((root (build-locator)))
        (with-state
            (parent root)
            (for-each
                (lambda (obj pos size)
                    (with-primitive obj
                        (identity)
                        (translate pos)))
                (stones-obj-list stones)
                (stones-pos-list stones)
                (stones-size-list stones)))
        (set-stones-root! stones root)))

(define (nudge stones amount)
    (pdata-map!
        (lambda (p)
            (foldl
                (lambda (pos size r)
                    (if (< (vdist p pos) size)
                        (vadd r (vmul (vnormalise (vsub p pos)) amount))
                        r))
                p
                (stones-pos-list stones)
                (stones-size-list stones)))
        "p"))

(define (shrink amount)
    (pdata-index-map!
        (lambda (i p)
            (if (or (zero? i) (eq? i (- (pdata-size) 1)))
                p
                (vadd (vmul p (- 1 amount)) 
                    (vmul (pdata-ref "p" (+ i 1)) (* 0.5 amount))
                    (vmul (pdata-ref "p" (- i 1)) (* 0.5 amount)))))
        "p"))

(define (mangle-normals amount)
    (pdata-map! 
        (lambda (n)
            (vadd n (vmul (srndvec) amount)))
        "n"))

(define (build-root x)
    (let ((root (with-state 
;                    (hint-unlit)
                    (hint-ignore-depth)
                    (colour (vmul (vector 0.2 0.4 0.2) (+ 0.5 (rndf))))
                    (build-ribbon 25))))
        
        (with-primitive root
            (pdata-index-map!
                (lambda (i p)
                    (vector (+ (crndf) x) (- (* i 0.3) 2) 0)
                    #;(vmul (vector (vx (srndvec)) (vy (srndvec)) 0) 3))
                "p")
            (pdata-index-map!
                (lambda (i w)
                    (* 0.4 (/ i (pdata-size))))
                "w"))
        root))

(clear)
(clear-colour 0)
(define l (make-light 'point 'free))
(light-position l (vector 0 10 -5))
(light-diffuse 0 (vector 0.2 0.2 0.2))
(light-diffuse l (vector 1 1 1))

(define s (stones-build (stones-init 100 5 1)))

(define roots (build-list 10 (lambda (_) (build-root (* 5 (crndf))))))

(define (animate)
    (for-each
        (lambda (root)
            (with-primitive root
                (nudge s 0.01)
                (shrink 0.01)))
        roots)
    (stones-relax s 0.01)
    (stones-update s))

(every-frame (animate))