;; lib/prolog/tests/string_agg.sx -- sub_atom/5 + aggregate_all/3

(define pl-sa-test-count 0)
(define pl-sa-test-pass 0)
(define pl-sa-test-fail 0)
(define pl-sa-test-failures (list))

(define
  pl-sa-test!
  (fn
    (name got expected)
    (begin
      (set! pl-sa-test-count (+ pl-sa-test-count 1))
      (if
        (= got expected)
        (set! pl-sa-test-pass (+ pl-sa-test-pass 1))
        (begin
          (set! pl-sa-test-fail (+ pl-sa-test-fail 1))
          (append!
            pl-sa-test-failures
            (str name "\n    expected: " expected "\n    got:      " got)))))))

(define
  pl-sa-goal
  (fn
    (src env)
    (pl-instantiate (nth (first (pl-parse (str "g :- " src "."))) 2) env)))

(define pl-sa-db (pl-mk-db))

(define
  pl-sa-num-val
  (fn (env key) (pl-num-val (pl-walk-deep (dict-get env key)))))

(define
  pl-sa-list-to-atoms
  (fn
    (t)
    (let
      ((w (pl-walk-deep t)))
      (cond
        ((and (pl-atom? w) (= (pl-atom-name w) "[]")) (list))
        ((and (pl-compound? w) (= (pl-fun w) ".") (= (len (pl-args w)) 2))
          (cons
            (pl-atom-name (first (pl-args w)))
            (pl-sa-list-to-atoms (nth (pl-args w) 1))))
        (true (list))))))

(define pl-sa-prog-src "member(X, [X|_]). member(X, [_|T]) :- member(X, T).")
(pl-db-load! pl-sa-db (pl-parse pl-sa-prog-src))

;; -- sub_atom/5 --

(pl-sa-test!
  "sub_atom ground: sub_atom(abcde,0,3,2,abc)"
  (pl-solve-once!
    pl-sa-db
    (pl-sa-goal "sub_atom(abcde, 0, 3, 2, abc)" {})
    (pl-mk-trail))
  true)

(pl-sa-test!
  "sub_atom ground: sub_atom(abcde,2,2,1,cd)"
  (pl-solve-once!
    pl-sa-db
    (pl-sa-goal "sub_atom(abcde, 2, 2, 1, cd)" {})
    (pl-mk-trail))
  true)

(pl-sa-test!
  "sub_atom ground mismatch fails"
  (pl-solve-once!
    pl-sa-db
    (pl-sa-goal "sub_atom(abcde, 0, 2, 3, cd)" {})
    (pl-mk-trail))
  false)

(pl-sa-test!
  "sub_atom empty sub at start"
  (pl-solve-once!
    pl-sa-db
    (pl-sa-goal "sub_atom(abcde, 0, 0, 5, '')" {})
    (pl-mk-trail))
  true)

(pl-sa-test!
  "sub_atom whole string"
  (pl-solve-once!
    pl-sa-db
    (pl-sa-goal "sub_atom(hello, 0, 5, 0, hello)" {})
    (pl-mk-trail))
  true)

(define pl-sa-env-b1 {})
(pl-solve-once!
  pl-sa-db
  (pl-sa-goal "sub_atom(abcde, B, 2, A, cd)" pl-sa-env-b1)
  (pl-mk-trail))
(pl-sa-test!
  "sub_atom bound SubAtom gives B=2"
  (pl-sa-num-val pl-sa-env-b1 "B")
  2)
(pl-sa-test!
  "sub_atom bound SubAtom gives A=1"
  (pl-sa-num-val pl-sa-env-b1 "A")
  1)

(define pl-sa-env-b2 {})
(pl-solve-once!
  pl-sa-db
  (pl-sa-goal "sub_atom(hello, B, L, A, ello)" pl-sa-env-b2)
  (pl-mk-trail))
(pl-sa-test! "sub_atom ello: B=1" (pl-sa-num-val pl-sa-env-b2 "B") 1)
(pl-sa-test! "sub_atom ello: L=4" (pl-sa-num-val pl-sa-env-b2 "L") 4)
(pl-sa-test! "sub_atom ello: A=0" (pl-sa-num-val pl-sa-env-b2 "A") 0)

(pl-sa-test!
  "sub_atom ab: 6 total solutions"
  (let
    ((env {}))
    (pl-solve-once!
      pl-sa-db
      (pl-sa-goal "aggregate_all(count, sub_atom(ab, _, _, _, _), N)" env)
      (pl-mk-trail))
    (pl-sa-num-val env "N"))
  6)

(pl-sa-test!
  "sub_atom a: 3 total solutions"
  (let
    ((env {}))
    (pl-solve-once!
      pl-sa-db
      (pl-sa-goal "aggregate_all(count, sub_atom(a, _, _, _, _), N)" env)
      (pl-mk-trail))
    (pl-sa-num-val env "N"))
  3)

;; -- aggregate_all/3 --

(pl-sa-test!
  "aggregate_all count member [a,b,c] = 3"
  (let
    ((env {}))
    (pl-solve-once!
      pl-sa-db
      (pl-sa-goal "aggregate_all(count, member(_, [a,b,c]), N)" env)
      (pl-mk-trail))
    (pl-sa-num-val env "N"))
  3)

(pl-sa-test!
  "aggregate_all count fail = 0"
  (let
    ((env {}))
    (pl-solve-once!
      pl-sa-db
      (pl-sa-goal "aggregate_all(count, fail, N)" env)
      (pl-mk-trail))
    (pl-sa-num-val env "N"))
  0)

(pl-sa-test!
  "aggregate_all count always succeeds"
  (pl-solve-once!
    pl-sa-db
    (pl-sa-goal "aggregate_all(count, fail, _)" {})
    (pl-mk-trail))
  true)

(define pl-sa-env-bag1 {})
(pl-solve-once!
  pl-sa-db
  (pl-sa-goal "aggregate_all(bag(X), member(X, [a,b,c]), L)" pl-sa-env-bag1)
  (pl-mk-trail))
(pl-sa-test!
  "aggregate_all bag [a,b,c]"
  (pl-sa-list-to-atoms (dict-get pl-sa-env-bag1 "L"))
  (list "a" "b" "c"))

(define pl-sa-env-bag2 {})
(pl-solve-once!
  pl-sa-db
  (pl-sa-goal "aggregate_all(bag(X), member(X, []), L)" pl-sa-env-bag2)
  (pl-mk-trail))
(pl-sa-test!
  "aggregate_all bag empty goal = []"
  (pl-sa-list-to-atoms (dict-get pl-sa-env-bag2 "L"))
  (list))

(pl-sa-test!
  "aggregate_all sum [1,2,3,4] = 10"
  (let
    ((env {}))
    (pl-solve-once!
      pl-sa-db
      (pl-sa-goal "aggregate_all(sum(X), member(X, [1,2,3,4]), S)" env)
      (pl-mk-trail))
    (pl-sa-num-val env "S"))
  10)

(pl-sa-test!
  "aggregate_all max [3,1,4,1,5,9,2,6] = 9"
  (let
    ((env {}))
    (pl-solve-once!
      pl-sa-db
      (pl-sa-goal "aggregate_all(max(X), member(X, [3,1,4,1,5,9,2,6]), M)" env)
      (pl-mk-trail))
    (pl-sa-num-val env "M"))
  9)

(pl-sa-test!
  "aggregate_all max empty fails"
  (pl-solve-once!
    pl-sa-db
    (pl-sa-goal "aggregate_all(max(X), member(X, []), M)" {})
    (pl-mk-trail))
  false)

(pl-sa-test!
  "aggregate_all min [3,1,4,1,5,9,2,6] = 1"
  (let
    ((env {}))
    (pl-solve-once!
      pl-sa-db
      (pl-sa-goal "aggregate_all(min(X), member(X, [3,1,4,1,5,9,2,6]), M)" env)
      (pl-mk-trail))
    (pl-sa-num-val env "M"))
  1)

(pl-sa-test!
  "aggregate_all min empty fails"
  (pl-solve-once!
    pl-sa-db
    (pl-sa-goal "aggregate_all(min(X), member(X, []), M)" {})
    (pl-mk-trail))
  false)

(define pl-sa-env-set1 {})
(pl-solve-once!
  pl-sa-db
  (pl-sa-goal
    "aggregate_all(set(X), member(X, [b,a,c,a,b]), S)"
    pl-sa-env-set1)
  (pl-mk-trail))
(pl-sa-test!
  "aggregate_all set [b,a,c,a,b] = [a,b,c]"
  (pl-sa-list-to-atoms (dict-get pl-sa-env-set1 "S"))
  (list "a" "b" "c"))

(define pl-sa-env-set2 {})
(pl-solve-once!
  pl-sa-db
  (pl-sa-goal "aggregate_all(set(X), fail, S)" pl-sa-env-set2)
  (pl-mk-trail))
(pl-sa-test!
  "aggregate_all set fail = []"
  (pl-sa-list-to-atoms (dict-get pl-sa-env-set2 "S"))
  (list))

(pl-sa-test!
  "aggregate_all sum empty = 0"
  (let
    ((env {}))
    (pl-solve-once!
      pl-sa-db
      (pl-sa-goal "aggregate_all(sum(X), fail, S)" env)
      (pl-mk-trail))
    (pl-sa-num-val env "S"))
  0)

(define pl-string-agg-tests-run! (fn () {:failed pl-sa-test-fail :passed pl-sa-test-pass :total pl-sa-test-count :failures pl-sa-test-failures}))