(ns org.aurellem.worm-learn

   "General worm creation framework."

   {:author "Robert McIntyre"}

   (:use (cortex world util import body sense

                 hearing touch vision proprioception movement

                 test))

   (:import (com.jme3.math ColorRGBA Vector3f))

   (:import java.io.File)

   (:import com.jme3.audio.AudioNode)

   (:import com.aurellem.capture.RatchetTimer)

   (:import (com.aurellem.capture Capture IsoTimer))

   (:import (com.jme3.math Vector3f ColorRGBA)))

 

 (import org.apache.commons.math3.transform.TransformType)

 (import org.apache.commons.math3.transform.FastFourierTransformer)

 (import org.apache.commons.math3.transform.DftNormalization)

 

 (use 'clojure.pprint)

 (use 'clojure.set)

 (dorun (cortex.import/mega-import-jme3))

 (rlm.rlm-commands/help)

 

 (load-bullet)

 

 (def hand "Models/test-creature/hand.blend")

 

 (defn worm-model []

   (load-blender-model "Models/worm/worm.blend"))

 

 (def output-base (File. "/home/r/proj/cortex/render/worm-learn/curl"))

 

 

 (defn motor-control-program

   "Create a function which will execute the motor script"

   [muscle-labels

    script]

   (let [current-frame (atom -1)

         keyed-script (group-by first script)

         current-forces (atom {}) ]

     (fn [effectors]

       (let [indexed-effectors (vec effectors)]

         (dorun 

          (for [[_ part force] (keyed-script (swap! current-frame inc))]

            (swap! current-forces (fn [m] (assoc m part force)))))

         (doall (map (fn [effector power]

                       (effector (int power)))

                     effectors

                     (map #(@current-forces % 0) muscle-labels)))))))

 

 (defn worm-direct-control

   "Create keybindings and a muscle control program that will enable

    the user to control the worm via the keyboard."

   [muscle-labels activation-strength]

   (let [strengths (mapv (fn [_] (atom 0)) muscle-labels)

         activator

         (fn [n]

           (fn [world pressed?]

             (let [strength (if pressed? activation-strength 0)]

               (swap! (nth strengths n) (constantly strength)))))

         activators

         (map activator (range (count muscle-labels)))

         worm-keys

         ["key-f" "key-r"

          "key-g" "key-t"

          "key-h" "key-y"

          "key-j" "key-u"

          "key-k" "key-i"

          "key-l" "key-o"]]

     {:motor-control

      (fn [effectors]

       (doall

        (map (fn [strength effector]

               (effector (deref strength)))

             strengths effectors)))

      :keybindings

      ;; assume muscles are listed in pairs and map them to keys.

      (zipmap worm-keys activators)}))

 

 ;; These are scripts that direct the worm to move in two radically

 ;; different patterns -- a sinusoidal wiggling motion, and a curling

 ;; motions that causes the worm to form a circle.

 

 (def curl-script

   [[150 :d-flex 40]

    [250 :d-flex 0]])

 

 (def period 18)

 

 (def worm-muscle-labels

   [:base-ex :base-flex

    :a-ex :a-flex

    :b-ex :b-flex

    :c-ex :c-flex

    :d-ex :d-flex])

 

 (defn gen-wiggle [[flexor extensor :as muscle-pair] time-base]

   (let [period period

         power 45]

     [[time-base flexor power]

      [(+ time-base period) flexor 0]

      [(+ time-base period 1) extensor power]

      [(+ time-base (+ (* 2 period) 2))  extensor 0]]))

   

 (def wiggle-script

   (mapcat gen-wiggle (repeat 4000 [:a-ex :a-flex])

                      (range 100 1000000 (+ 3 (* period 2)))))

 

 

 (defn shift-script [shift script]

   (map (fn [[time label power]] [(+ time shift) label power])

        script))

 

 (def do-all-the-things 

   (concat

    curl-script

    [[300 :d-ex 40]

     [320 :d-ex 0]]

    (shift-script 280 (take 16 wiggle-script))))

 

 ;; Normally, we'd use unsupervised/supervised machine learning to pick

 ;; out the defining features of the different actions available to the

 ;; worm. For this project, I am going to explicitely define functions

 ;; that recognize curling and wiggling respectively. These functions

 ;; are defined using all the information available from an embodied

 ;; simulation of the action. Note how much easier they are to define

 ;; than if I only had vision to work with. Things like scale/position

 ;; invariance are complete non-issues here. This is the advantage of

 ;; body-centered action recognition and what I hope to show with this

 ;; thesis.

 

 

 ;; curled? relies on proprioception, resting? relies on touch,

 ;; wiggling? relies on a fourier analysis of muscle contraction, and

 ;; grand-circle? relies on touch and reuses curled? as a gaurd.

 

 (defn curled?

   "Is the worm curled up?"

   [experiences]

   (every?

    (fn [[_ _ bend]]

      (> (Math/sin bend) 0.64))

    (:proprioception (peek experiences))))

 

 (defn rect-region [[x0 y0] [x1 y1]]

   (vec

    (for [x (range x0 (inc x1))

          y (range y0 (inc y1))]

      [x y])))

 

 (def worm-segment-bottom (rect-region [8 15] [14 22]))

 

 (defn contact

   "Determine how much contact a particular worm segment has with

    other objects. Returns a value between 0 and 1, where 1 is full

    contact and 0 is no contact."

   [touch-region [coords contact :as touch]]

   (-> (zipmap coords contact)

       (select-keys touch-region)

       (vals)

       (#(map first %))

       (average)

       (* 10)

       (- 1)

       (Math/abs)))

 

 (defn resting?

   "Is the worm resting on the ground?"

   [experiences]

   (every?

    (fn [touch-data]

      (< 0.9 (contact worm-segment-bottom touch-data)))

    (:touch (peek experiences))))

 

 (defn vector:last-n [v n]

   (let [c (count v)]

     (if (< c n) v

         (subvec v (- c n) c))))

 

 (defn fft [nums]

   (map

    #(.getReal %)

    (.transform

     (FastFourierTransformer. DftNormalization/STANDARD)

     (double-array nums) TransformType/FORWARD)))

 

 (def indexed (partial map-indexed vector))

 

 (defn max-indexed [s]

   (first (sort-by (comp - second) (indexed s))))

 

 (defn wiggling?

   "Is the worm wiggling?"

   [experiences]

   (let [analysis-interval 0x40]

     (when (> (count experiences) analysis-interval)

       (let [a-flex 3

             a-ex   2

             muscle-activity

             (map :muscle (vector:last-n experiences analysis-interval))

             base-activity

             (map #(- (% a-flex) (% a-ex)) muscle-activity)]

         (= 2

            (first

             (max-indexed

              (map #(Math/abs %)

                   (take 20 (fft base-activity))))))))))

 

 (def worm-segment-bottom-tip (rect-region [15 15] [22 22]))

 

 (def worm-segment-top-tip (rect-region [0 15] [7 22]))

 

 (defn grand-circle?

   "Does the worm form a majestic circle (one end touching the other)?"

   [experiences]

   (and (curled? experiences)

        (let [worm-touch (:touch (peek experiences))

              tail-touch (worm-touch 0)

              head-touch (worm-touch 4)]

          (and (< 0.55 (contact worm-segment-bottom-tip tail-touch))

               (< 0.55 (contact worm-segment-top-tip    head-touch))))))

 

 

 (declare phi-space phi-scan)

 

 (defn debug-experience

   [experiences]

   (cond

    (grand-circle? experiences) (println "Grand Circle")

    (curled? experiences)       (println "Curled")

    (wiggling? experiences)     (println "Wiggling")

    (resting? experiences)      (println "Resting")))

 

 

 (def standard-world-view

   [(Vector3f. 4.207176, -3.7366982, 3.0816958)

    (Quaternion. 0.11118768, 0.87678415, 0.24434438, -0.3989771)])

 

 (def worm-side-view

   [(Vector3f. 4.207176, -3.7366982, 3.0816958)

    (Quaternion. -0.11555642, 0.88188726, -0.2854942, -0.3569518)])

 

 (def degenerate-worm-view

   [(Vector3f. -0.0708936, -8.570261, 2.6487997)

    (Quaternion. -2.318909E-4, 0.9985348, 0.053941682, 0.004291452)])

 

 (defn worm-world-defaults []

   (let [direct-control (worm-direct-control worm-muscle-labels 40)]

     (merge direct-control     

            {:view worm-side-view

             :record nil

             :experiences (atom [])

             :experience-watch debug-experience

             :worm-model worm-model

             :end-frame nil})))

 

 (defn dir! [file]

   (if-not (.exists file)

     (.mkdir file))

   file)

 

 (defn record-experience! [experiences data]

   (swap! experiences #(conj % data)))

 

 (defn enable-shadows [world]

   (let [bsr (doto

                 (BasicShadowRenderer. (asset-manager) 512)

               (.setDirection (.normalizeLocal (Vector3f. 1 -1 -1))))]

     (.addProcessor (.getViewPort world) bsr)))

 

 (defn enable-good-shadows [world]

   (let [pssm

         (doto (PssmShadowRenderer. (asset-manager) 1024 3)

           (.setDirection  (.normalizeLocal (Vector3f. -1 -3 -1)))

           (.setLambda (float 0.55))

           (.setShadowIntensity (float 0.6))

           (.setCompareMode PssmShadowRenderer$CompareMode/Software)

           (.setFilterMode  PssmShadowRenderer$FilterMode/Bilinear))]

     (.addProcessor (.getViewPort world) pssm)))

         

 

 (defn worm-world

   [& {:keys [record motor-control keybindings view experiences

              worm-model end-frame experience-watch] :as settings}]

   (let [{:keys [record motor-control keybindings view experiences

                 worm-model end-frame experience-watch]}

         (merge (worm-world-defaults) settings)

         worm (doto (worm-model) (body!))

         touch   (touch! worm)

         prop    (proprioception! worm)

         muscles (movement! worm)

         

         touch-display  (view-touch)

         prop-display   (view-proprioception)

         muscle-display (view-movement)

         

         floor

         (box 5 1 5 :position (Vector3f. 0 -10 0)

              :mass 0

              :texture "Textures/aurellem.png"

              :material "Common/MatDefs/Misc/Unshaded.j3md")

         timer (IsoTimer. 60)]

 

     (world

      (nodify [worm floor])

        (merge standard-debug-controls keybindings)

        (fn [world]

          (enable-good-shadows world)

          (.setShadowMode worm RenderQueue$ShadowMode/CastAndReceive)

          (.setShadowMode floor RenderQueue$ShadowMode/Receive)

                   

          (.setBackgroundColor (.getViewPort world) (ColorRGBA/White))

          (.setDisplayStatView world false)

          (.setDisplayFps world false)

          (position-camera world view)

          (.setTimer world timer)

          (display-dilated-time world timer)

          (when record

            (Capture/captureVideo

             world

             (dir! (File. record "main-view"))))

          (speed-up world)

          ;;(light-up-everything world)

          )

        (fn [world tpf]

          (if (and end-frame (> (.getTime timer) end-frame))

            (.stop world))

          (let [muscle-data (vec (motor-control muscles))

                proprioception-data (prop)

                touch-data (mapv #(% (.getRootNode world)) touch)]

            (when experiences

              (record-experience!

               experiences {:touch touch-data

                            :proprioception proprioception-data

                            :muscle muscle-data}))

            (when experience-watch

              (experience-watch @experiences))

            (muscle-display

             muscle-data

             (when record (dir! (File. record "muscle"))))

            (prop-display

             proprioception-data

             (when record (dir! (File. record "proprio"))))

            (touch-display 

             touch-data

             (when record (dir! (File. record "touch")))))))))

 

 

 

 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

 ;;;;;;;;   Phi-Space   ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

 

 (defn generate-phi-space []

   (let [experiences (atom [])]

     (run-world

      (apply-map 

       worm-world

       (merge

        (worm-world-defaults)

        {:end-frame 700

         :motor-control

         (motor-control-program worm-muscle-labels do-all-the-things)

         :experiences experiences})))

     @experiences))

 

 (defn bin [digits]

   (fn [angles]

     (->> angles

          (flatten)

          (map (juxt #(Math/sin %) #(Math/cos %)))

          (flatten)

          (mapv #(Math/round (* % (Math/pow 10 (dec digits))))))))

 

 ;; k-nearest neighbors with spatial binning. Only returns a result if

 ;; the propriceptive data is within 10% of a previously recorded

 ;; result in all dimensions.

 (defn gen-phi-scan [phi-space]

   (let [bin-keys (map bin [3 2 1])

         bin-maps

         (map (fn [bin-key]

                (group-by

                 (comp bin-key :proprioception phi-space)

                 (range (count phi-space)))) bin-keys)

         lookups (map (fn [bin-key bin-map]

                       (fn [proprio] (bin-map (bin-key proprio))))

                     bin-keys bin-maps)]

     (fn lookup [proprio-data]

       (set (some #(% proprio-data) lookups)))))

 

 

 (defn longest-thread

   "Find the longest thread from phi-index-sets. The index sets should

    be ordered from most recent to least recent."

   [phi-index-sets]

   (loop [result '()

          [thread-bases & remaining :as phi-index-sets] phi-index-sets]

     (if (empty? phi-index-sets)

       (vec result)

       (let [threads

             (for [thread-base thread-bases]

               (loop [thread (list thread-base)

                      remaining remaining]

                 (let [next-index (dec (first thread))]

                   (cond (empty? remaining) thread

                         (contains? (first remaining) next-index)

                         (recur

                          (cons next-index thread) (rest remaining))

                         :else thread))))

             longest-thread

             (reduce (fn [thread-a thread-b]

                       (if (> (count thread-a) (count thread-b))

                         thread-a thread-b))

                     '(nil)

                     threads)]

         (recur (concat longest-thread result)

                (drop (count longest-thread) phi-index-sets))))))

 

 

 (defn init []

   (def phi-space (generate-phi-space))

   (def phi-scan (gen-phi-scan phi-space))

   )

 

 ;; (defn infer-nils-dyl [s]

 ;;   (loop [closed ()

 ;;          open s

 ;;          anchor 0]

 ;;     (if-not (empty? open)

 ;;       (recur (conj closed

 ;;                    (or (peek open)

 ;;                        anchor))

 ;;              (pop open)

 ;;              (or (peek open) anchor))

 ;;        closed)))

       

 ;; (defn infer-nils [s]

 ;;   (for [i (range (count s))]

 ;;     (or (get s i)

 ;;         (some (comp not nil?) (vector:last-n (- (count s) i)))

 ;;         0)))

 

 

 (defn infer-nils

   "Replace nils with the next available non-nil element in the

    sequence, or barring that, 0."

   [s]

   (loop [i (dec (count s))

          v (transient s)]

     (if (zero? i) (persistent! v)

         (if-let [cur (v i)]

           (if (get v (dec i) 0)

             (recur (dec i) v)

             (recur (dec i) (assoc! v (dec i) cur)))

           (recur i (assoc! v i 0))))))

 

 ;; tests

 

 ;;(infer-nils [1 nil 1 1]) [1 1 1 1]

 ;;(infer-nils [1 1 1 nil]) [1 1 1 0]

 ;;(infer-nils [nil 2 1 1]) [2 2 1 1]       

   

 

 (defn debug-experience-phi []

   (let [proprio (atom ())]

     (fn

       [experiences]

       (let [phi-indices (phi-scan (:proprioception (peek experiences)))]

         (swap! proprio (partial cons phi-indices))

         (let [exp-thread (longest-thread (take 300 @proprio))

               phi-loop (mapv phi-space (infer-nils exp-thread))]

           (println-repl (vector:last-n exp-thread 22))

           (cond

            (grand-circle? phi-loop) (println "Grand Circle")

            (curled? phi-loop)       (println "Curled")

            (wiggling? phi-loop)     (println "Wiggling")

            (resting? phi-loop)      (println "Resting")

            :else                    (println "Unknown")))))))

 

 

 (defn init-interactive []

   (def phi-space

     (let [experiences (atom [])]

       (run-world

        (apply-map 

         worm-world

         (merge

          (worm-world-defaults)

          {:experiences experiences})))

       @experiences))

   (def phi-scan (gen-phi-scan phi-space)))

     

 

 (defn run-experiment-1 []

   (.start (worm-world :experience-watch (debug-experience-phi))))