(datatype Math
  (Num Rational)
  (Var String)
  (Add Math Math)
  (Div Math Math)
  (Mul Math Math))

(let zero (Num (rational 0 1)))
(let one (Num (rational 1 1)))
(let two (Num (rational 2 1)))

(rewrite (Add a b) (Add b a))
(rewrite (Add a zero) a)
(rewrite (Add (Num r1) (Num r2))
         (Num (+ r1 r2)))

(let one-two (Add one two))

(push)
(run 1)
;; yay, constant folding works
(check (= one-two (Num (rational 3 1))))
;; also, commutativity works
(check (= (Add two one) one-two))
(pop)

(push)
;; rule is like rewrite, but more general
;; the following rule doesn't union (Num r) with the result:
(rule ((Num r))
      ((union one (Div (Num r) (Num r)))))
;; uh oh, division by zero!
(run 1)

(pop)

;; we need to detect when things are non-zero
(function lower-bound (Math) Rational :merge (max old new))
(function upper-bound (Math) Rational :merge (min old new))

(rule ((Num r))
      ((set (lower-bound (Num r)) r)
       (set (upper-bound (Num r)) r)))
(rule ((= e (Add a b)) (= x (lower-bound a)) (= y (lower-bound b)))
      ((set (lower-bound e) (+ x y))))
(rule ((= e (Add a b)) (= x (upper-bound a)) (= y (upper-bound b)))
      ((set (upper-bound e) (+ x y))))
(rule ((= e (Mul a b)))
      ((set (lower-bound e)
            (min (* (lower-bound a) (lower-bound b))
                 (min (* (lower-bound a) (upper-bound b))
                 (min (* (upper-bound a) (lower-bound b))
                      (* (upper-bound a) (upper-bound b))))))
       (set (upper-bound e)
            (max (* (lower-bound a) (lower-bound b))
                 (max (* (lower-bound a) (upper-bound b))
                 (max (* (upper-bound a) (lower-bound b))
                      (* (upper-bound a) (upper-bound b))))))))

(rule ((= e (Add a b))
       (> (lower-bound e) (rational 0 1)))
      ((union one (Div (Add a b) (Add a b)))))

(let x (Var "x"))
(let x1 (Add x one))

(push)
(set (lower-bound x) (rational 0 1))
(set (upper-bound x) (rational 1 1))

(run 3)

(query-extract (lower-bound x1))
(query-extract (upper-bound x1))
(check (= one (Div x1 x1)))

(pop)


;; Set the variable x to a particular input value 200/201
(set (lower-bound x) (rational 200 201))
(set (upper-bound x) (rational 200 201))

(run 3)

(query-extract (lower-bound x1))
(query-extract (upper-bound x1))

(function true-value (Math) f64)

(rule ((= (to-f64 (lower-bound e))
          (to-f64 (upper-bound e))))
      ((set (true-value e)
            (to-f64 (lower-bound e)))))

(run 1)
(query-extract (true-value x1))

(function best-error (Math) f64 :merge new :default (to-f64 (rational 10000 1)))

(rule ((Num n))
      ((set (best-error (Num n)) (to-f64 n))))
(rule ((Add a b)) ((best-error (Add a b))))

;; finally, the mega rule for finding more accurate programs
(rule ((= expr (Add a b))
       (= (best-error a) va)
       (= (best-error b) vb)
       (= true-v (true-value (Add a b)))
       (= computed (+ va vb))
       (< (abs (- computed true-v))
          (best-error (Add a b))))
      ((set (best-error (Add a b)) computed)))



(push)

(let target
  (Add 
    (Add (Num (rational 1 100)) (Num (rational 1 100)))
    (Num (rational -2 100))))

(run 1)

;; set a default
(best-error target)
;; error is bad, constant folding hasn't fired enough
(query-extract (best-error target))

(run 1)

;; error is good, constant folding has fired enough
(query-extract (best-error target))


(pop)