   (:require [edu.berkeley.ai.util :as util]
             [edu.berkeley.ai.util  [graphs :as graphs] ]
             [exp [env :as env]  [hierarchy :as hierarchy] 
-                 [sas :as sas] [sas-analysis :as sas-analysis]]))
+                 [sas :as sas] [sas-analysis :as sas-analysis]])
+  (:import [java.util HashMap]))
 
 
 ; For now, pretend analysis is free, just brute-force it. 


;; TODO: also look at using landmarks to structure search ? 
; (but where's the state abstraction?)

;; NOTE: greedy should come in when we choose a parent for a var. 
; If some parent can use current var RIGHT NOW, should assign to it, not branch.

   0))


; Try to make an action that greedily fires the action scheduled on effect-var,
; effectively representing a composition of set-parent and fire-action actions.
(defn make-greedy-fire-action [s effect-var]
  (let [{:keys [name precond-map effect-map reward] :as factored-primitive}
          (->> effect-var action-var (env/get-var s))
        precond-vars (keys (dissoc precond-map effect-var))
        [free-pv unfree-pv] (util/separate #(env/get-var s (free-var %)) precond-vars)]
    (when (and (env/state-matches-map? s precond-map)
               (every? #(env/get-var s (parent-var % effect-var)) unfree-pv))
        (assert (every? #(nil? (env/get-var s (action-var %))) precond-vars))
        (env/FactoredPrimitive
         [::GreedyFire name]
         (into precond-map 
               (concat [[(action-var effect-var) factored-primitive]]
                       (for [v free-pv]   [(free-var v)              true])
                       (for [v unfree-pv] [(parent-var v effect-var) true])))
         (into effect-map 
               (concat [[(action-var effect-var) nil]]
                       (for [v unfree-pv] [(free-var v)              true])
                       (for [v unfree-pv] [(parent-var v effect-var) false])))
         0))))



;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;; Primitive Environment ;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; Helpers ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defn backward-reachable-nodes-and-preds [#^HashMap cache simple-dtgs var-name to-val]
  (util/cache-with cache [:backward var-name to-val]
    (graphs/compute-reachable-nodes-and-necessary-predecessors
     (map reverse (util/safe-get simple-dtgs var-name)) to-val)))

(defn possibly-acyclic-successors
  "Return a list of successors of cur-val, which can potentially lead to dst-val
   without a cycle."  
  [cache simple-dtgs var cur-val dst-val]
  (let [backward-sets (backward-reachable-nodes-and-preds cache simple-dtgs var dst-val)]
    (for [[s t] (util/safe-get simple-dtgs var)
          :when (and (= s cur-val) (not (contains? (backward-sets t) s)))]
      t)))

; TODO: take reduced arcs into account, etc.  
(defn activation-actions 
  "Return a list of activation actions for var v, ideally which are supported by 
   current state; i.e., should take reduced arcs into account, etc. "
  [s child-map v]
  (for [c (util/safe-get child-map v)]
    (make-set-parent-var-action v c)))

(defn current-child [s child-var-map p-var]
  (when-not (env/get-var s (free-var p-var))
    (some #(env/get-var s (parent-var p-var %))
          (child-var-map p-var))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;; Actions fn, actual env ;;;;;;;;;;;;;;;;;;;;;;;;;;;

; Simplest way to express mechanics (in order of preference):

;  Schematics:
;  a* => A   goes to a* -> a   (greedy fire)
;  a => A    goes to  A => A   (bottom-up action)
;  a -> A    goes to  a => A   (bottom-up activate)
;  a => a => goes to => A =>   (top-down action)          
;  a => a -> goes to a => a => (top-down activate)          

; Ideally, should prefer top-most top-down, bottom-most bottom-up, or something...


;; We make this a real type to make accessing precomputed stuff easier...
(deftype ASPlanEnv [init actions-fn g-map]
  env/Env 
    (initial-state [] init)
    (actions-fn    [] actions-fn)
    (goal-fn       [] #(when (env/state-matches-map? % g-map) (env/solution-and-reward %)))
  env/FactoredEnv
    (goal-map      [] g-map))

(defn make-asplan-env [sas-problem] 
  (let [causal-graph  (remove #(apply = %) (sas-analysis/standard-causal-graph sas-problem))
        child-var-map (util/map-vals #(map second %) (util/group-by first causal-graph))
        vars          (keys (:vars sas-problem))
        dtgs          (sas-analysis/domain-transition-graphs sas-problem)
        simple-dtgs   (util/map-vals (fn [dtg] (for [[pval emap] dtg, [eval _] emap] [pval eval])) dtgs)
        acyclic-succ-fn (partial possibly-acyclic-successors (HashMap.) simple-dtgs)]
    (assert (graphs/dag? causal-graph))    
    (ASPlanEnv 
     (expand-initial-state (env/initial-state sas-problem) child-var-map (goal-action dtgs))
     (fn asplan-actions [s]
       (let [[aa-vars na-vars] (split-with #(env/get-var s (action-var %)) vars)
             aa-parent-edges   (for [av aa-vars
                                      :let [pm (:precond-map (env/get-var s (action-var av)))]
                                      pv    (remove #{av} (keys pm))
                                      :when (and (not (env/get-var s (action-var pv)))
                                                 (not (= (pm pv) (env/get-var s pv))))]
                                  [pv av])
             na-tuples         (for [nav na-vars
                                     :let  [child (current-child s child-var-map nav)]
                                     :when (and child (not (env/get-var s (action-var child))))]
                                 [nav child])]
        (util/cond-let [x]
          ;; Greedy action -- all preconditions satisfied and not assigned elsewhere
          (some #(make-greedy-fire-action s %) aa-vars) 
            [x]

          ;; Active bottom-up var -- assigned, needs its value changed.               
          (util/find-first (fn [[pv cv]] (env/get-var s (parent-var pv cv))) aa-parent-edges)
            (let [[pv cv] x, 
                  p-val (env/get-var s pv)
                  d-val (util/safe-get-in (env/get-var s (action-var cv)) [:precond-map pv])]
              (apply concat (map (get-in dtgs [pv p-val]) (acyclic-succ-fn pv p-val d-val))))

          ;; Inactive bottom-up var -- needs to be assigned.  
          (first aa-parent-edges)
            (activation-actions s child-var-map (first x))

          ;; Active top-down var -- add actions
          (util/find-first #(not (env/get-var s (free-var %))) (map second na-tuples))  
            (apply concat (vals (util/safe-get-in dtgs [x (env/get-var s x)])))

          ;; Inactive top-down var -- activate it
          (first (map second na-tuples))
            (activation-actions s child-var-map x))))
     (env/goal-map sas-problem))))

(defn asplan-solution-name [sol]
  (map second (filter #(= (first %) ::GreedyFire) (map env/action-name sol))))

(defn asplan-solution-pair-name [[sol rew]]
  [(asplan-solution-name sol) rew])



;; TODO: cutoff when top-down and bottom-up meet, don't match ? 

;; TODO: Take advantage of "greedy-chain" condition, don't assign


; (use '[exp env hierarchy taxi-infinite ucs asplan hierarchical-incremental-search] 'edu.berkeley.ai.util)
; (let [e (make-random-infinite-taxi-sas2 1 2 1 2)] (println (run-counted #(uniform-cost-search e))) (println (run-counted #(asplan-solution-pair-name (debug 0 (sahucs-fast-flat (make-asplan-henv e )))))))




;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;;;;;;;;;;;;;;;;;;;; First attempt: "Skip" hierarchy ;;;;;;;;;;;;;;;;;;;;;;;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;


;; OK, now we need a hierarchy.    Problems:
