#| -*-Scheme-*-

$Id$

Copyright 2006 Massachusetts Institute of Technology

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or (at
your option) any later version.

This program is distributed in the hope that it will be useful, but
WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301,
USA.

|#


;;;; Simple catch-as-catch-can solver #42.  
;;;  By GJS, November 2003.  For use with SCMUTILS.

(declare (usual-integrations))

;;; Assumes standardized equations, as produced below.

#|
(define (solve-incremental equations variables) ...)

;;; Variables list gives variables to be eliminated in the preferred
;;; order of elimination.

;;; To access the parts of a solution returned by solve-incremental

(define (residual-equations solution) ...)
(define (residual-variables solution) ...)
(define (substitutions solution) ...)
(define (hopeless-variables solution) ...)


;;; To access substitutions produced by solve-incremental

(define (substitution-variable subst) ...)
(define (substitution-expression subst) ...)
(define (substitution-justifications subst) ...)


;;; To access residual equations of a solution

(define (equation-expression eqn) ...)
(define (equation-justifications eqn) ...)
|#

(define (residual-equations solution) (car solution))
(define (residual-variables solution) (cadr solution))
(define (substitutions solution) (caddr solution))
(define (hopeless-variables solution) (cadddr solution))

(define (solve-incremental equations variables)
  (let lp ((residual-eqs      equations)
	   (residual-vars     variables)
	   (substitutions    '())
	   (hopeless-vars    '()))
    (if (or (null? residual-vars) (null? residual-eqs))	;we are done.
	(list residual-eqs residual-vars substitutions hopeless-vars)
	(isolate-var (car residual-vars)
		     (sort residual-eqs
			   (lambda (eqn1 eqn2)
			     (< (length (equation-variables eqn1))
				(length (equation-variables eqn2)))))
		     (lambda (new-substitution equation-used)	
		       (lp (filter (lambda (eqn) ;flush tautologies
				     (let ((expr (equation-expression eqn)))
				       (not (and (number? expr) (= expr 0)))))
				   (map (lambda (equation)
					  (backsubstitute-equation new-substitution
								   equation))
					(delete equation-used residual-eqs)))
			   (cdr residual-vars)
			   (cons new-substitution
				 (map (lambda (substitution)
					(backsubstitute new-substitution
							substitution))
				      substitutions))
			   hopeless-vars))
		     (lambda ()		; fail
		       (lp residual-eqs
			   (cdr residual-vars)
			   substitutions
			   (cons (car residual-vars) hopeless-vars)))))))

(define (isolate-var var eqs succeed fail)
  ;; succeed = (lambda (new-substitution equation-used) ...)
  ;; fail    = (lambda () ...)
  (let lp ((eqs-to-scan eqs))
    (cond ((null? eqs-to-scan) (fail))
	  ((member var (equation-variables (car eqs-to-scan)))
	   (isolatable? var (car eqs-to-scan)
			(lambda (value)
			  (succeed (make-substitution var value
				      (equation-justifications (car eqs-to-scan)))
				   (car eqs-to-scan)))
			(lambda ()
			  (lp (cdr eqs-to-scan)))))
	  (else (fail)))))
	
(define (isolatable? var eqn succeed fail)
  (let lp ((expr (equation-expression eqn)))
    (cond ((equal? var expr) (succeed 0))
	  ((and (expt? expr)
		(number? (cadr (operands expr)))
		(> (cadr operands) 0))
	   (lp (car (operands expr))))
	  ((product? expr)
	   ;; If var^n in operands then could be zero.
	   (let lp ((factors (operands expr)))
	     (cond ((null? factors) (fail))
		   ((symbol? (car factors))
		    (if (eq? var (car factors))
			(succeed 0)
			(lp (cdr factors))))
		   ((and (expt? (car factors))
			 (number? (cadr (operands (car factors))))
			 (> (cadr (car factors)) 0)
			 (eq? var (caddr (operands (car factors)))))
		    (succeed 0))
		   (else (lp (cdr factors))))))
	  ((sum? expr)
	   ;; Split operands into ones with var^1 and without...
	   (let lp ((addends (operands expr)) (with '()) (without '()))
	     (cond ((null? addends)
		    (if (null? without)
			(succeed 0)
			(succeed
			 (symb:quo (symb:negate (symb:add:n without))
				   (symb:add:n with)))))
		   ((occurs? var (car addends))
		    (let ((addend (car addends)))
		      (cond ((equal? var addend)
			     (lp (cdr addends) (cons 1 with) without))
			    ((product? addend)
			     (let ((factors (operands addend)))
			       (if (member var factors)
				   (lp (cdr addends)
				       (cons (symb:mul:n (delete var factors))
					     with)
				       without)
				   (fail)))) ;occurs more painfully
			    (else (fail)))))
		   (else
		    (lp (cdr addends) with (cons (car addends) without))))))
	  (else (fail)))))
	   
(define (occurs? var expr)
  (or (equal? var expr)
      (and (pair? expr)
	   (or (occurs? var (car expr))
	       (occurs? var (cdr expr))))))

(define (make-substitution var value justs)
  (list (list '= var (new-simplify value)) justs))

(define (substitution-variable subst) (lhs (car subst)))
(define (substitution-expression subst) (rhs (car subst)))
(define (substitution-justifications subst) (cadr subst))

(define (lhs exp) (cadr exp))
(define (rhs exp) (caddr exp))


;;; expr should have value zero
;;; justs is a list of tms nodes

(define (make-equation expr justs)
  (if *tms-present*
      (let* ((specs (standardize-equation expr '() '() #f list))
	     (pexpr (car specs))
	     (vspecs (cadr specs))
	     (tms (tms:node-tms (car justs))))
	(if (and (number? pexpr) (not (= pexpr 0)))
	    (begin (tms:justify-node (tms:contradiction tms) 'equation justs)
		   (list pexpr justs vspecs))
	    (let ((node (tms:make-node tms (cp:make-equation-datum pexpr))))
	      (tms:justify-node node 'equation justs)
	      (list pexpr (list node) vspecs))))
      (let* ((specs (standardize-equation expr '() '() #f list))
	     (pexpr (car specs))
	     (vspecs (cadr specs)))
	(if (and (number? pexpr) (not (= pexpr 0)))
	    (begin (write-line `(contradiction ,justs))
		   (list pexpr justs vspecs))
	    (list pexpr justs vspecs)))))

;;; Set to #f for independent use of equation solver.
(define *tms-present* #t)

(define-record-type <cp:equation-datum>
    (cp:make-equation-datum residual)
    cp:equation-datum?
  (residual cp:equation-datum-residual))

(define (cp:equation? x)
  (and (tms:node? x)
       (cp:equation-datum? (tms:node-datum x))))


(define (equation-expression eqn) (car eqn))
(define (equation-justifications eqn) (cadr eqn))
(define (equation-variables eqn) (caddr eqn))

(define (backsubstitute new-substitution substitution)
  (if (occurs? (substitution-variable new-substitution)
	       (substitution-expression substitution))
      (make-substitution
       (substitution-variable substitution)
       (substitute (substitution-expression new-substitution)
		   (substitution-variable new-substitution)
		   (substitution-expression substitution))
       (list-union (substitution-justifications substitution)
		   (substitution-justifications new-substitution)))
      substitution))

(define (backsubstitute-equation substitution equation)
  (if (occurs? (substitution-variable substitution)
	       (equation-expression equation))
      (make-equation
       (substitute (substitution-expression substitution)
		   (substitution-variable substitution)
		   (equation-expression equation))
       (list-union (equation-justifications equation)
		   (substitution-justifications substitution)))
      equation))
  
(define (standardize-equation residual variables functions variable continue)
  ;; continue = (lambda (new-residual new-map functions) ...)
  (let ((redo #f))
    (define (walk-expression expression map functions continue)
      (cond ((pair? expression)
	     (let ((rator (operator expression))
		   (rands (operands expression)))
	       (cond ((and (= (length rands) 1) (eq? (car rands) variable))
		      (walk-expression rator
				       (list-adjoin rator map)
				       (list-adjoin rator functions)
				       continue))
		     ((or (D? expression) (D2? expression))
		      (continue expression
				(list-adjoin expression map)
				(list-adjoin expression functions)))
		     (else
		      (walk-expression rator map functions
			(lambda (rator-result rator-map rator-functions)
			  (walk-list rands rator-map rator-functions
				     (lambda (rands-result rands-map rands-functions)
				       (continue (cons rator-result
						       rands-result)
						 rands-map
						 rands-functions)))))))))
	    ((number? expression)
	     (continue (if (and (inexact? expression)
				(< (abs expression) *zero-threshold*))
			   (begin (set! redo #t) 0)
			   expression)
		       map
		       functions))
	    ((memq expression '(+ - / * D expt exp sin cos))
	     (continue expression map functions))
	    (else
	     (continue expression
		       (list-adjoin expression map)
		       functions))))
    (define (walk-list elist map functions continue)
      (if (pair? elist)
	  (walk-expression (car elist) map functions
	    (lambda (car-result car-map car-functions)
	      (walk-list (cdr elist) car-map car-functions 
			 (lambda (cdr-result cdr-map cdr-functions)
			   (continue (cons car-result cdr-result)
				     cdr-map
				     cdr-functions)))))
	  (continue elist map functions)))
    (let lp ((residual (new-simplify residual)))
      (walk-expression (if (quotient? residual)
			   (symb:dividend residual)
			   residual)
		       variables
		       functions
		       (lambda (expression map funs)
			 (if redo
			     (begin (set! redo #f)
				    (lp (new-simplify expression)))
			     (continue expression map funs)))))))

(define *zero-threshold* 1e-15)		;for small numbers

(define (D? x)
  (and (pair? x)
       (eq? (car x) 'D)))

(define (D2? x)
  (and (pair? x) 
       (equal? (car x) '(expt D 2))))

#|
;;; Shows signs of life.  

(set! *tms-present* #f)

(pp (solve-incremental
     (list (make-equation '(+ (* 3 x)     y  -7)  (list 'A))
	   (make-equation '(+ (* 3 x) (- y)  -5)  (list 'B)))
     '(x y)))
(() () (((= y 1) (B A)) ((= x 2) (B A))) ())


(pp (solve-incremental
     (list (make-equation '(+  x   y   z  1)  (list 'A))
	   (make-equation '(+  x   y      2)  (list 'B))
	   (make-equation '(+  x          1)  (list 'C)))
     '(x y z)))
(() () (((= z 1) (A B C)) ((= y -1) (B C)) ((= x -1) (C))) ())


(pp (solve-incremental
     (list (make-equation '(+  x          1)  (list 'C))
	   (make-equation '(+  x   y      2)  (list 'B))
	   (make-equation '(+  x   y   z  1)  (list 'A)))
     '(x y z)))
(() () (((= z 1) (A B C)) ((= y -1) (B C)) ((= x -1) (C))) ())

;;; The following signals a contradiction, as it should:

(pp (solve-incremental
     (list (make-equation '(+ (* 3 x)     y  -7)  (list 'A))
	   (make-equation '(+ (* 3 x)     y  -5)  (list 'B)))
     '(x y)))
(contradiction (B A))
(((2 (B A))) () (((= x (+ 7/3 (* -1/3 y))) (A))) (y))


;;; Some slightly nonlinear systems can be solved:

(pp (solve-incremental
     (list (make-equation '(-  3 (+ x y))  (list 'A))
	   (make-equation '(-  5 (- x y))  (list 'B))
	   (make-equation '(-  3 (+ (* (sqrt x) z) (square y)))  (list 'C)))
     '(x y z)))
(() () (((= z 1) (C B A)) ((= y -1) (B A)) ((= x 4) (B A))) ())

;;; Underdetermined systems can be reduced:

(pp (solve-incremental
     (list (make-equation '(+ (* (+ a b) (- a c)) c)  (list 'A))
	   (make-equation '(- 3 (+ a b))  (list 'B)))
     '(a b c)))
(() (c) (((= b (+ 3 (* -2/3 c))) (A B)) ((= a (* 2/3 c)) (A B))) ())


(pp (solve-incremental
     (list (make-equation '(+ (* (+ a b) (- a c)) c)  (list 'A))
	   (make-equation '(- 3 (- a c))  (list 'B)))
     '(a b c)))
(() (c) (((= b (+ -3 (* -4/3 c))) (A B)) ((= a (+ 3 c)) (B))) ())


;;; Even very hard ones are clarified.

(pp (solve-incremental
     (list (make-equation '(+ (* (+ a b) (- a c)) c)  (list 'A))
	   (make-equation '(- 3 (- a b))  (list 'B)))
     '(a b c)))
(()
 ()
 (((= c (/ (+ 9/2 (expt b 2) (* 9/2 b)) (+ 1 b))) (A B))
  ((= a (+ 3 b)) (B)))
 (b))
|#