;;   forces.scm - compute forces acting on atoms, based on MM2 energy terms
;;   Copyright (C) 1996,1997 Will Ware
;;   
;;   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., 675 Mass Ave, Cambridge, MA 02139, USA.
;;
;;   I can be reached via email at .

;; ******************************************************************
;; Table of difference vectors and distances

;; At the beginning of compute-forces, we compute the difference vectors and
;; distances for all pairs of atoms. These would otherwise be repeatedly
;; calculated by various energy/force terms, so it's a big win to precalculate
;; all that stuff in a big batch and make it available for other functions.
;; For N atoms, we will need to compute N*(N-1)/2 vector/distance pairs.
;; But to keep my head from collapsing, I'll keep it simple and use N^2
;; things, the bookkeeping is too ugly otherwise.

(define diff-dist-table '())

(define (lookup-diff-dist m n)
  (vector-ref (vector-ref diff-dist-table n) m))

(define (build-diff-dist-table)
  (set! diff-dist-table (make-vector (length atom-list)))
  (dotimes
   (i (length atom-list))
   (vector-set! diff-dist-table i
		(make-vector (length atom-list)))))

(define (fill-diff-dist-table)
  (do ((L1 atom-list (cdr L1))
       (i 0 (+ i 1)))
      ((null? L1))
    (do ((L2 atom-list (cdr L2))
	 (j 0 (+ j 1)))
	((null? L2))
      (vector-set! (vector-ref diff-dist-table i)
		   j
		   (create-diff-dist (car L2) (car L1))))))

;; ******************************************************************
;; Coefficients for MM2 energy/force terms

(define (lookup-helper compare-func result-func default-result the-list)
  (cond ((null? the-list) default-result)
	((funcall compare-func (car the-list))
	 (funcall result-func (car the-list)))
	(else (lookup-helper compare-func result-func default-result
			     (cdr the-list)))))

;; A length-coefficient entry is two MM2 atom types, followed by a
;; k_s value, followed by an r0 value.

(define length-coefficients
  '((1 5 460 1.113)
    (1 1 440 1.523)
    (2 2 960 1.337)
    (4 4 1560 1.212)
    (1 6 536 1.402)
    (1 8 510 1.438)
    (3 7 1080 1.208)
    (1 11 510 1.392)
    (1 12 323 1.795)
    (1 13 230 1.949)
    (1 14 220 2.149)
    (8 20 610 0.6)
    (8 8 560 1.381)
    (6 20 460 0.6)
    (6 21 460 0.942)
    (6 6 781 1.470)
    (1 19 297 1.880)
    (1 25 291 1.856)
    (1 15 321.3 1.815)
    (19 19 185 2.332)
    (22 22 440 1.501)))

(define (lookup-length-coeffs m n)
  (lookup-helper
   (make-lambda (x) (or (and (= m (car x))
			(= n (cadr x)))
		   (and (= n (car x))
			(= m (cadr x)))))
   (make-lambda (x) (cddr x))
   '(400 1.3)
   length-coefficients))

;; An angle-coefficient entry is three MM2 atoms types, followed by
;; a k_th value, followed by a th0 value.

(define angle-coefficients
  '((1 1 1 0.45 1.911)
    (1 1 5 0.36 1.909)
    (5 1 5 0.32 1.909)
    (1 1 11 0.65 1.911)
    (11 1 11 1.07 1.869)
    (1 2 1 0.45 2.046)
    (2 1 2 0.45 1.911)
    (1 2 2 0.55 2.119)
    (2 2 5 0.36 2.094)
    (2 2 2 0.43 2.094)
    (1 4 4 0.2 3.142)
    (1 3 7 0.46 2.138)
    (1 6 1 0.77 1.864)
    (1 8 1 0.63 1.88)
    (1 1 6 0.57 1.911)
    (1 6 20 0.35 1.835)
    (1 8 20 0.5 1.906)
    (20 6 20 0.24 2.286)
    (19 19 19 0.35 1.943)
    (19 1 19 0.4 2.016)
    (1 19 1 0.48 1.934)
    (12 1 12 1.08 1.95)
    (1 1 15 0.55 1.902)
    (1 15 1 0.72 1.902)
    ;; I think I made up these last few entries, they should be
    ;; considered suspect. Really, I should pick up the more complete
    ;; tables from Jay Ponder's ftp site.
    (4 4 5 0.4 3.142)
    (7 2 1 0.4 2.0944)
    (7 2 2 0.4 2.0944)
    (5 8 5 1.1 1.85)
    (1 2 6 0.4 2.0944)
    (1 2 5 0.4 2.0944)
    (5 2 6 0.4 2.0944)))

(define (lookup-angle-coeffs m n p)
  (lookup-helper
   (make-lambda (x) (or (and (= m (car x))
			(= n (cadr x))
			(= p (caddr x)))
		   (and (= p (car x))
			(= n (cadr x))
			(= m (caddr x)))))
   (make-lambda (x) (cdddr x))
   '(0.4 1.9)
   angle-coefficients))

;; An torsion-coefficient entry is four MM2 atoms types, followed by
;; three values for v1, v2, and v3 respectively.

(define torsion-coefficients
  '((1 1 1 1 1.39 1.88 0.65)
    (1 1 1 5 0.0 0.0 1.85)
    (5 1 1 5 0.0 0.0 1.65)
    (1 1 1 11 0.0 -0.6 6.46)
    (1 2 2 1 -0.69 69.47 0.0)
    (2 2 2 2 -6.46 55.58 0.0)
    (5 2 2 5 0.0 104.21 0.0)))

(define (lookup-torsion-coeffs m n p q)
  (lookup-helper
   (make-lambda (x) (or (and (= m (car x))
			(= n (cadr x))
			(= p (caddr x))
			(= q (cadddr x)))
		   (and (= q (car x))
			(= p (cadr x))
			(= n (caddr x))
			(= m (cadddr x)))))
   (make-lambda (x) (cddddr x))
   '(0.0 0.0 0.0)
   torsion-coefficients))

;; ******************************************************************
;; Geometric fun with vectors

(define (vplus v1 v2)
  (vector (+ (vector-ref v1 0) (vector-ref v2 0))
	  (+ (vector-ref v1 1) (vector-ref v2 1))
	  (+ (vector-ref v1 2) (vector-ref v2 2))))

(define (vdiff v1 v2)
  (vector (- (vector-ref v1 0) (vector-ref v2 0))
	  (- (vector-ref v1 1) (vector-ref v2 1))
	  (- (vector-ref v1 2) (vector-ref v2 2))))

(define (dot-product v1 v2)
  (+ (* (vector-ref v1 0) (vector-ref v2 0))
     (* (vector-ref v1 1) (vector-ref v2 1))
     (* (vector-ref v1 2) (vector-ref v2 2))))

(define (cross-product v1 v2)
  (vector (- (* (vector-ref v1 1) (vector-ref v2 2))
	     (* (vector-ref v1 2) (vector-ref v2 1)))
	  (- (* (vector-ref v1 2) (vector-ref v2 0))
	     (* (vector-ref v1 0) (vector-ref v2 2)))
	  (- (* (vector-ref v1 0) (vector-ref v2 1))
	     (* (vector-ref v1 1) (vector-ref v2 0)))))

(define (safe-sqrt x)
  (real-part (sqrt x)))

(define (safe-acos z)
  (if (> z 1.0) (set! z 1.0))
  (if (< z -1.0) (set! z -1.0))
  (acos z))

(define (vlen v)
  (safe-sqrt (dot-product v v)))

(define (vscale v x)
  (vector (* (vector-ref v 0) x)
	  (* (vector-ref v 1) x)
	  (* (vector-ref v 2) x)))

;; find the component of v1 perpendicular to v2
(define (perpendicular-component v1 v2)
  (vdiff v1 (vscale v2 (/ (dot-product v1 v2) (dot-product v2 v2)))))

(define (v-angle dd1 dd2)
  (safe-acos
   (/ (dot-product (dd-diff dd1) (dd-diff dd2))
      (* (dd-dist dd1) (dd-dist dd2)))))

;; i-torsion uses an obsolete version of v-angle, don't use torsion forces
;; for a while til I fix it

(define (i-torsion v1 v2 v3 v4)
  (let* ((d12 (lookup-diff-dist v1 v2))
	 (d23 (lookup-diff-dist v2 v3))
	 (d43 (lookup-diff-dist v4 v3))
	 (p (dd-diff d12))
	 (q (dd-diff d23))
	 (r (dd-diff d43)))
    (v-angle (perpendicular-component p q)
	     (perpendicular-component r q))))

;; ******************************************************************
;; Forces computed by energy terms

;; Functions with names like 'build-???-term' run during setup-terms, and
;; they build zero-argument functions which are stored in the term list.
;; setup-terms runs fairly infrequently (only when there are changes in
;; the structure: add-atom, delete-atom, add-bond, delete-bond, load-structure)
;; so the building functions don't need a lot of optimization.

;; Functions with names like '???-force' run during compute-forces, and
;; they compute the force contributions for each kind of MM2 energy term.
;; These forces are applied to all atoms that participate in computing the
;; energy term. compute-forces runs frequently, so it's worth a lot of
;; effort to optimize the '???-force' functions.

(define (build-length-term m n)
  (entering "build-length-term")
  (let* ((ma (list-ref atom-list m))
	 (na (list-ref atom-list n))
	 (coeffs
	  (lookup-length-coeffs
	   (species-mm2index (atm-species ma))
	   (species-mm2index (atm-species na))))
	 (ks (* 0.01 (car coeffs)))
	 (r0 (cadr coeffs)))
    (make-lambda ()
      (length-force ks r0 ma na m n))))

(define (length-force ks r0 ma na m n)
  (entering "length-force")
  (let* ((rdd (lookup-diff-dist m n))
	 (rd (dd-diff rdd))
	 (r (dd-dist rdd))
	 (du-dr (* ks (- r r0)))
	 (mult (/ (- du-dr) r)))
    (atm-add-force ma (vscale rd mult))
    (atm-add-force na (vscale rd (- mult)))))

(define (old-length-force ks r0 ma na)
  (entering "length-force")
  (let* ((rd (vdiff (atm-position ma) (atm-position na)))
	 (r (safe-sqrt (dot-product rd rd)))
	 (du-dr (* ks (- r r0)))
	 (mult (/ (- du-dr) r)))
    (atm-add-force ma (vscale rd mult))
    (atm-add-force na (vscale rd (- mult)))))

;; OK, for angles it gets a bit trickier. Let the atom positions be vectors
;; r1, r2, r3, with r2 the vertex. Define the following:
;; L1 = -dotproduct(r1-r2,r3-r2) * (r1-r2) + dotproduct(r1-r2,r1-r2) * (r3-r2)
;; L3 = -dotproduct(r1-r2,r3-r2) * (r3-r2) + dotproduct(r3-r2,r3-r2) * (r1-r2)
;; m1 = L1/|L1|, m3 = L3/|L3|
;; Potential energy is given by U = f(theta), force on atom 1 is
;; f1 = -m1 * dU/dm1 = (-m1/|r1-r2|) * dU/dtheta
;; Likewise f3 = (-m3/|r3-r2|) * dU/dtheta
;; Conservation: f2 = -f1-f3

(define (build-angle-term m n p)
  (entering "build-angle-term")
  (let* ((ma (list-ref atom-list m))
	 (na (list-ref atom-list n))
	 (pa (list-ref atom-list p))
	 (coeffs
	  (lookup-angle-coeffs
	   (species-mm2index (atm-species ma))
	   (species-mm2index (atm-species na))
	   (species-mm2index (atm-species pa))))
	 (kth (car coeffs))
	 (th0 (cadr coeffs)))
    (make-lambda ()
      (angle-force kth th0 ma na pa m n p))))

(define (angle-force kth th0 ma na pa m n p)
  (entering "angle-force")
  (let* ((dd12 (lookup-diff-dist m n))
	 (dd32 (lookup-diff-dist p n))
	 (th (v-angle dd32 dd12))
	 (tdif (- th th0))
	 (du-dth (* kth (* tdif (+ 1 (* 1.508 tdif tdif)))))
	 (r1r2 (dd-diff dd12))
	 (r3r2 (dd-diff dd32))
	 (L1 (vdiff (vscale r3r2 (dot-product r1r2 r1r2))
		    (vscale r1r2 (dot-product r1r2 r3r2))))
	 (f1 (vscale L1 (/ du-dth (* (vlen L1) (vlen r1r2)))))
	 (L3 (vdiff (vscale r1r2 (dot-product r3r2 r3r2))
		    (vscale r3r2 (dot-product r1r2 r3r2))))
	 (f3 (vscale L3 (/ du-dth (* (vlen L3) (vlen r3r2)))))
	 (f2 (vscale (vplus f1 f3) -1.0)))
    (atm-add-force ma f1)
    (atm-add-force na f2)
    (atm-add-force pa f3)))

;; To think about torsions, think of the projection of the whole thing into
;; the plane perpendicular to the torqued bond.

(define torsion-whine false)
(define use-torsion-forces false)

(define (build-torsion-term m n p q)
  (entering "build-torsion-term")
  (let* ((ma (list-ref atom-list m))
	 (na (list-ref atom-list n))
	 (pa (list-ref atom-list p))
	 (qa (list-ref atom-list q))
	 (coeffs
	  (lookup-torsion-coeffs
	   (species-mm2index (atm-species ma))
	   (species-mm2index (atm-species na))
	   (species-mm2index (atm-species pa))
	   (species-mm2index (atm-species qa))))
	 (v1 (car coeffs))
	 (v2 (cadr coeffs))
	 (v3 (caddr coeffs)))
    (make-lambda ()
      (torsion-force v1 v2 v3 ma na pa qa m n p q))))

(define (torsion-force v1 v2 v3 ma na pa qa m n p q)
  (entering "torsion-force")
  (if use-torsion-forces
      (set! torsion-whine true)))

;; Actually, torsions appear to contribute a negligible amount of force
;; compared even to van der Waals, but they require this atrociously complex
;; calculation. I'm thinking about just not bothering to compute them at all.
;; I think the math here is correct, just woefully inefficient.

;; torsion is broken for the time being
(define (do-not-use-torsion-force v1 v2 v3 ma na pa qa m n p q)
  (let* ((ddmn (lookup-diff-dist m n))
	 (ddpn (lookup-diff-dist p n))
	 (ddqp (lookup-diff-dist q p))
	 (pv (dd-diff ddmn))
	 (qv (dd-diff ddpn))
	 (rv (dd-diff ddqp))
	 (pp (dd-dist ddmn))
	 (qq (dd-dist ddpn))
	 (rr (dd-dist ddqp))
	 (pq (dot-product pv qv))
	 (qr (dot-product qv rv))
	 (alpha (safe-sqrt (/ (* pq pq) qq)))
	 (beta (safe-sqrt (* qq qq)))
	 (gamma (safe-sqrt (/ (* qr qr) qq)))
	 (vm1 (cross-product qv pv))
	 (vq1 (cross-product rv qv))
	 (vm2 (vscale vm1 (/ 1.0 (* (vlen vm1)))))
	 (vq2 (vscale vq1 (/ 1.0 (* (vlen vq1)))))
	 (w (safe-acos (dot-product vm2 vq2)))
	 (du-dw (* -0.0005 (+ (* v1 (sin w))
			      (* -2 v2 (sin (* 2 w)))
			      (* 3 v3 (sin (* 3 w))))))
	 (fm (vscale vm2 (/ du-dw
			    (safe-sqrt (- pp (/ (* pq pq) qq))))))
	 (fq (vscale vq2 (/ du-dw
			    (safe-sqrt (- rr (/ (* qr qr) qq)))))))
    (atm-add-force ma fm)
    (atm-add-force qa fq)
    (atm-add-force pa
	(vdiff (vscale fm (/ alpha beta))
	       (vscale fq (/ (+ gamma beta) beta))))
    (atm-add-force na
	(vdiff (vscale fq (/ gamma beta))
	       (vscale fm (/ (+ alpha beta) beta))))))

;; vdw is similar to length force
;; du/dr = 6*evdw*[(r/rvdw)^-7 - (r/rvdw)^-13]
;; Don't forget the factor of 0.001

(define use-vdw-forces true)

(define (build-vdw-term m n)
  (entering "build-vdw-term")
  (let* ((ma (list-ref atom-list m))  ;; members of atom-list
	 (na (list-ref atom-list n))
	 (evdw (* 0.5 (+ (species-evdw (atm-species ma))
			 (species-evdw (atm-species na)))))
	 (rvdw (+ (element-rvdw (atm-element ma))
		  (element-rvdw (atm-element na)))))
    (make-lambda ()
      (vdw-force evdw rvdw ma na m n))))

(define (vdw-force evdw rvdw ma na m n)
  (entering "vdw-force")
  (if use-vdw-forces
      (let* ((rdd (lookup-diff-dist m n))
	     (rd (dd-diff rdd))
	     (r (dd-dist rdd))
	     (rsq_recip (/ (* rvdw rvdw) (* r r))) ;; (r/rvdw)^-2
	     (r_recip (safe-sqrt rsq_recip))
	     (r6 (* rsq_recip rsq_recip rsq_recip))    ;; (r/rvdw)^-6
	     (du-dr (* 0.006 evdw r_recip r6 (- 1 (* 2 r6))))
	     (mult (/ (- du-dr) r)))
	(atm-add-force ma (vscale rd mult))
	(atm-add-force na (vscale rd (- mult))))))

;; We only need to run setup-terms when the structure changes: add-atom,
;; delete-atom, add-bond, delete-bond, or load-structure. This flag keeps
;; track of all that.
(define need-to-resetup-terms true)

(define (compute-forces)
  (entering "compute-forces")
  (if need-to-resetup-terms
      (let ()
	(setup-terms)
	(set! need-to-resetup-terms false)))
  ;; precalculate difference vectors and distances for all pairs of atoms
  (dbgprintf "1~%")
  (fill-diff-dist-table)
  ;; set all forces to zero, one force vector for each atom
  (dbgprintf "2~%")
  (dolist (a atom-list) (atm-zero-force a))
  ;; if the user defined external forces, put them in now
  (dbgprintf "3~%")
  (external-forces)
  ;; compute force contributions for each energy term
  (dbgprintf "4~%")
  (dolist (x term-list) (funcall x)))

;; ******************************************************************
;; Building up a term list from an atom list and bond list

(define (other-end bond atm)
  (cond ((= atm (bond-first bond)) (bond-second bond))
	((= atm (bond-second bond)) (bond-first bond))
	(else false)))

(define whine-about-bond-count false)

(define (extend-chain v lst)
  (let ((result '()))
    (dolist
     (L lst)
     (let* ((partial '())
	    (x (vector-ref v (car L))))
       (if (not (null? x))
	   (dolist
	    (z x)
	    (let ((flag true))
	      (dolist (z1 L)
		      (if (= z z1) (set! flag false)))
	      (if flag
		  (set! partial (cons (cons z L) partial))))))
       (set! result (append result partial))))
    result))

(define (start-chain)
  (do ((r '())
       (n (- (length atom-list) 1) (- n 1)))
      ((< n 0) r)
    (set! r (cons (list n) r))))

(define (unique-chains ch)
  (let ((n (- (length (car ch)) 1)))
    (remove-from-list-if-test
     ch
     (make-lambda (L)
		  (< (car L) (list-ref L n))))))

(define (complex-member m lst)
  (cond ((null? lst) false)
	((equal? m (car lst)) true)
	(else (complex-member m (cdr lst)))))

(define (setup-terms)
  (entering "setup-terms")
  (build-diff-dist-table)

  ;; for each atom, count its bonds, verify against valence, and reassign
  ;; its mm2 table index according to hybridization
  (let ((n 0) (bc false))
    (dolist (a atom-list)
	    (set! bc (create-bond-count n))
	    (if (not (= (element-total-bonds (atm-element a))
			(bond-count-total-bonds bc)))
		(set! whine-about-bond-count true))
	    (hybridize-atom a bc)
	    (set! n (+ n 1))))

  (let ((angle-chains '())
	(torsion-chains '())
	(not-unbonded-chains '()))
    (let* ((na (length atom-list))
	   (v (make-vector na))
	   (ch (start-chain)))
      (dotimes
       (i na)
       (vector-set! v i '()))
      (dolist
       (b bond-list)
       (vector-set! v (bond-first b)
		    (cons (bond-second b) (vector-ref v (bond-first b))))
       (vector-set! v (bond-second b)
		    (cons (bond-first b) (vector-ref v (bond-second b)))))
      (set! ch (extend-chain v (extend-chain v ch)))
      (set! angle-chains (unique-chains ch))
      (set! torsion-chains (unique-chains (extend-chain v ch))))
    ;; VDW and electrostatic terms apply to atoms that are unbonded, i.e.
    ;; separated by more than two bonds; this information is buried in
    ;; in the bond list and angle-chains, but must be extracted, and
    ;; duplications thrown away.
    (set! not-unbonded-chains
	  (append
	   (mapcar (make-lambda
		    (b)
		    (let ((b1 (bond-first b))
			  (b2 (bond-second b)))
		      (if (< b1 b2)
			  (list b2 b1)
			  (list b1 b2))))
		   bond-list)
	   (mapcar (make-lambda (ch)
				(list (car ch) (caddr ch)))
		   angle-chains)))
    (do ((L not-unbonded-chains (cdr L))
	 (accum '()))
	((null? L) (set! not-unbonded-chains accum))
      (let ((flag true))
	(mapcar (make-lambda
		 (c)
		 (if (and flag
			  (or (and (= (car c) (car (car L)))
				   (= (cadr c) (cadr (car L)))))
			  (or (and (= (cadr c) (car (car L)))
				   (= (car c) (cadr (car L))))))
		     (set! flag false)))
		(cdr L))
	(if flag (set! accum (cons (car L) accum)))))
    (set! term-list '())
    ;; enumerate length terms
    (dolist (x bond-list)
	    (set! term-list
		  (cons (build-length-term (bond-first x)
					   (bond-second x))
			term-list)))
    ;; enumerate angle terms
    (dolist (a angle-chains)
	    (set! term-list
		  (cons (build-angle-term (car a) (cadr a) (caddr a))
			term-list)))
    ;; enumerate torsion terms
    (dolist (a torsion-chains)
	    (set! term-list
		  (cons (build-torsion-term
			 (car a) (cadr a) (caddr a) (cadddr a))
			term-list)))
    ;; enumerate unbonded terms
    (do ((AL atom-list (cdr AL))
	 (x 0 (+ x 1)))
	((null? AL))
      (do ((A2L (cdr AL) (cdr A2L))
	   (y (+ x 1) (+ y 1)))
	  ((null? A2L))
	(if (not (complex-member (list y x) not-unbonded-chains))
	    (let ()
	      ;;(set! term-list
		;;    (cons (build-electrostatic-term x y) term-list))
	      (set! term-list (cons (build-vdw-term x y) term-list))))))))

;; ******************************************************************
;; Potential energy minimization

;; We compute force vectors, and then scale them until to be displacement
;; vectors, where the largest displacement has a magnitude equal to
;; emin-factor (which is in angstroms).

;; currently, these represent angstroms
(define coarse-emin-factor 0.2)
(define fine-emin-factor 0.01)
(define negligible-force 0.0001)

(define emin-factor fine-emin-factor)

(define (emin-step)
  (entering "emin-step")
  (force-rotate-mode)
  (set! whine-about-bond-count false)
  (set! torsion-whine false)
  (dotimes
   (i 20)
   (compute-forces)
   (dolist (a atom-list)
	   (let ((f (atm-force a)))
	     (if (> (vlen f) negligible-force)
		 (atm-add-pos a
		    (vscale f (/ emin-factor (vlen f))))))))
  (if whine-about-bond-count
      (error-msg "Mismatch between bonds and valences"))
  (if torsion-whine
      (warning-msg "Torsion forces temporarily disabled")))