; lib/artdag/schedule.sx — relational scheduling on lib/minikanren CLP(FD).
; Each node gets a slot var in [1..max-slots]; every edge (input->node) imposes
; `fd-lt slot(input) slot(node)`. `fd-label` searches the finite domains; a solution
; is a {node-id -> slot} assignment respecting all dependencies. Grouping by slot
; gives parallel batches (plan.sx's batch shape). Labeling picks smallest slots
; first, so the FIRST solution is the ASAP leveling — it agrees with plan.sx's greedy
; Kahn waves; the relational extra is enumerating EVERY valid schedule. The
; parallelism cap is a cardinality property, enforced by filtering labeled solutions
; (the FD core handles precedence only). lib/minikanren is a READ-ONLY consumed
; substrate: make-var, fd-in, fd-lt, fd-label, mk-conj, reify, stream-take, empty-s.

(define
  artdag/range1
  (fn (n) (map (fn (i) (+ i 1)) (range 0 n))))

(define
  artdag/-zip-assoc
  (fn
    (ids vals)
    (reduce
      (fn (m p) (assoc m (first p) (nth p 1)))
      {}
      (zip ids vals))))

; build the constraint goal + the ordered slot vars for a dag over domain 1..maxslots.
(define
  artdag/sched-goal-and-vars
  (fn
    (dag maxslots)
    (let
      ((ids (artdag/dag-order dag)))
      (let
        ((vars (map (fn (id) (make-var)) ids)))
        (let
          ((id->var (artdag/-zip-assoc ids vars))
            (dom (artdag/range1 maxslots)))
          (let
            ((in-goals (map (fn (v) (fd-in v dom)) vars))
              (lt-goals
                (reduce
                  (fn
                    (acc id)
                    (concat
                      acc
                      (map
                        (fn
                          (inp)
                          (fd-lt (get id->var inp) (get id->var id)))
                        (artdag/node-inputs (artdag/dag-get dag id)))))
                  (list)
                  ids)))
            {:goal (apply mk-conj (concat in-goals lt-goals (list (fd-label vars)))) :vars vars :ids ids}))))))

(define
  artdag/-sched-solutions
  (fn
    (g limit)
    (map
      (fn (sol) (artdag/-zip-assoc (get g :ids) sol))
      (map
        (fn (s) (reify (get g :vars) s))
        (stream-take limit ((get g :goal) empty-s))))))

; all valid dependency-respecting slot assignments within 1..maxslots.
(define
  artdag/schedules
  (fn
    (dag maxslots)
    (artdag/-sched-solutions
      (artdag/sched-goal-and-vars dag maxslots)
      -1)))

; one valid assignment (ASAP within the bound), or nil if maxslots is too small.
(define
  artdag/schedule
  (fn
    (dag maxslots)
    (let
      ((ss (artdag/-sched-solutions (artdag/sched-goal-and-vars dag maxslots) 1)))
      (if (empty? ss) nil (first ss)))))

; ASAP schedule: node-count slots are always sufficient (a linear chain is the worst
; case), and smallest-first labeling yields the tightest leveling.
(define
  artdag/schedule-asap
  (fn (dag) (artdag/schedule dag (artdag/node-count dag))))

(define
  artdag/schedule-makespan
  (fn
    (assignment)
    (reduce
      (fn (m id) (max m (get assignment id)))
      0
      (keys assignment))))

; group node-ids by slot (ascending), each batch id-sorted for determinism.
(define
  artdag/schedule->batches
  (fn
    (dag assignment)
    (let
      ((mx (artdag/schedule-makespan assignment)))
      (filter
        (fn (b) (not (empty? b)))
        (map
          (fn
            (slot)
            (artdag/sort-strings
              (filter
                (fn (id) (= (get assignment id) slot))
                (keys assignment))))
          (artdag/range1 mx))))))

; independent check: every input is scheduled strictly before its consumer.
(define
  artdag/schedule-valid?
  (fn
    (dag assignment)
    (every?
      (fn
        (id)
        (every?
          (fn (inp) (< (get assignment inp) (get assignment id)))
          (artdag/node-inputs (artdag/dag-get dag id))))
      (artdag/dag-order dag))))

; schedules whose every slot holds <= cap nodes (parallelism cap as a post-filter).
(define
  artdag/schedules-capped
  (fn
    (dag maxslots cap)
    (filter
      (fn
        (asn)
        (every?
          (fn (b) (<= (len b) cap))
          (artdag/schedule->batches dag asn)))
      (artdag/schedules dag maxslots))))