(load "mk.scm")

(define print
  (lambda (exp)
    (display exp)
    (newline)))

(define peanoo
  (lambda (n)
    (conde
      [(== n 'z)]
      [(fresh (p)
        (== n `(s ,p))
        (peanoo p))])))

; ported from coq
; https://coq.inria.fr/library/Coq.Init.Peano.html
; https://coq.inria.fr/library/Coq.Numbers.Natural.Peano.NPeano.html

; Fixpoint divmod x y q u :=
;   match x with
;     | 0 => (q,u)
;     | S x´ => match u with
;                 | 0 => divmod x´ y (S q) y
;                 | S u´ => divmod x´ y q u´
;               end
;   end.
(define divmodo
  (lambda (x y q u qo uo)
    (conde
      [(== x 'z) (== q qo) (== u uo) (peanoo y) (peanoo q) (peanoo u)]
      [(fresh (x^)
        (== x `(s ,x^))
        (peanoo x^)
        (conde
          [(== u 'z) (divmodo x^ y `(s ,q) y qo uo)]
          [(fresh (u^)
            (== u `(s ,u^))
            (peanoo u^)
            (divmodo x^ y q u^ qo uo))]))])))

; Definition div x y :=
;   match y with
;     | 0 => y
;     | S y´ => fst (divmod x y´ 0 y´)
;   end.
(define divo 
  (lambda (x y out)
    (conde
      [(== y 'z) (== out y)]
      [(fresh (y^ uo)
        (== y `(s ,y^))
        (divmodo x y^ 'z y^ out uo))])))

; Fixpoint minus (n m:nat) : nat :=
;   match n, m with
;   | O, _ => n
;   | S k, O => n
;   | S k, S l => k - l
;   end
(define minuso
  (lambda (n m out)
    (conde
      [(== n 'z) (== out n)]
      [(fresh (n^)
        (== n `(s ,n^)) (== m 'z) (== out n))]
      [(fresh (n^ m^)
        (== n `(s ,n^)) (== m `(s ,m^))
        (minuso n^ m^ out))])))

; Definition modulo x y :=
;   match y with
;     | 0 => y
;     | S y´ => y´ - snd (divmod x y´ 0 y´)
;   end.
(define moduloo
  (lambda (x y out)
    (conde
      [(== y 'z) (== out y)]
      [(fresh (y^ qo uo)
        (== y `(s ,y^))
        (minuso y^ uo out)
        (divmodo x y^ 'z y^ qo uo))])))

(define fizz-buzzo
  (lambda (n out)
    (fresh (m3 m5)
      (peanoo n)
      (moduloo n '(s (s (s z))) m3)
      (moduloo n '(s (s (s (s (s z))))) m5)
      (conde
        [(== m3 'z) (=/= m5 'z) (== out 'fizz)]
        [(=/= m3 'z) (== m5 'z) (== out 'buzz)]
        [(== m3 'z) (== m5 'z) (== out 'fizzbuzz)]
        [(=/= m3 'z) (=/= m5 'z) (== out n)]))))

(define succo
  (lambda (n)
    (fresh (p)
      (== n `(s ,p))
      (peanoo n)
      (peanoo p))))

; (print (run 10 (q) (peanoo q)))
; (print (run 3 (x y q u qo uo) (divmodo x y q u qo uo)))
; (print (run 10 (x y out) (divo x y out)))
; (print (run 1 (q) (divo '(s (s (s (s z)))) '(s (s z)) q)))
; (print (run 1 (q) (minuso '(s (s (s (s z)))) '(s (s z)) q)))
; (print (run 3 (q) (moduloo '(s (s (s (s z)))) '(s (s z)) q)))
; (print (run 3 (q) (moduloo '(s z) '(s (s z)) q)))
; (print (run 10 (x y out) (moduloo x y out)))
; (print (run 5 (q) (moduloo q '(s (s (s z))) 'z)))
; (print (run 1 (q) (fizz-buzzo 'z q)))
; (print (run 1 (q) (fizz-buzzo '(s z) q)))
; (print (run 1 (q) (fizz-buzzo '(s (s z)) q)))
; (print (run 1 (q) (fizz-buzzo '(s (s (s z))) q)))
; (print (run 2 (q) (fizz-buzzo '(s (s (s z))) q)))
; this takes 11 minutes to run... why?
; (map print (run 100 (n q) (succo n) (fizz-buzzo n q)))
(map print (run 20 (n q) (succo n) (fizz-buzzo n q)))