   (: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. 

   0))


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; Flat Hierarchy HLAs ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;


; TODO: ideally, we should avoid cyclies in the value of effect-var,
 ; NOT in whole state-abstracted states.

;; TODO: add a way to stop, when we reach the right var val. 

;; TODO: make sure top-down setup is right, e.g., we always add a target 
; action, then we can work on the precondition

;; TODO: is there an invariant that whenever a var is "truly active" it 
; has an active var, or is at right val?  If so, we need to enforce this as 
; we activate a top-down var, since it'll appear dead.  


; Simplest way to express mechanics:
;  If any action can greedily fire,
;    fire it.
;  Else, if any vars without actions but with explicit child with action and wrong value,
;    add a directed action.
;  Else, if any vars without actions or explicit children but have child with action,
;    activate precondition
;  Else, if any var without action but with explicit parent, 
;    add an undirected action.

;  Schematics:
;  a => A  goes to  A => A  or a -> a
;  a -> A  goes to  a => A
;  a => a  goes to  a => A
;  A => a  is not allowed ? 



(defn current-child [hierarchy s p-var]
  (some #(env/get-var s (parent-var p-var %))
        (util/safe-get-in hierarchy [:child-var-map p-var])))

(defn possible-next-actions-to [hierarchy var c-val d-val]
  (apply concat 
    (map (util/safe-get-in hierarchy [:dtgs var c-val])
         ((:acyclic-edges-fn hierarchy) var c-val d-val))))

(defn possible-next-actions [hierarchy var c-val]
  (apply concat (vals (util/safe-get-in hierarchy [:dtgs var c-val]))))

; Add any applicable action, using the target val if available.
; If not, goal is to ease deadlock by using a parent.
; TODO: use parent commitments to decide where to head ?

; Ideally, we stop as soon as we've used up at least one parent.
; One idea is just to stop after each action.
; This means we must take care when doing top-down; activation may
; not be required. 
(comment 
(deftype AddSomeActionHLA [hierarchy effect-var]
  env/Action
   (action-name [] [::AddSomeAction effect-var])
   (primitive?  [] false)
  env/ContextualAction 
   (precondition-context [s] 
     (let [cc (util/make-safe (current-child hierarchy s effect-var))]
       #{effect-var (parent-var effect-var cc) (action-var cc)}))
  hierarchy/HighLevelAction
   (immediate-refinements- [s] 
     (assert (not (env/get-var s (action-var effect-var))))
     (let [cc    (util/make-safe (current-child hierarchy s effect-var))
           c-val (env/get-var s effect-var)]
       (if-let [ca (env/get-var s (action-var cc))] ;; driven bottom-up?
         (let [d-val (util/safe-get-in ca :precond-map effect-var)]
           (if (= c-val d-val)
               [[]]) 
             (for [a (possible-next-action-to hierarchy effect-var c-val d-val)]
               [(make-add-action-action a)]))
         (for [a (possible-next-action hierarchy effect-var c-val)]
           [(make-add-action-action a)]))))
   (cycle-level-           [s] nil))
 ;)

; Choose a parent for a var; assumes any greedy vars have already
; fired, so all parents are allowed.  Prereq for adding actions on a var.
(deftype ActivateVarHLA [hierarchy effect-var]
  env/Action
   (action-name [] [::ActivateVar effect-var])
   (primitive?  [] false)
  env/ContextualAction 
   (precondition-context [s] #{})
  hierarchy/HighLevelAction
   (immediate-refinements- [s]
     (assert (env/get-var s (free-var effect-var)))
     (let [asa (AddSomeActionHLA hierarchy effect-var)]
      (for [c (util/safe-get-in hierarchy [:child-var-map effect-var])]
        [(make-set-parent-var-action effect-var c) asa])))
   (cycle-level-           [s] nil))


;; TODO: assert is wrong.?

(defn make-greedy-fire-plan
  "Return an HLA to greedily fire, or nil if cannot greedily fire."
  [hierarchy s effect-var]
  (let [a            (->> effect-var action-var (env/get-var s))
        _            (assert a)
        pm           (:precond-map a)
        precond-vars (keys (dissoc pm effect-var))
        [free-pv unfree-pv] (util/separate #(env/get-var s (free-var %)) precond-vars)]
    (when (and (env/state-matches-map? s pm)
               (every? #(env/get-var s (parent-var % effect-var)) unfree-pv))
      (assert (every? #(nil? (env/get-var s (action-var %))) precond-vars))
      (let [fire (make-fire-action-action a)]
        [(hierarchy/SimpleHLA [::GreedyFire ] 
          (env/precondition-context fire s)
          [(conj (vec (for [p free-pv] (make-set-parent-var-action p effect-var))) 
                 fire
                 (AddSomeActionHLA hierarchy effect-var))])]))))

(defn make-a-greedy-fire-plan [hierarchy s]
  (some #(make-greedy-fire-plan hierarchy s %) 
        (filter #(env/get-var s (action-var %))
                (util/safe-get hierarchy :vars))))


(defn bottom-up-supported-vars [hierarchy s]
  (mapcat (fn [v] 
            (when-let [a (env/get-var s (action-var v))] 
              (for [[var val] (:precond-map a)
                    :when (not (= (env/get-var s var) val))]
                var)))
          (util/safe-get hierarchy :vars)))

(defn top-down-supported-vars [hierarchy s]
  (for [vp (util/safe-get hierarchy :vars)
        vc (util/safe-get-in hierarchy [:child-var-map vp])
        :when (env/get-var s (parent-var vp vc))]
    vc))

(defn choose-open-var 
  "Choose a var to start changing.  Var must be open, i.e. parent-var and action-var both nil.
   Var should be supported, either 
     -from below by open precondition of action on potential parent-var, or
     -from above by parent-var link to this var."
  [hierarchy s]
  (first (filter #(env/get-var s (free-var %)) 
                 (concat (bottom-up-supported-vars hierarchy s)
                         (top-down-supported-vars hierarchy s)))))

(deftype FlatASPlanHLA [hierarchy name pc] :as this
  env/Action
   (action-name [] name)
   (primitive?  [] false)
  env/ContextualAction 
   (precondition-context [s] pc)
  hierarchy/HighLevelAction
   (immediate-refinements- [s] 
     (if-let [g (make-a-greedy-fire-plan hierarchy s)] [(conj (vec g) this)]
      (if-let [v (choose-open-var hierarchy s)] [[(ActivateVarHLA hierarchy v) this]] 
        [[]])))
   (cycle-level-           [s] nil))

(defn make-flat-asplan-hla [hierarchy full-context]
  (FlatASPlanHLA hierarchy [::FlatASPlan] full-context))


;;;;;;;;;;;;;;;;;;;;;;;;; Flat Hierarchy precomputations ;;;;;;;;;;;;;;;;;;;;;;;

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

;; TODO: note, should be able to use acylic dtgs for, e.g., single parent. 

;; TODO: Take advantage of "greedy-chain" condition, don't assign
;   val to parent too high up the chain too early. 

;; Memoized partial computations to speed up acylic edge generation.
(def #^HashMap *forward-reachability-cache* nil)
(def #^HashMap *backward-reachability-cache* nil)

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

(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 get-possibly-acyclic-edges
  "Return a list of edges, where those provably on no acyclic s-t path have been eliminated.
   Do so in polynomial time, possibly missing some always-cyclic edges."  
  [cache simple-dtgs var from-val to-val]
  (let [forward-sets  (forward-reachable-nodes-and-preds cache simple-dtgs var from-val)
        backward-sets (backward-reachable-nodes-and-preds cache simple-dtgs var to-val)]
    (filter
     (fn [[a b]]
       (let [f (forward-sets a)
             b (backward-sets b)]
         (and f b (empty? (clojure.set/intersection f b)))))
     (util/safe-get simple-dtgs var))))


(defn get-child-var-map [sas-problem]
  (let [causal-graph (remove #(apply = %) (sas-analysis/standard-causal-graph sas-problem))]
    (assert (graphs/dag? causal-graph))
    (util/map-vals #(map second %) (util/group-by first causal-graph))))

; Cannot be used iwth safe-get...
;(deftype ASPlanFlatHierarchy [sas-problem vars dtgs child-var-map possible-next-map])

(defn make-asplan-env [sas-problem child-var-map dtgs]
  (reify env/Env
   (initial-state [] 
     (expand-initial-state (env/initial-state sas-problem) child-var-map (goal-action dtgs)))
   (actions-fn    [] :no-actions-fn)
   (goal-fn       [] (env/goal-fn sas-problem))))

(defn make-asplan-flat-henv [sas-problem] 
  (let [child-var-map (get-child-var-map 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)
        reach-cache   (HashMap.)
        env           (make-asplan-env sas-problem child-var-map dtgs)]
    (hierarchy/SimpleHierarchicalEnv 
     env
     [(make-flat-asplan-hla 
       {:type ASPlanFlatHierarchy
        :sas-problem      sas-problem 
        :vars             (keys (:vars sas-problem))
        :dtgs             dtgs
        :child-var-map    child-var-map
        :acyclic-edges-fn (partial get-possibly-acylic-edges cache simple-dtgs)} 
       (env/current-context (env/initial-state env)))])))

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

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

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



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

     {(action-var evar) factored-primitive}
     reward)))


(comment ;; Not actually needed -- use below instead.
  (defn make-fire-action-action [{:keys [name precond-map effect-map reward] :as factored-primitive}]
   (let [[evar eval] (util/safe-singleton (seq effect-map))]
     (env/FactoredPrimitive
      [::Fire name]
      (into precond-map 
            (cons [(action-var evar) factored-primitive]
                  (for [[pvar pval] (dissoc precond-map evar)]
                    [(parent-var pvar evar) true])))
      (into effect-map 
            (cons [(action-var evar) nil]
                  (apply concat
                         (for [[pvar pval] (dissoc precond-map evar)]
                           [[(parent-var pvar evar) false]
                            [(free-var evar) true]]))))
      0))))

(defn make-set-parent-var-action [p-var c-var]
  (env/FactoredPrimitive 
