; Depends on th_real.plf, th_int.plf, smt.plf, sat.plf

; Some axiom arguments are marked "; Omit", because they can be deduced from
; other arguments and should be replaced with a "_" when invoking the axiom.

;; ====================================== ;;
;; Arith Terms, Predicates, & Conversions ;;
;; ====================================== ;;

; Types for arithmetic variables
; Specifically a real
(declare real_var type)
; Specifically an integer
(declare int_var type)

; Functions to map them to terms
(declare term_real_var (! v real_var (term Real)))
(declare term_int_var (! v int_var (term Int)))

; A function to cast an integer term to real.
(declare term_int_to_real (! i (term Int) (term Real)))


; The recursive functions `reify_int_term` and `reify_real_term` work
; together to  reify or "make real" an integer term. That is, to convert it to
; a real term.  More precisely, they take an integer term and return a real
; term in which any integer variables are immediately converted to real terms,
; and all non-leaves in the term are real-sorted.
;
; They explicitly do not work on integer division, because such a conversion
; would not be correct when integer division is involved.

; This function recursively converts an integer term to a real term.
(program reify_int_term ((t (term Int))) (term Real)
  (match t
    ((term_int_var v) (term_int_to_real (term_int_var v)))
    ((a_int i) (a_real (mpz_to_mpq i)))
    ((+_Int x y) (+_Real (reify_int_term x) (reify_int_term y)))
    ((-_Int x y) (-_Real (reify_int_term x) (reify_int_term y)))
    ((u-_Int x) (u-_Real (reify_int_term x)))
    ((*_Int x y) (*_Real (reify_int_term x) (reify_int_term y)))
    ; Reifying integer division is an error, since it changes the value!
    ((/_Int x y) (fail (term Real)))
  ))

; This function recursively converts a real term to a real term.
; It will never change the top-level node in the term (since that node is
; real), but it may change subterms...
(program reify_real_term ((t (term Real))) (term Real)
  (match t
    ((term_real_var v) (term_real_var v))
    ((a_real v) (a_real v))
    ; We've found an integer term -- reify it!
    ((term_int_to_real t') (reify_int_term t'))
    ((+_Real x y) (+_Real (reify_real_term x) (reify_real_term y)))
    ((-_Real x y) (-_Real (reify_real_term x) (reify_real_term y)))
    ((u-_Real x) (u-_Real (reify_real_term x)))
    ((*_Real x y) (*_Real (reify_real_term x) (reify_real_term y)))
    ((/_Real x y) (/_Real (reify_real_term x) (reify_real_term y)))
  ))

; Predicates of the form (term Integer) (comparison) (term Real)
(define arithpred_IntReal (! x (term Int)
                          (! y (term Real)
                          formula)))
(declare >_IntReal arithpred_IntReal)
(declare >=_IntReal arithpred_IntReal)
(declare <_IntReal  arithpred_IntReal)
(declare <=_IntReal arithpred_IntReal)

; From an arith predicate, compute the equivalent real predicate
; All arith predicates are (possibly negated) >='s with a real on the right.
; Technically it's a real literal on the right, but we don't assume that here.
(program reify_arith_pred ((p formula)) formula
  (match p
         ((not p') (not (reify_arith_pred p')))
         ((>=_Real x y) (>=_Real (reify_real_term x) (reify_real_term y)))
         ((>=_Int x y) (>=_Real (reify_int_term x) (reify_int_term y)))
         ((>=_IntReal x y) (>=_Real (reify_int_term x) (reify_real_term y)))
         (default (fail formula))
         ))

; From an arith predicate, prove the equivalent real predicate
(declare pf_reified_arith_pred
  (! p formula
  (! p' formula
    (! pf (th_holds p)
      (! reify_sc (^ (reify_arith_pred p) p')
         (th_holds p'))))))

;; ========================== ;;
;; Int Bound Tightening Rules ;;
;; ========================== ;;

; Returns whether `ceil` is the ceiling of `q`.
(program is_ceil ((q mpq) (ceil mpz)) bool
  (let diff (mp_add (mpz_to_mpq ceil) (mp_neg q))
    (mp_ifneg diff
              ff
              (mp_ifneg (mp_add diff (~ 1/1))
                        tt
                        ff))))

; Returns whether `n` is the greatest integer less than `q`.
(program is_greatest_integer_below ((n mpz) (q mpq)) bool
  (is_ceil q (mp_add n 1)))


; Negates terms of the form:
; (a) k     OR
; (b) x     OR
; (c) k * x
; where k is a constant and x is a variable.
; Otherwise fails.
; This aligns closely with the LFSCArithProof::printLinearMonomialNormalizer
; function.
(program negate_linear_monomial_int_term ((t (term Int))) (term Int)
  (match t
    ((term_int_var v) (*_Int (a_int (~ 1)) (term_int_var v)))
    ((a_int k) (a_int (mp_neg k)))
    ((*_Int x y)
     (match x
            ((a_int k)
             (match y
                    ((term_int_var v) (*_Int (a_int (mp_neg k)) y))
                    (default (fail (term Int)))))
            (default (fail (term Int)))))
    (default (fail (term Int)))
  ))

; This function negates linear interger terms---sums of terms of the form
; recognized by `negate_linear_monomial_int_term`.
(program negate_linear_int_term ((t (term Int))) (term Int)
  (match t
    ((term_int_var v) (negate_linear_monomial_int_term t))
    ((a_int i) (negate_linear_monomial_int_term t))
    ((+_Int x y) (+_Int (negate_linear_int_term x) (negate_linear_int_term y)))
    ((*_Int x y) (negate_linear_monomial_int_term t))
    (default (fail (term Int)))
  ))

; Statement that z is the greatest integer less than z'.
(declare holds_neg_of_greatest_integer_below_int
  (! z mpz
  (! z' mpz
    type)))

; For proving statements of the above form.
(declare check_neg_of_greatest_integer_below_int
  (! z mpz
  (! z' mpz
    (! sc_check (^ (is_greatest_integer_below (mp_neg z) (mpz_to_mpq z')) tt)
       (holds_neg_of_greatest_integer_below_int z z')))))

; Axiom for tightening [Int] < i into -[Int] >= -(i - 1).
; Note that [Int] < i is actually not([Int] >= i)
(declare tighten_not_>=_IntInt
  (! t       (term Int)  ; Omit
  (! neg_t   (term Int)  ; Omit
  (! old_bound     mpz ; Omit
  (! neg_int_bound mpz ; Omit
    (! pf_step (holds_neg_of_greatest_integer_below_int neg_int_bound old_bound)
    ; Note that even when the RHS is an integer, we convert it to real and use >_IntReal
    (! pf_real_bound (th_holds (not (>=_IntReal t (term_int_to_real (a_int old_bound)))))
      (! sc_neg (^ (negate_linear_int_term t) neg_t)
        (th_holds (>=_IntReal neg_t (term_int_to_real (a_int neg_int_bound))))))))))))

;; ======================================== ;;
;; Linear Combinations and Affine functions ;;
;; ======================================== ;;

; Unifying type for both kinds of arithmetic variables
(declare arith_var type)
(declare av_from_int (! v int_var arith_var))
(declare av_from_real (! v real_var arith_var))

; Total order type -- return value for the comparison of two things
(declare ord type)
(declare ord_lt ord)
(declare ord_eq ord)
(declare ord_gt ord)

; Compare two arith vars. Integers come before reals, and otherwise we use the
; LFSC ordering
(program arith_var_cmp ((v1 arith_var) (v2 arith_var)) ord
  (match v1
    ((av_from_int  i1)
      (match v2
        ((av_from_int  i2)
          (ifequal i1 i2
            ord_eq
            (compare i1 i2 ord_lt ord_gt)))
        ((av_from_real r2) ord_lt)))
    ((av_from_real r1)
      (match v2
        ((av_from_int  i2) ord_gt)
        ((av_from_real r2)
          (ifequal r1 r2
            ord_eq
            (compare r1 r2 ord_lt ord_gt)))))))

; Type for linear combinations of variables
; NB: Functions below will assume that the list is always sorted by variable!
(declare lc type)
(declare lc_null lc)
(declare lc_cons (! c mpq (! v arith_var (! rest lc lc))))

; Sum of linear combinations.
(program lc_add ((l1 lc) (l2 lc)) lc
  (match l1
    (lc_null l2)
    ((lc_cons c1 v1 l1')
      (match l2
        (lc_null l1)
        ((lc_cons c2 v2 l2')
          (match (arith_var_cmp v1 v2)
            (ord_lt (lc_cons c1 v1 (lc_add l1' l2)))
            (ord_eq
              (let c (mp_add c1 c2)
                (mp_ifzero c
                  (lc_add l1' l2')
                  (lc_cons c v1 (lc_add l1' l2')))))
            (ord_gt (lc_cons c2 v2 (lc_add l1 l2')))))))))

; Scaling a linear combination
(program lc_mul_c ((l lc) (c mpq)) lc
  (match l
    (lc_null l)
    ((lc_cons c' v' l') (lc_cons (mp_mul c c') v' (lc_mul_c l' c)))))

; Negating a linear combination
(program lc_neg ((l lc)) lc
         (lc_mul_c l (~ 1/1)))

; An affine function of variables (a linear combination + a constant)
(declare aff type)
(declare aff_cons (! c mpq (! l lc aff)))

; Sum of affine functions
(program aff_add ((p1 aff) (p2 aff)) aff
  (match p1
    ((aff_cons c1 l1)
      (match p2
        ((aff_cons c2 l2) (aff_cons (mp_add c1 c2) (lc_add l1 l2)))))))

; Scaling an affine function
(program aff_mul_c ((p aff) (c mpq)) aff
  (match p
    ((aff_cons c' l') (aff_cons (mp_mul c' c) (lc_mul_c l' c)))))

; Negating an affine function
(program aff_neg ((p aff)) aff
  (aff_mul_c p (~ 1/1)))

; Subtracting affine functions
(program aff_sub ((p1 aff) (p2 aff)) aff
  (aff_add p1 (aff_neg p2)))

;; ================================= ;;
;; Proving (Real) terms to be affine ;;
;; ================================= ;;

; truth type for some real term being affine
; * `t` the real term
; * `a` the equivalent affine function
(declare is_aff (! t (term Real) (! a aff type)))

; Constants are affine
(declare is_aff_const
  (! x mpq
    (is_aff (a_real x) (aff_cons x lc_null))))

; Real variables are affine
(declare is_aff_var_real
  (! v real_var
    (is_aff (term_real_var v)
            (aff_cons 0/1 (lc_cons 1/1 (av_from_real v) lc_null)))))

; Int variables are affine
(declare is_aff_var_int
  (! v int_var
    (is_aff (term_int_to_real (term_int_var v))
            (aff_cons 0/1 (lc_cons 1/1 (av_from_int v) lc_null)))))

; affine functions are closed under addition
(declare is_aff_+
  (! x (term Real)      ; Omit
  (! aff_x aff          ; Omit
  (! y (term Real)      ; Omit
  (! aff_y aff          ; Omit
  (! aff_z aff          ; Omit
    (! is_affx (is_aff x aff_x)
    (! is_affy (is_aff y aff_y)
      (! a (^ (aff_add aff_x aff_y) aff_z)
        (is_aff (+_Real x y) aff_z))))))))))

; affine functions are closed under subtraction
(declare is_aff_-
  (! x (term Real)      ; Omit
  (! aff_x aff          ; Omit
  (! y (term Real)      ; Omit
  (! aff_y aff          ; Omit
  (! aff_z aff          ; Omit
    (! is_affx (is_aff x aff_x)
    (! is_affy (is_aff y aff_y)
      (! a (^ (aff_sub aff_x aff_y) aff_z)
        (is_aff (-_Real x y) aff_z))))))))))

; affine functions are closed under left-multiplication by scalars
(declare is_aff_mul_c_L
  (! y (term Real)      ; Omit
  (! aff_y aff          ; Omit
  (! aff_z aff          ; Omit
    (! x mpq
    (! is_affy (is_aff y aff_y)
      (! a (^ (aff_mul_c aff_y x) aff_z)
        (is_aff (*_Real (a_real x) y) aff_z))))))))

; affine functions are closed under right-multiplication by scalars
(declare is_aff_mul_c_R
  (! y (term Real)      ; Omit
  (! aff_y aff          ; Omit
  (! aff_z aff          ; Omit
    (! x mpq
    (! is_affy (is_aff y aff_y)
      (! a (^ (aff_mul_c aff_y x) aff_z)
        (is_aff (*_Real y (a_real x)) aff_z))))))))

;; ========================== ;;
;; Bounds on Affine Functions ;;
;; ========================== ;;

; Bounds that an affine function might satisfy
(declare bound type)
(declare bound_pos bound)           ; > 0
(declare bound_non_neg bound)       ; >= 0

; formulas over affine functions
; the statement that `a` satisfies `b` for all inputs
(declare bounded_aff (! a aff (! b bound formula)))

; Sum of two bounds (the bound satisfied by the sum of two functions satifying
; the input bounds)
(program bound_add ((b bound) (b2 bound)) bound
  (match b
    (bound_pos bound_pos)
    (bound_non_neg b2)))

; The implication of `a1` satisfying `b` and `a2` satisfying `b2`, obtained by
; summing the inequalities.
(program bounded_aff_add ((a1 aff) (b bound) (a2 aff) (b2 bound)) formula
  (bounded_aff (aff_add a1 a2) (bound_add b b2)))


; The implication of scaling the inequality of `a1` satisfying `b`.
(program bounded_aff_mul_c ((a1 aff) (b bound) (c mpq)) formula
  (match (mpq_ispos c)
    (tt (bounded_aff (aff_mul_c a1 c) b))
    (ff (fail formula))))

; Does an affine function actuall satisfy a bound, for some input?
(program bound_respected ((b bound) (a aff)) bool
  (match a
    ((aff_cons c combo)
      (match combo
        (lc_null
          (match b
            (bound_pos (mpq_ispos c))
            (bound_non_neg (mp_ifneg c ff tt))))
        (default tt)))))

;; =================================== ;;
;; Axioms for bounded affine functions ;;
;; =================================== ;;

; Always true (used as a initial value when summing many bounds together)
(declare bounded_aff_ax_0_>=_0
  (th_holds (bounded_aff (aff_cons 0/1 lc_null) bound_non_neg)))

; Contradiction axiom: an affine function that does not respect its bounds
(declare bounded_aff_contra
  (! a aff      ; Omit
  (! b bound    ; Omit
    (! pf (th_holds (bounded_aff a b))
      (! sc (^ (bound_respected b a) ff)
         (th_holds false))))))

; Rule for summing two affine bounds to get a third
(declare bounded_aff_add
  (! a1 aff             ; Omit
  (! a2 aff             ; Omit
  (! b bound            ; Omit
  (! b2 bound           ; Omit
  (! ba_sum formula     ; Omit
    (! pf_a1 (th_holds (bounded_aff a1 b))
    (! pf_a2 (th_holds (bounded_aff a2 b2))
       (! sc (^ (bounded_aff_add a1 b a2 b2) ba_sum)
         (th_holds ba_sum))))))))))

; Rule for scaling an affine bound
(declare bounded_aff_mul_c
  (! a aff          ; Omit
  (! b bound        ; Omit
  (! ba formula     ; Omit
    (! c mpq
    (! pf_a (th_holds (bounded_aff a b))
       (! sc (^ (bounded_aff_mul_c a b c) ba)
         (th_holds ba))))))))


; [y >= x] implies that the aff. function y - x is >= 0
(declare aff_>=_from_term
  (! y (term Real)  ; Omit
  (! x (term Real)  ; Omit
  (! p aff          ; Omit
    (! pf_affine (is_aff (-_Real y x) p)
    (! pf_term_bound (th_holds (>=_Real y x))
      (th_holds (bounded_aff p bound_non_neg))))))))

; not [y >= x] implies that the aff. function -(y - x) is > 0
(declare aff_>_from_term
  (! y (term Real)  ; Omit
  (! x (term Real)  ; Omit
  (! p aff          ; Omit
  (! p_n aff        ; Omit
    (! pf_affine (is_aff (-_Real y x) p)
    (! pf_term_bound (th_holds (not (>=_Real y x)))
      (! sc_neg (^ (aff_neg p) p_n)
        (th_holds (bounded_aff p_n bound_pos))))))))))