;#lang scheme/base
;(require fluxus-016/drflux)

(define max-tokens 2000)
(define additive-growth #t)

(define (texpath s)
    (string-append "/home/dave/flotsam/groworld/game-prototypes/danceplant/" s))

(define-struct token (char gen age id (root #:mutable)))

(define nxt-id 0)

(define (token-gid) 
    (set! nxt-id (+ nxt-id 1))
    nxt-id)

(define (reset-token-gid)
    (set! nxt-id 0))

(define (token-eq? a b)
    (eq? (token-char a) (token-char b)))

(define (token-copy a gen)
    (make-token (token-char a) gen (token-age a) (token-gid) 0))

(define (make-token-list str)
    (map
        (lambda (char)
            (make-token char 0 0 0 0))
        (string->list str)))

(define (copy-token-list str gen)
    (map 
        (lambda (token)
            (token-copy token gen))
        str))

(define (print-token-list tl)
    (for-each
        (lambda (token)
            (printf "~a" (token-char token)))
        tl)
    (newline))

;~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

(define-struct ls-rule (from to))

; make a rule from a char and a string
(define (make-rule from to)
    (when (not (char? from)) (error ""))
    (make-ls-rule
        (make-token from 0 0 0 0)
        (make-token-list to)))

; convert a list of strings into a list of rules
(define (make-rule-list l)
    (map
        (lambda (r)
            (make-rule (string-ref (car r) 0) (cadr r)))
        l))

;~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

(define-struct plant ((rules #:mutable) 
        axiom
        angle
        branch-scale
        (str #:mutable) 
        (gen #:mutable) 
        (root #:mutable)        
        base-root
        (components #:mutable)
        (anim-time #:mutable)
        (col #:mutable)
        (decay #:mutable)))

(define (build-plant axiom angle branch-scale rules)
    (let ((base-root (build-locator)))
        (with-state
            (parent base-root)
            (make-plant (make-rule-list rules) (make-token-list axiom)
                angle branch-scale (make-token-list axiom) 0 (build-locator) 
                base-root '() (time) (vector 1 1 1) #f))))

(define (plant-clear plant)
    (set-plant-gen! plant 0)
    (set-plant-str! plant (copy-token-list (plant-axiom plant) 0))
    (set-plant-components! plant '()))

(define (plant-inc-gen plant)
    (when (not (plant-decay plant))
        (set-plant-anim-time! plant (time))
        (set-plant-gen! plant (+ (plant-gen plant) 1))
        (destroy (plant-root plant))
        (with-state
            (parent (plant-base-root plant))
            (set-plant-root! plant (build-locator)))
        (reset-token-gid)
        (plant-generate plant)        
        ; (plant-print plant)        
        (set-plant-components! plant (build-components plant (plant-root plant) 
                (plant-angle plant) (plant-branch-scale plant)))))

(define (plant-update plant time)            
    (components-update plant (plant-components plant) time))

(define (plant-print plant)
    (print-token-list (plant-str plant)))

(define (plant-add-to-rule plant n s)
    
    (cond ((plant-decay plant)
            (component-do-decay plant (plant-components plant)))
        (else
            
            (let ((c 0))
                (set-plant-rules! plant 
                    (map 
                        (lambda (rule)
                            (set! c (+ c 1))
                            (if (eq? (- c 1) n)
                                (make-ls-rule (ls-rule-from rule)
                                    (append (ls-rule-to rule) 
                                        (map
                                            (lambda (char)
                                                (make-token char -1 (plant-gen plant) 0 0))
                                            (string->list s))))
                                rule))
                        (plant-rules plant)))))))

; try all the rules on this token - returns #f if none are matched
(define (plant-run-rules plant token gen)
    (foldl
        (lambda (rule str)
            (if str ; if str is not #f
                str ; then we have already found a rule, return it
                (if (and (token-eq? token (ls-rule-from rule)) ; check this rule
                        (< (length (plant-str plant)) max-tokens))
                    (copy-token-list (ls-rule-to rule) gen) ; return the result
                    #f))) ; no match
        #f
        (plant-rules plant)))

; runs the lsystem rules on every token in a list, returns the new list
(define (plant-search-replace plant pos result gen)
    (cond 
        ((>= pos (length (plant-str plant))) result)
        (else
            (let ((ret (plant-run-rules plant (list-ref (plant-str plant) pos) gen)))
                (if (list? ret)
                    (plant-search-replace plant (+ pos 1) (append result ret) gen)
                    (plant-search-replace plant (+ pos 1) 
                        (append result (list (list-ref (plant-str plant) pos))) gen))))))

(define (plant-generate plant)
    (when (< (length (plant-str plant)) max-tokens)
        ;    (print-token-list (plant-str plant))
        (cond (additive-growth
                (set-plant-str! plant (plant-search-replace plant 0 
                        '() (plant-gen plant))))
            (else
                ; reset the string
                (set-plant-str! plant (copy-token-list (plant-axiom plant) 0))
                (for ((i (in-range 0 (plant-gen plant))))
                    (set-plant-str! plant (plant-search-replace plant 0 '() (+ i 1))))))))

;~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

(define-struct component (type root gen rot sc id (decay #:mutable) fall))

(define (components-update plant components time)
    (for-each
        (lambda (component)
            (cond ((component-decay component)
                    (with-primitive (component-root component)
                        (translate (component-fall component))))
                (else
                    (let ((growing (eq? (plant-gen plant) (component-gen component)))
                            (t (- time (plant-anim-time plant))))
                        (with-primitive (component-root component)
                            (identity)
                            (translate (vector 1 0 0))        
                            
                            (rotate (vector 0 (+ (* 0.5 (sin (+ (component-id component) 
                                                    (* 1 time))))) 0))
                            
                            (if (and growing (< t 1))
                                (rotate (vector 0 (* t (component-rot component)) 0))                        
                                (rotate (vector 0 (component-rot component) 0)))
                            
                            (scale (component-sc component))
                            (when (and growing (< t 1)) (scale (vector t 1 1))))))))
        components))


(define (component-do-decay plant components)
    (for-each
        (lambda (component)
            (when (zero? (random 100))
                ;(null? (with-primitive (component-root component) (get-children)))
                (set-component-decay! component #t)
                (with-primitive (component-root component)
                    (detach-parent)
                    ;                    (apply-transform)
                    (parent (plant-root plant)))))
        components))

; builds objects from a list of tokens
(define (build-components plant root angle branch-scale)
    (let ((parent-stack '(root))
            (last-obj root)
            (rot 0)
            (sc 1)
            (sc-stack '(1))
            (rot-stack '(0))
            (bo (with-state                    
                    (rotate (vector 90 0 0))
                    (translate (vector 0.5 0 0))
                    (scale (vector 1.1 0.5 0.5))
                    (build-plane)))
            (lo (with-state                    
                    (rotate (vector 90 0 0))
                    (translate (vector 0.5 0 0))
                    (build-plane))))
        
        
        (define (build-obj-component token)
            (let ((obj (with-state
                            (colour (plant-col plant))
                            (hint-ignore-depth)
                            (cond 
                                ((char=? #\F (token-char token))
                                    (texture (load-texture (texpath "textures/branch-a.png"))))
                                ((char=? #\G (token-char token))
                                    (texture (load-texture (texpath "textures/branch-b.png"))))
                                ((char=? #\L (token-char token))
                                    (texture (load-texture (texpath "textures/leaf-a.png"))))
                                ((char=? #\M (token-char token))
                                    (texture (load-texture (texpath "textures/leaf-b.png")))))
                            (hint-unlit)
                            (cond ((or (char=? #\F (token-char token))
                                        (char=? #\G (token-char token)))
                                    (build-copy bo))
                                ((or (char=? #\L (token-char token))
                                        (char=? #\M (token-char token)))
                                    (build-copy lo))))))
                (with-primitive obj
                    (hide 0)
                    (when (not (zero? last-obj))
                        (parent last-obj)))
                (set! last-obj obj)
                (make-component (token-char token) obj
                    ;                    (if (> (token-age token) (token-gen token))
                        ;                        (+ (token-age token) 1)
                        (token-gen token);)
                    rot sc (token-id token) #f (srndvec))))
        
        (with-primitive bo
            (hide 1)
            (pdata-set! "t" 0 (vector 0 0 0))
            (pdata-set! "t" 1 (vector 0 1 0))
            (pdata-set! "t" 2 (vector 1 1 0))
            (pdata-set! "t" 3 (vector 1 0 0))
            (apply-transform))
        
        (with-primitive lo
            (hide 1)
            (pdata-set! "t" 0 (vector 1 1 0))
            (pdata-set! "t" 1 (vector 1 0 0))
            (pdata-set! "t" 2 (vector 0 0 0))
            (pdata-set! "t" 3 (vector 0 1 0))
            (apply-transform))
        
        (foldl
            (lambda (token ret)
                (let ((char (token-char token)))
                    (cond 
                        ((or (char=? #\F char) (char=? #\G char) (char=? #\L char) (char=? #\M char))
                            (let ((r (build-obj-component token)))
                                (set! rot 0)
                                (cons r ret)))
                        ((char=? #\+ char) (set! rot (+ rot angle)) ret)
                        ((char=? #\- char) (set! rot (- rot angle)) ret)
                        ((char=? #\[ char)
                                (set! parent-stack (cons last-obj parent-stack))
                                (set! sc-stack (cons sc sc-stack))
                                (set! sc (* sc branch-scale))
                                (set! rot-stack (cons rot rot-stack))
                                ret)
                            ((char=? #\] char)
                            (set! last-obj (car parent-stack))                        
                            (set! parent-stack (cdr parent-stack))
                            (set! sc (car sc-stack))                        
                            (set! sc-stack (cdr sc-stack))
                            (set! rot (car rot-stack))                        
                            (set! rot-stack (cdr rot-stack))
                            ret)
                        (else ret))))
            '()
            (plant-str plant))))

;~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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

(define (build-foliage count rules)
    (build-list count 
        (lambda (_) 
            (with-state
                (translate (vector 0 (+ 0 (* (rndf) 10)) (* 15 (crndf))))
                (build-plant "FA" (* 40 (crndf)) 0.96 (choose rules))))))


(define (foliage-grow foliage)
    (for-each
        (lambda (plant)
            (for ((i (in-range 4 (random 10))))
                (set-plant-col! plant (vmul (vector 0 0.6 0.6) (rndf)))
                (plant-inc-gen plant)))
        foliage))

(define (foliage-update foliage time)
    (for-each
        (lambda (plant)
            (plant-update plant time))
        foliage))

;~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

(define debounce-vec (build-vector 10 (lambda (_) #f)))

(define (debounce n)
    (vector-ref debounce-vec n))

(define (debounce! n s)
    (vector-set! debounce-vec n s))

(define (on-key p n proc)
    (if p
        (when (debounce n)
            (proc)
            (debounce! n #f))
        (debounce! n #t)))

;~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

(clear)
(clear-colour (vector 0 0.5 0.4))
(set-camera-transform (mmul (mtranslate (vector 0 -5 -10))
        (mrotate (vector 90 90 180))))

(with-state
    (hint-unlit)
    (colour (vector 0 0.4 0.3))
    (scale 100)
    (rotate (vector 0 90 0))
    (build-plane))

(define foliage (build-foliage 30
        '((("A" "F[-[+++FM]G[-FL]+A]"))
            (("A" "F-F[++AL]G-A"))
            (("A" "F++[+A+FFL]")))))

(foliage-grow foliage)

(define rules '(("a" "F") ("b" "F") ("A" "[a][b]")))
(define p (build-plant "FA" 40 0.96 rules))
(set-plant-col! p (vector 0.5 1 0.5))
(define decay-count 0)

(define (animate)
    
    (when (and (zero? decay-count)
            (or (> (length (plant-str p)) max-tokens) (> (plant-gen p) 20)))
        (set! decay-count 1)
        (set-plant-decay! p #t))
    
    (when (> decay-count 0)
        (set! decay-count (+ decay-count 1))
        ;        (display decay-count)(newline)
        )
    
    (when (> decay-count 100)
        (set! decay-count 0) 
        (set-plant-decay! p #f)
        (plant-clear p)
        (set-plant-rules! p (make-rule-list rules))
        (plant-inc-gen p))
    
    (on-key (key-special-pressed 101) 0 ; up
        (lambda ()
            (plant-add-to-rule p 0 "G+a")
            (plant-add-to-rule p 1 "G-b")
            (plant-inc-gen p)))
    
    (on-key (key-special-pressed 100) 1 ; left
        (lambda ()
            (plant-add-to-rule p 0 "+F+b")
            (plant-add-to-rule p 1 "+F+a")
            (plant-inc-gen p)))
    
    (on-key (key-special-pressed 102) 2 ; right
        (lambda ()
            (plant-add-to-rule p 0 "-F-a")
            (plant-add-to-rule p 1 "-F-b")
            (plant-inc-gen p)))
    
    (on-key (key-special-pressed 103) 3 ; down
        (lambda ()
            (plant-add-to-rule p 0 "[F++b-FL]")
            (plant-add-to-rule p 1 "[F--a+FL]")
            (plant-inc-gen p)))
    
    (on-key (key-pressed "q") 4
        (lambda ()
            (plant-inc-gen p)
            (plant-add-to-rule p 0 "b-a")))
    
    (on-key (key-pressed "e") 4
        (lambda ()
            (plant-inc-gen p)
            (plant-add-to-rule p 1 "a-b")))
    
    (on-key (key-pressed "z") 5
        (lambda ()
            (plant-inc-gen p)
            (plant-add-to-rule p 1 "b[+M]")))
    
    (on-key (key-pressed "c") 6
        (lambda ()
            (plant-add-to-rule p 0 "a+[--M]")
            (plant-inc-gen p)))
    
    (on-key (key-pressed "s") 7
        (lambda ()
            (plant-add-to-rule p 0 "A")
            (plant-add-to-rule p 1 "A")
            (plant-inc-gen p)))           
    
    (on-key (key-pressed (string #\newline)) 8
        (lambda ()
            (plant-clear p)
            (set-plant-rules! p (make-rule-list rules))))
    
    
    (plant-update p (time))
    (foliage-update foliage (time)))

(every-frame (animate))