 ; Seems we need all the angelic machinery we had before, in general case.
 
 ;; State variables are (all keyworded, note!):
-; :const {[size] [freespace] [legal-go] [goal o] [map]...}
+; :const {[size] [freespace] [legal-go] [goal o] [base-offsets] [n-base-samples] [map]...}
 ;    freespace is set of free [x y] pairs (for base).
 ;    legal-go is set of allowed [gxo gyo] pairs
-;    goal is [[minx miny] [maxx maxy]] range -- inclusive.
-; [bx], [by] - current "base" coordinates.
-; [gxo], [gyo] - gripper x and y offsets.  
-; [x o], [y o] - location of object
-; [at-goal? o] - 
-; [object-at x y] - nil, object, or :border
-; [holding]
+;    goal is now set of allowed positions. Previously: [[minx miny] [maxx maxy]] range -- inclusive.
+;    base-offsets is randomized list of offset positions to drasp/drop at a position.
+; [base],  current "base" coordinates.
+; [gripper-offset], gripper coordinate offsets
+; [pos o], location of object
+; [at-goal? o] - true/false 
+; [object-at [x y]] - nil, object, or :border
+; [holding] - object or nil
 ; [parked?]
 
 ; implicit obstacle border.

  (env/FactoredPrimitive [:park] {[:parked?] false} {[:parked?] true} -5))

(defn make-unpark [s]
  (env/FactoredPrimitive [:unpark] {[:parked?] true [:gxo] 0 [:gyo] 0} {[:parked?] false} -5))

(defn make-base-dir [s [dirname dx dy]]
  (let [const (env/get-var s :const)
        cx    (env/get-var s [:bx])
        cy    (env/get-var s [:by])
        nx    (+ cx dx)
        ny    (+ cy dy)]
    (when (contains? (get const [:freespace]) [nx ny])
      (env/FactoredPrimitive 
       [:base dirname]
       {[:bx] cx [:by] cy [:parked?] false}
       {[:bx] nx [:by] ny}
       -2))))

(defn make-gripper-dir [s [dirname dx dy]]
  (let [const (env/get-var s :const)
        cbx   (env/get-var s [:bx])
        cby   (env/get-var s [:by])
        cgxo  (env/get-var s [:gxo])
        cgyo  (env/get-var s [:gyo])        
        ngxo  (+ cgxo dx)
        ngyo  (+ cgyo dy)]
    (when (contains? (get const [:legal-go]) [ngxo ngyo])
      (env/FactoredPrimitive 
       [:gripper dirname]
       {[:bx] cbx [:by] cby [:gxo] cgxo [:gyo] cgyo [:parked?] true 
        [:object-at (+ cbx ngxo) (+ cby ngyo)] nil}
       {[:gxo] ngxo [:gyo] ngyo}
       -1))))

(defn make-pickup-dir [s [dirname dx dy]]
  (let [cbx   (env/get-var s [:bx])
        cby   (env/get-var s [:by])
        cgxo  (env/get-var s [:gxo])
        cgyo  (env/get-var s [:gyo])        
        ox    (+ cbx cgxo dx)
        oy    (+ cby cgyo dy)
        o     (env/get-var s [:object-at ox oy])]
    (when (and o (not (= o :border)))
      (env/FactoredPrimitive 
       [:pickup dirname]
       {[:bx] cbx [:by] cby [:gxo] cgxo [:gyo] cgyo [:parked?] true 
        [:x o] ox [:y o] oy [:holding] nil [:at-goal? o] false}
       {[:x o] nil [:y o] nil [:holding] o [:object-at ox oy] nil}
       -1))))

(defn region-contains? [[[minx miny] [maxx maxy]] [x y]]
  (and (<= minx x maxx) (<= miny y maxy)))

(defn make-putdown-dir [s [dirname dx dy]]
  (let [const (env/get-var s :const)
        o     (env/get-var s [:holding])
        cbx   (env/get-var s [:bx])
        cby   (env/get-var s [:by])
        cgxo  (env/get-var s [:gxo])
        cgyo  (env/get-var s [:gyo])        
        ox    (+ cbx cgxo dx)
        oy    (+ cby cgyo dy)]
    (when (and o (region-contains? (get const [:goal o]) [ox oy]))
      (env/FactoredPrimitive 
       [:putdown dirname]
       {[:bx] cbx [:by] cby [:gxo] cgxo [:gyo] cgyo [:parked?] true 
        [:holding] o [:object-at ox oy] nil}
       {[:x o] ox [:y o] oy [:holding] nil [:object-at ox oy] o [:at-goal? o] true}
       -1))))


; obstacles and goals are [[minx miny] [maxx maxy]]
; objects are [name [cx cy] goal] -- no goal means static 
; g-rad is gripper "radius" -- diamond shaped.

(defn region-cells [[[minx miny] [maxx maxy]]]
  (for [x (range minx (inc maxx)), y (range miny (inc maxy))] [x y]))

(defn make-2d-manipulation-env [size start obstacles objects g-rad]
  (let [[width height] size
        [bx by]        start
        border-regions [[[0 0] [0 height]] [[width 0] [width height]]
                        [[0 0] [width 0]] [[0 height] [width height]]]
        border-cells   (set (mapcat region-cells border-regions))
        obstacle-cells (set (mapcat region-cells obstacles))
        all-cells      (set (region-cells [[0 0] size]))
        freespace      (clojure.set/difference all-cells border-cells obstacle-cells)
        legal-go       (set (filter (fn [[x y]] (<= (+ (util/abs x) (util/abs y)) g-rad))
                                    (region-cells [[(- g-rad) (- g-rad)] [g-rad g-rad]])))]
    (env/SimpleFactoredEnv
     (into {:const (into {[:size] size [:freespace] freespace [:legal-go] legal-go [:objects] objects
                          [:map] (reduce (fn [m [x y]] (assoc-in m [y x] \.))
                                   (reduce (fn [m [x y]] (assoc-in m [y x] \*))
                                      (vec (repeat (inc height) (vec (repeat (inc width) \ ))))
                                      border-cells)
                                   obstacle-cells)}
                         (for [[o _ goal] objects :when goal] [[:goal o] goal]))
            [:bx] bx [:by] by [:gxo] 0 [:gyo] 0 [:holding] nil [:parked?] true}
           (apply concat
            (for [[x y] all-cells]     [[:object-at x y] nil])
            (for [[x y] border-cells] [[:object-at x y] :border])
            (for [[o [ox oy] goal] objects]
              [[[:x o] ox] [[:y o] oy] 
               [[:at-goal? o] (if goal false true)] 
               [[:object-at ox oy] o]])))
     (fn [s]
       (filter #(and % (env/applicable? % s))
          (apply concat

(defn print-state [s] (print (state-str s)))



(defn reach-context [s]
  (into #{[:bx] [:by] [:gxo] [:gyo]}
        (let [bx (env/get-var s [:bx])
              by (env/get-var s [:by])]
          (for [[xo yo] (util/safe-get (env/get-var s :const) [:legal-go])]
            [:object-at (+ bx xo) (+ by yo)]))))

(deftype ReachHLA [env dx dy] :as this
  env/Action                (action-name [] [:reach dx dy])
                            (primitive? [] false)
  env/ContextualAction      (precondition-context [s] (reach-context s))
  hierarchy/HighLevelAction (immediate-refinements- [s]
                             (if (and (= dx (env/get-var s [:bx])) 
                                      (= dy (env/get-var s [:by])))
                               [[]]
                               (for [dir dirs 
                                     :let [a (make-gripper-dir s dir)]
                                     :when (and a (env/applicable? a s))]
                                 [a this])))
                            (cycle-level- [s] 1))
; manhattan, SLD?  range?

(deftype GraspHLA [env o] :as this
  env/Action                (action-name [] [:grasp o])
                            (primitive? [] false)
  env/ContextualAction      (precondition-context [s] 
                              (conj (reach-context s) [:x o] [:y o] [:at-goal? o] [:holding]))
  hierarchy/HighLevelAction (immediate-refinements- [s]
                             (let [bx (env/get-var s [:bx]), by (env/get-var s [:by])
                                   ox (env/get-var s [:x o]), oy (env/get-var s [:y o])]
                               (for [[dirname dx dy] dirs]
                                 [(ReachHLA env (- ox dx) (- oy dy))
                                  (make-pickup dirname bx by (- ox dx) (- oy dy) o ox oy)])))
                            (cycle-level- [s] nil))
;manhattan
;pess: SLD?  


(deftype DropAtHLA [env o dx dy] :as this
  env/Action                (action-name [] [:dropat o dx dy])
                            (primitive? [] false)
  env/ContextualAction      (precondition-context [s] 
                              (conj (reach-context s) [:holding]))
  hierarchy/HighLevelAction (immediate-refinements- [s]
                             (let [bx (env/get-var s [:bx]), by (env/get-var s [:by])
                                   ox (env/get-var s [:x o]), oy (env/get-var s [:y o])]
                               (for [[dirname dx dy] dirs]
                                 [(ReachHLA env (- ox dx) (- oy dy))
                                  (make-pickup dirname bx by (- ox dx) (- oy dy) o ox oy)])))
                            (cycle-level- [s] nil))
;manhattan
;pess: SLD?  







;; TODO: should nav be done by external alg, (fairest between angelic/not) or explicit? 
(deftype NavHLA [env dx dy] :as this
  env/Action                (action-name [] [:nav dx dy])
                            (primitive? [] false)
  env/ContextualAction      (precondition-context [s] #{[:bx] [:by]})
  hierarchy/HighLevelAction (immediate-refinements- [s]
                             (if (and (= dx (env/get-var s [:bx])) 
                                      (= dy (env/get-var s [:by])))
                               [[]]
                               (for [dir dirs 
                                     :let [a (make-base-dir s dir)]
                                     :when (and a (env/applicable? a s))]
                                 [a this])))
                            (cycle-level- [s] 1))
; manhattan
; pess? SLD? (try direct line, if occ then -inf).

(deftype MoveBaseHLA [env dx dy] :as this
  env/Action                (action-name [] [:movebase dx dy])
                            (primitive? [] false)
  env/ContextualAction      (precondition-context [s] #{[:bx] [:by] [:parked?] [:gxo] [:gyo]})
  hierarchy/HighLevelAction (immediate-refinements- [s]
                             (let [bx (env/get-var s [:bx]), by (env/get-var s [:by])]
                               (if (and (= dx bx) (= dy by))
                                  [[]]
                                 [[(ReachHLA env bx by) (make-unpark s) 
                                   (NavHLA dx dy) (make-park s)]])))
                            (cycle-level- [s] nil))
; manhattan + park (if appl)
; pess? SLD?




; DropFrom
; Drop
; GoGrasp
; GoDropAt
; GoDrop
; MoveToGoal
; Act

;; Do we sample base, drop positions, or enumerate ? 




; (print-state (initial-state (make-2d-manipulation-env [10 10] [1 1] [ [ [4 4] [6 6] ] ] [ [:a [5 5] [ [4 4] [4 4] ] ] ] 2))) 