;;; Copyright 2024 Peter McGoron
;;; Licensed under the Apache License, Version 2.0 (the "License");
;;; you may not use this file except in compliance with the License.
;;; You may obtain a copy of the License at
;;; 
;;;     http://www.apache.org/licenses/LICENSE-2.0
;;; 
;;; Unless required by applicable law or agreed to in writing, software
;;; distributed under the License is distributed on an "AS IS" BASIS,
;;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
;;; implied. See the License for the specific language governing
;;; permissions and limitations under the License.
;;; 
;;; Computer algebra system written using COND-VALUES.

(cond-expand
 (chicken (import r7rs)))

(load "cond-values.r7rs.scm")
(load "values-lib.r7rs.scm")

(import (mcgoron.com cond-values)
        (mcgoron.com values-lib)
        srfi-1)

;;; Record type that stores information about a known function. The CAS
;;; allows for unknown functions.
(define-record-type/destructor <function-data>
  (function-data procedure properties)
  function-data?
  destruct-function-data
  (procedure %get-procedure)
  (properties %get-properties))

;;; Database for known functions as an assocation list.
(define *functions*
  (list (cons '/ (function-data / '()))
        (cons 'exp (function-data exp '()))
        (cons 'expt (function-data expt '()))
        (cons '+ (function-data + '(abelian)))
        (cons '* (function-data * '(abelian)))))

;;; Destructures a function application.
(define function=> pair=>)

;;; Destructures a known function into
;;; (VALUES HEAD EXPRS PROCEDURE ARGNUM PROPERTIES)
;;; 
;;; where (HEAD EXPRS) was the application and the rest of the arguments
;;; are from the known function record.
(define (known-function=> expr)
  (after ((let (function=> expr) => (head arguments))
          (let (assq=> head *functions*) => (desc)))
    (cond-values
      (after ((let (destruct-function-data desc) => desc))
        (apply values head arguments desc))
      (else (error "improper function description" expr desc)))))

;;; Destructures an abelian function into
;;; (VALUES HEAD EXPRS PROC)
;;; where (HEAD . EXPRS) was the original application and PROC is the
;;; implementation of HEAD.
(define (abelian=> expr)
  (after ((let (known-function=> expr) => (head arguments proc props))
          (when (memq 'abelian props)))
    (values head arguments proc)))

;;; ABELIAN-FOLD implements constant folding on abelian operations (such
;;; as +, *, etc.).
;;; 
;;; The function scans ARGUMENTS for arguments, and:
;;; 
;;; * If an argument is an application of the same function, then the
;;;   arguments to that function are lifted up to the arguments of the
;;;   calling function.
;;; * If an argument is an application to a different function, call
;;;   the constant folder on it.
;;; * If the argument is an atom, do nothing to it.
;;; 
;;; Once all the arguments are constant folded, the function returns
;;; 
;;; * a number, which is the application of PROC to the folded arguments
;;;   when all arguments reduce to numbers, or
;;; * (CONS HEAD-TO-MATCH FOLDED), where FOLDED has at most one constant
;;;   argument in front, and the other folded arguments in arbitrary order.
(define (abelian-fold head-to-match arguments proc)
  (define (fold-expr expr constants non-constants)
    (let ((expr (constant-fold expr)))
      (cond-values
        (after ((let (function=> expr) => (head tail))
                (when (eq? head head-to-match)))
          (fold-values fold-expr (list constants non-constants) tail))
        (after ((when (number? expr)))
          (values (cons expr constants) non-constants))
        (else (values constants (cons expr non-constants))))))
  (let*-values (((constants non-constants)
                 (fold-values fold-expr '(() ()) arguments)))
    (if (null? constants)
        (cons head-to-match non-constants)
        (let ((constant (apply proc constants)))
          (if (null? non-constants)
              constant
              (cons head-to-match (cons constant non-constants)))))))

;;; Fold constant expressions in EXPR.
;;; 
;;; NOTE: This will apply algebraic operations like commutation and
;;; distribution. This can cause cancellations using floating point
;;; numbers. This function should only be used with exact numbers.
(define (constant-fold expr)
  (cond-values
    (apply-after (abelian=> expr) abelian-fold)
    (else expr)))

(constant-fold '(+ 5 (+ x 6)))
(constant-fold '(+ 5 (* 11 20 (+ 16 x))))