;; lib/datalog/tests/aggregates.sx — count / sum / min / max.

(define dl-at-pass 0)
(define dl-at-fail 0)
(define dl-at-failures (list))

(define
  dl-at-deep=?
  (fn
    (a b)
    (cond
      ((and (list? a) (list? b))
        (and (= (len a) (len b)) (dl-at-deq-l? a b 0)))
      ((and (dict? a) (dict? b))
        (let ((ka (keys a)) (kb (keys b)))
          (and (= (len ka) (len kb)) (dl-at-deq-d? a b ka 0))))
      ((and (number? a) (number? b)) (= a b))
      (else (equal? a b)))))

(define
  dl-at-deq-l?
  (fn
    (a b i)
    (cond
      ((>= i (len a)) true)
      ((not (dl-at-deep=? (nth a i) (nth b i))) false)
      (else (dl-at-deq-l? a b (+ i 1))))))

(define
  dl-at-deq-d?
  (fn
    (a b ka i)
    (cond
      ((>= i (len ka)) true)
      ((let ((k (nth ka i)))
         (not (dl-at-deep=? (get a k) (get b k))))
        false)
      (else (dl-at-deq-d? a b ka (+ i 1))))))

(define
  dl-at-set=?
  (fn
    (a b)
    (and
      (= (len a) (len b))
      (dl-at-subset? a b)
      (dl-at-subset? b a))))

(define
  dl-at-subset?
  (fn
    (xs ys)
    (cond
      ((= (len xs) 0) true)
      ((not (dl-at-contains? ys (first xs))) false)
      (else (dl-at-subset? (rest xs) ys)))))

(define
  dl-at-contains?
  (fn
    (xs target)
    (cond
      ((= (len xs) 0) false)
      ((dl-at-deep=? (first xs) target) true)
      (else (dl-at-contains? (rest xs) target)))))

(define
  dl-at-test!
  (fn
    (name got expected)
    (if
      (dl-at-deep=? got expected)
      (set! dl-at-pass (+ dl-at-pass 1))
      (do
        (set! dl-at-fail (+ dl-at-fail 1))
        (append!
          dl-at-failures
          (str
            name
            "\n    expected: " expected
            "\n    got:      " got))))))

(define
  dl-at-test-set!
  (fn
    (name got expected)
    (if
      (dl-at-set=? got expected)
      (set! dl-at-pass (+ dl-at-pass 1))
      (do
        (set! dl-at-fail (+ dl-at-fail 1))
        (append!
          dl-at-failures
          (str
            name
            "\n    expected (set): " expected
            "\n    got:            " got))))))

(define
  dl-at-throws?
  (fn
    (thunk)
    (let
      ((threw false))
      (do
        (guard
          (e (#t (set! threw true)))
          (thunk))
        threw))))

(define
  dl-at-run-all!
  (fn
    ()
    (do
      ;; count
      (dl-at-test-set! "count siblings"
        (dl-query
          (dl-program
            "parent(p, bob). parent(p, alice). parent(p, charlie).
             sibling(X, Y) :- parent(P, X), parent(P, Y), !=(X, Y).
             sib_count(N) :- count(N, S, sibling(bob, S)).")
          (list (quote sib_count) (quote N)))
        (list {:N 2}))

      ;; sum
      (dl-at-test-set! "sum prices"
        (dl-query
          (dl-program
            "price(apple, 5). price(pear, 7). price(plum, 3).
             total(T) :- sum(T, X, price(F, X)).")
          (list (quote total) (quote T)))
        (list {:T 15}))

      ;; min
      (dl-at-test-set! "min score"
        (dl-query
          (dl-program
            "score(alice, 80). score(bob, 65). score(carol, 92).
             lo(M) :- min(M, S, score(P, S)).")
          (list (quote lo) (quote M)))
        (list {:M 65}))

      ;; max
      (dl-at-test-set! "max score"
        (dl-query
          (dl-program
            "score(alice, 80). score(bob, 65). score(carol, 92).
             hi(M) :- max(M, S, score(P, S)).")
          (list (quote hi) (quote M)))
        (list {:M 92}))

      ;; count over derived relation (stratification needed).
      (dl-at-test-set! "count over derived"
        (dl-query
          (dl-program
            "parent(a, b). parent(a, c). parent(b, d). parent(c, e).
             ancestor(X, Y) :- parent(X, Y).
             ancestor(X, Z) :- parent(X, Y), ancestor(Y, Z).
             num_ancestors(N) :- count(N, X, ancestor(a, X)).")
          (list (quote num_ancestors) (quote N)))
        (list {:N 4}))

      ;; count with no matches → 0.
      (dl-at-test-set! "count empty"
        (dl-query
          (dl-program
            "p(1). p(2).
             zero(N) :- count(N, X, q(X)).")
          (list (quote zero) (quote N)))
        (list {:N 0}))

      ;; sum with no matches → 0.
      (dl-at-test-set! "sum empty"
        (dl-query
          (dl-program
            "p(1). p(2).
             total(T) :- sum(T, X, q(X)).")
          (list (quote total) (quote T)))
        (list {:T 0}))

      ;; min with no matches → rule does not fire.
      (dl-at-test-set! "min empty"
        (dl-query
          (dl-program
            "p(1). p(2).
             lo(M) :- min(M, X, q(X)).")
          (list (quote lo) (quote M)))
        (list))

      ;; Aggregate with comparison filter on result.
      (dl-at-test-set! "popularity threshold"
        (dl-query
          (dl-program
            "post(p1). post(p2).
             liked(u1, p1). liked(u2, p1). liked(u3, p1).
             liked(u1, p2). liked(u2, p2).
             popular(P) :- post(P), count(N, U, liked(U, P)), >=(N, 3).")
          (list (quote popular) (quote P)))
        (list {:P (quote p1)}))

      ;; findall: collect distinct values into a list.
      (dl-at-test-set! "findall over EDB"
        (dl-query
          (dl-program
            "p(a). p(b). p(c).
             all_p(L) :- findall(L, X, p(X)).")
          (list (quote all_p) (quote L)))
        (list {:L (list (quote a) (quote b) (quote c))}))

      (dl-at-test-set! "findall over derived"
        (dl-query
          (dl-program
            "parent(a, b). parent(b, c). parent(c, d).
             ancestor(X, Y) :- parent(X, Y).
             ancestor(X, Z) :- parent(X, Y), ancestor(Y, Z).
             desc(L) :- findall(L, X, ancestor(a, X)).")
          (list (quote desc) (quote L)))
        (list {:L (list (quote b) (quote c) (quote d))}))

      (dl-at-test-set! "findall empty"
        (dl-query
          (dl-program
            "p(1).
             all_q(L) :- findall(L, X, q(X)).")
          (list (quote all_q) (quote L)))
        (list {:L (list)}))

      ;; Aggregate vs single distinct.
      ;; Group-by via aggregate-in-rule-body. Per-user friend count
      ;; over a friends relation. The U var is bound by the prior
      ;; positive lit u(U) so the aggregate counts only U-rooted
      ;; friends per group.
      (dl-at-test-set! "group-by per-user friend count"
        (dl-query
          (dl-program
            "u(alice). u(bob). u(carol).
             f(alice, x). f(alice, y). f(bob, x).
             counts(U, N) :- u(U), count(N, X, f(U, X)).")
          (list (quote counts) (quote U) (quote N)))
        (list
          {:U (quote alice) :N 2}
          {:U (quote bob) :N 1}
          {:U (quote carol) :N 0}))

      ;; Stratification: recursion through aggregation is rejected.
      ;; Indirect recursion through aggregation also rejected.
      ;; q -> r (via positive lit), r -> q (via aggregate body).
      ;; The aggregate edge counts as negation for stratification.
      (dl-at-test! "indirect agg cycle rejected"
        (dl-at-throws?
          (fn ()
            (let ((db (dl-make-db)))
              (do
                (dl-add-rule! db
                  {:head (list (quote q) (quote N))
                   :body (list (list (quote r) (quote N)))})
                (dl-add-rule! db
                  {:head (list (quote r) (quote N))
                   :body (list (list (quote count) (quote N) (quote X)
                                     (list (quote q) (quote X))))})
                (dl-saturate! db)))))
        true)

      (dl-at-test! "agg recursion rejected"
        (dl-at-throws?
          (fn ()
            (let ((db (dl-make-db)))
              (do
                (dl-add-rule! db
                  {:head (list (quote q) (quote N))
                   :body (list (list (quote count) (quote N) (quote X)
                                     (list (quote q) (quote X))))})
                (dl-saturate! db)))))
        true)

      ;; Negation + aggregation in the same body — different strata.
      (dl-at-test-set! "neg + agg coexist"
        (dl-query
          (dl-program
            "u(a). u(b). u(c). banned(b).
             active(X) :- u(X), not(banned(X)).
             cnt(N) :- count(N, X, active(X)).")
          (list (quote cnt) (quote N)))
        (list {:N 2}))

      ;; Min over a derived empty relation: no result.
      (dl-at-test-set! "min over empty derived"
        (dl-query
          (dl-program
            "s(50). s(60).
             score(N) :- s(N), >(N, 100).
             low(M) :- min(M, X, score(X)).")
          (list (quote low) (quote M)))
        (list))

      ;; Aggregates as the top-level query goal (regression for
      ;; dl-match-lit aggregate dispatch and projection cleanup).
      (dl-at-test-set! "count as query goal"
        (dl-query
          (dl-program "p(1). p(2). p(3). p(4).")
          (list (quote count) (quote N) (quote X) (list (quote p) (quote X))))
        (list {:N 4}))

      (dl-at-test-set! "findall as query goal"
        (dl-query
          (dl-program "p(1). p(2). p(3).")
          (list (quote findall) (quote L) (quote X)
                (list (quote p) (quote X))))
        (list {:L (list 1 2 3)}))

      (dl-at-test-set! "distinct counted once"
        (dl-query
          (dl-program
            "rated(alice, x). rated(alice, y). rated(bob, x).
             rater_count(N) :- count(N, U, rated(U, F)).")
          (list (quote rater_count) (quote N)))
        (list {:N 2})))))

(define
  dl-aggregates-tests-run!
  (fn
    ()
    (do
      (set! dl-at-pass 0)
      (set! dl-at-fail 0)
      (set! dl-at-failures (list))
      (dl-at-run-all!)
      {:passed dl-at-pass
       :failed dl-at-fail
       :total (+ dl-at-pass dl-at-fail)
       :failures dl-at-failures})))