(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)))

 

 (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-y" "key-h"

          "key-j" "key-u"

          "key-i" "key-k"

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

     {: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

   [[370 :d-up 40]

    [600 :d-up 0]])

 

 (def period 18)

 

 (def worm-muscle-labels

   [:base-up :base-down

    :a-down :a-up

    :b-up :b-down

    :c-down :c-up

    :d-up :d-down])

 

 (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-down :a-up])

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

 

 

 ;; 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.

 

 

 (defn straight?

   "Is the worm straight?"

   [experiences]

   (every?

    (fn [[_ _ bend]]

      (< (Math/sin bend) 0.05))

    (:proprioception (peek experiences))))

 

 (defn curled?

   "Is the worm curled up?"

   [experiences]

   (every?

    (fn [[_ _ bend]]

      (> (Math/sin bend) 0.64))

    (:proprioception (peek experiences))))

 

 (defn grand-circle?

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

   [experiences]

   (and (curled? experiences)

        true)) ;; TODO: add code here.

 

 (defn vector:last-n [v n]

   (let [c (count v)]

     (if (< c n) v

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

 

 (defn touch-average [[coords touch]]

   (/ (average (map first touch)) (average (map second touch))))

 

 (def worm-segment-touch-bottom

   [[8 15] [8 16] [8 17] [8 18] [8 19] [8 20] [8 21] [8 22] [9 15]

    [9 16] [9 17] [9 18] [9 19] [9 20] [9 21] [9 22] [10 15] [10 16]

    [10 17] [10 18] [10 19] [10 20] [10 21] [10 22] [11 15] [11 16]

    [11 17] [11 18] [11 19] [11 20] [11 21] [11 22] [12 15] [12 16]

    [12 17] [12 18] [12 19] [12 20] [12 21] [12 22] [13 15] [13 16]

    [13 17] [13 18] [13 19] [13 20] [13 21] [13 22] [14 15] [14 16]

    [14 17] [14 18] [14 19] [14 20] [14 21] [14 22]])

 

 

 

 (defn floor-contact [[coords contact :as touch]]

   (let [raw-average

         (average

          (map

           first

           (vals

            (select-keys

             (zipmap coords contact)

             ))))]

     (Math/abs (- 1. (* 10 raw-average)))))

 

 

 (defn wiggling?

   "Is the worm wiggling?"

   [experiences]

   (vector:last-n experiences 200)

   

 )

 

 (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)]

     {:view worm-side-view

      :motor-control (:motor-control direct-control)

      :keybindings (:keybindings direct-control)

      :record nil

      :experiences nil

      :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 worm-world

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

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

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

                 worm-model end-frame]}

         (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 10 1 10 :position (Vector3f. 0 -10 0)

                    :color ColorRGBA/Gray :mass 0)

         timer (IsoTimer. 60)]

 

     (world

        (nodify [worm floor])

        (merge standard-debug-controls keybindings)

        (fn [world]

          (position-camera world view)

          (.setTimer world timer)

          (display-dilated-time world timer)

          (if record

            (Capture/captureVideo

             world

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

          (speed-up world)

          (light-up-everything world))

        (fn [world tpf]

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

            (.stop world))

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

                proprioception-data (prop)

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

            (when experiences

              (record-experience!

               experiences {:touch touch-data

                            :proprioception proprioception-data

                            :muscle muscle-data})

              ;;(if (curled? @experiences) (println "Curled"))

              ;;(if (straight? @experiences)    (println "Straight"))

              ;; (println-repl

              ;;  (apply format "%.2f %.2f %.2f %.2f %.2f\n"

              ;;         (map floor-contact touch-data)))

              

              )

            (muscle-display

             muscle-data

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

            (prop-display

             proprioception-data

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

            (touch-display 

             touch-data

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

 

 

 ;; A demonstration of self organiging touch maps through experience. 

 

 (def single-worm-segment-view

   [(Vector3f. 2.0681207, -6.1406755, 1.6106138)

    (Quaternion. -0.15558705, 0.843615, -0.3428654, -0.38281822)])

 

 (def worm-single-segment-muscle-labels

   [:lift-1 :lift-2 :roll-1 :roll-2])

 

 (defn touch-kinesthetics []

   [[170 :lift-1 40]

    [190 :lift-1 20]

    [206 :lift-1  0]

 

    [400 :lift-2 40]

    [410 :lift-2  0]

 

    [570 :lift-2 40]

    [590 :lift-2 20]

    [606 :lift-2  0]

 

    [800 :lift-1 30]

    [809 :lift-1 0]

 

    [900 :roll-2 40]

    [905 :roll-2 20]

    [910 :roll-2  0]

 

    [1000 :roll-2 40]

    [1005 :roll-2 20]

    [1010 :roll-2  0]

    

    [1100 :roll-2 40]

    [1105 :roll-2 20]

    [1110 :roll-2  0]

    ])

 

 (defn single-worm-segment []

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

 

 (defn worm-segment-defaults []

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

     (merge (worm-world-defaults)

            {:worm-model single-worm-segment

             :view single-worm-segment-view

             :motor-control

             (motor-control-program

              worm-single-segment-muscle-labels

              (touch-kinesthetics))

             :end-frame 1200})))

 

 (def full-contact [(float 0.0) (float 0.1)])

 

 (defn pure-touch?

   "This is worm specific code to determine if a large region of touch

    sensors is either all on or all off."

   [[coords touch :as touch-data]]

   (= (set (map first touch)) (set full-contact)))

 

 (defn remove-similar

   [coll]

   (loop [result () coll (sort-by (comp - count) coll)]

     (if (empty? coll) result

         (let  [x  (first coll)

                xs (rest coll)

                c (count x)]

           (if (some

                (fn [other-set]

                  (let [oc (count other-set)]

                    (< (- (count (union other-set x)) c) (* oc 0.1))))

                xs)

             (recur result xs)

             (recur (cons x result) xs))))))

 

 

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

   (vec

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

          y (range y0 (inc y1))]

      [x y])))

 

 (def all-touch-coordinates

   (concat

    (rect-region [0  15] [7  22])

    (rect-region [8   0] [14 29])

    (rect-region [15 15] [22 22])))

 

 (defn view-touch-region [coords]

   (let [touched-region

         (reduce

          (fn [m k]

            (assoc m k [0.0 0.1]))

          (zipmap all-touch-coordinates (repeat [0.1 0.1])) coords)

         data

         [[(vec (keys touched-region)) (vec (vals touched-region))]]

         touch-display (view-touch)]

     (touch-display data)

     (touch-display data)))

 

 (defn learn-touch-regions []

   (let [experiences (atom [])

         world (apply-map

                worm-world

                (assoc (worm-segment-defaults)

                  :experiences experiences))]

     (run-world world)

     (->>

      @experiences

      (drop 175)

      ;; access the single segment's touch data

      (map (comp first :touch))

      ;; only deal with "pure" touch data to determine surfaces

      (filter pure-touch?)

      ;; associate coordinates with touch values

      (map (partial apply zipmap))

      ;; select those regions where contact is being made

      (map (partial group-by second))

      (map #(get % full-contact))

      (map (partial map first))

      ;; remove redundant/subset regions

      (map set)

      remove-similar)))

 

 (defn learn-and-view-touch-regions []

   (map view-touch-region

        (learn-touch-regions)))