;
; Some extremely convenient functions
;
(define (repeat val n)
"(repeat val n) Return a list of length n filled with val"
  (if (zero? n) '() (cons val (repeat val (- n 1)))))

(define (gl n k . len)
"Args: <number of levels> <number of replicates> [<result length>]"
  (let ((res  (apply append (map (lambda (x) (repeat x k)) (seq 1 n)))))
    (if (or (null? len) (= (car len) (* n k)))
      res
      (apply append (repeat res (floor (/ (car len) (length res))))))))

(define first car)
(define second cadr)
(define third caddr)
(define fourth cadddr)
(define (head x . m) (let ((n (if (pair? m) (car m) 5))) 
                       (list-select x (seq 0 (min (- n 1) (- (length x) 1))))))
(define (tail x . m) (let ((s (if (pair? m) (car m) 5)) (n (- (length x) 1)))
                       (list-select x (seq (max 0 (- n s)) n))))

(define (flatten x)
"Args: x
Flattens a list of lists to a list - matrices by row"
     (let loop ((x x) (r '()))
      (cond ((pair? x) (loop (car x)
                             (loop (cdr x) r)))
            ((null? x) r)
            (else (cons x r)))))

(define (print . x) (display (paste x)) (newline))

(define (count-items item l) 
"Args: <val> <list>
Counts number of occurrences of value in list"
(apply + (map (lambda (x) (if (eqv? item x) 1 0)) l)))

(define (tagval name args)
"(tagval name args) Return argument following tag in list"
  (let* ((n (length args))
         (tagpos (which (map (lambda (x) 
                                 (if (symbol? x) 
                                   (string=? name (symbol->string x)) #f)) args)))
         (pos (if tagpos (+ 1 (apply max tagpos)) n)))
    (if (= pos n) #f (list-ref args pos))))


(define (rowcol nrows ncols len) 
"Args: (nr nc|'rows nr|'cols nc) len
Return suitable nr and nc for a matrix of len elements"
  (cond ((and (number? nrows) (number? ncols)) 
           (list nrows ncols))
        ((and (symbol? nrows) (number? ncols)) 
           (cond ((equal? nrows 'rows) (list ncols (/ len ncols)))
                 ((equal? nrows 'cols) (list (/ len ncols) ncols))
                 (else '())))
        (else '())))

(define (matrix vals nrows ncols)
"Args: (<num>|<number list>) (nr nc)|'rows nr|'cols nc
Returns a matrix."
  (let ((dims (rowcol nrows ncols (if (list? vals) (length vals) 1))))
    (if (null? dims)
       (error "Wrong number of values for matrix dimensions!")
    (let* ((nr (car dims)) (nc (cadr dims)) (len (* nr nc)))
      (if (pair? vals)
       (if (= (length vals) len)
          (let loop ((v vals) (i nc) (row '()) (res '()))
             (if (pair? v)
               (if (zero? i)
                 (loop v nc '() (append res (list (reverse row))))
                 (loop (cdr v) (- i 1) (cons (car v) row) res))
               (append res (list (reverse row)))))
           (begin
             (display "nr=") (display nr) (display " nc=") (display nc) (newline)
             (display "len(vals)=") (display (length vals)) (newline)
             (error "Wrong number of values for matrix dimensions!")))
       (repeat (repeat vals nc) nr))))))

(define (ncol A)
"Args: A
Returns the number of columns in matrix A."
  (if (and (pair? A) (not (null? A)))
     (let ((rowlen (map (lambda (x) (if (pair? x) (length x) 0)) A)))
        (cond ((= (length rowlen) 1) (car rowlen))
              ((apply = rowlen) (car rowlen))
              (else 0)))
     0))
(define (nrow A) 
"Args: A
Returns the number of rows in matrix A."
  (if (pair? A) (length A) 0))

(define (dim A)
  (list (nrow A) (ncol A)))

(define (as-matrix l)
"(as-matrix l) Convert vector to 1-D matrix"
  (if (atomic-list? l) (matrix l (length l) 1) #f))

(define (homogenous-list? l test)
  (if (pair? l) 
     (let loop ((a (car l)) (b (cdr l)))
        (if (not (test a))
           #f
           (if (pair? b)
             (loop (car b) (cdr b))
             #t)))
     #f))
(define (atomic-list? l)
  (homogenous-list? l (lambda (x) (and (not (null? x)) (not (pair? x))))))
(define (numeric-list? l)
  (homogenous-list? l number?))
(define (matrix? A) 
"Args: A
Tests if A is a matrix."
  (if (> (ncol A) 0) 
     (cond ((= (length A) 1) (atomic-list? (car A)))
            (else (apply equal? (map atomic-list? A))))
     #f))

(define (numeric-matrix? A) 
"Args: A
Tests if A is a numeric matrix."
   (if (> (ncol A) 0) 
      (cond ((= (length A) 1) (numeric-list? (car A)))
             (else (apply equal? (map numeric-list? A))))
      #f))
;
; Extensions of apply/map to matrices
;
(define (matrix-apply f X)
   (apply f (map (lambda (x) (apply f x)) X)))

(define (matrix-map f X)
"Args: f X
map f to elements of a matrix"
   (map (lambda (x) (map f x)) X))

(define (matrix-matrix-map f X Y)
"Args: f X Y
map dyadic f to elements of two matrices"
  (if (and (= (nrow X) (nrow Y)) (= (ncol X) (ncol Y)))
     (map (lambda (row1 row2) (map f row1 row2)) X Y)
     (error "Nonconformable arrays!")))

(define (matrix-scalar op X y)
"Args: op X y
element-wise application of operation to matrix and scalar."
  (matrix-map (lambda (x) (op x y)) X))

(define (rowsum X) 
  (map (lambda (x) (apply + x)) X))

(define (colsum X)
  (apply map + X))

(define (matrix-filter-row f X)
"Args: f X
keep eligible rows of a matrix"
  (if (matrix? X)
    (let loop ((rest X) (result '()))
      (if (null? rest)
         result
         (let ((thisrow (car rest)))
           (if (equal? thisrow (filter f thisrow))
              (loop (cdr rest) (cons thisrow result))
              (loop (cdr rest) result)))))
    (error "Not a matrix!")))

(define (matrix-filter-col f X)
"Args f X
keep eligible columns of a matrix"
  (if (matrix? X)
    (let loop ((nc (ncol X)) (rest X) (result '()))
      (if (= nc 0)
         result
         (let ((thiscol (map car rest)))
(display thiscol) (display " f= ") (display (filter f thiscol)) (newline)
           (if (equal? thiscol (filter f thiscol))
              (loop (- nc 1) (map cdr rest) (cons thiscol result))
              (loop (- nc 1) (map cdr rest) result)))))
    (error "Not a matrix!")))

(define (transpose A)
"Args: A
Transposes matrix A."
  (let loop ((W A) (res '()))
     (if (pair? (car W))
        (loop (map cdr W) (append res (list (map car W))))
        res)))
;
; Get or set matrix elements
;
(define (row A r) 
"(row A r) returns rth row of matrix A"
  (list-ref A (- r 1)))
(define (col A c) 
"(col A c) returns cth column of matrix A"
  (map (lambda (x) (list-ref x (- c 1))) A))

(define (matrix-ref A i j) 
"(matrix-ref A i j) returns (i,j)th element of matrix A"
  (list-ref (list-ref A (- i 1)) (- j 1)))

(define (list-element-set! l k val)
  (let loop ((i 1) (w l) (newl '()))
     (if (< i k)
        (loop (+ i 1) (cdr w) (cons (car w) newl))
        (append (reverse newl) (list val) (cdr w)))))

(define (matrix-row-set! A r newrow)
"(matrix-row-set! A r newrow) Set the rth row of A"
  (let loop ((i 1) (W A) (newA '()))
     (if (< i r)
        (loop (+ i 1) (cdr W) (append newA (list (car W))))
        (append newA (list newrow) (cdr W)))))

(define (matrix-col-set! A c newcol)
"(matrix-col-set! A r newcol) Set the cth column of A"
  (let loop ((col newcol) (W A) (newA '()))
     (if (pair? col)
        (loop (cdr col) (cdr W) (append newA (list (list-element-set! (car W) c (car col)))))
        newA)))

(define (matrix-set! A i j val) 
"(matrix-set! A i j) Sets (i,j)th element of matrix A"
  (matrix-row-set! A i (list-element-set! (row A i) j val)))
;
; Vertical and horizontal concatenation
;
(define (rbind A B)
"Args: A B
Concatenate matrices A and B by rows."
  (if (= (ncol A) (ncol B))
     (append A B)
     (error "Nonconformable arrays!")))

(define (cbind A B)
"Args: A B
Concatenate matrices A and B by cols."
  (if (= (nrow A) (nrow B))
     (map append A B)
     (error "Nonconformable arrays!")))

(define (concat . vlist)
"Args: <list of vectors or matrices>
Concatenate column vectors and matrices horizontally."
  (if (list? vlist)
    (if (list? (car vlist))
      (let loop ((n (length (car vlist)))
                 (res (if (matrix? (car vlist)) 
                         (car vlist) (matrix (car vlist) 'cols 1)))
                 (rest (cdr vlist)))
;;;;      (display "<-") (display res) (newline)
;;;;      (display "->") (display rest) (newline)
          (cond ((null? rest) res)
                ((and (matrix? (car rest)) (= n (ncol (car rest))))
                    (loop n (cbind res (car rest)) (cdr rest)))
                ((and (list? (car rest)) (= n (length (car rest))))
                    (loop n (cbind res (matrix (car rest) n 1)) (cdr rest)))
                (else (error "Nonconformant number of rows!"))))
      (error "Not a vector!"))
    (error "Expecting a list of vectors/matrices!")))
;
; Inner product
;
(define (dotprod a b) "inner product" (apply + (map * a b)))
;
; Outer product
;
(define (outer-map op a b)
"Args: op a b
Generalized outer product of vectors a and b"
  (cond ((and (pair? a) (pair? b))
            (let loop ((a a) (res '()))
              (if (pair? a)
                 (let ((i (car a)))
                   (loop (cdr a) (append res (list (map (lambda (x) (op i x)) b)))))
                 res)))
         ((pair? a)
            (map (lambda (x) (list (op b x))) a))
         ((pair? b)
            (map (lambda (x) (op a x)) b))
         (else (list (op a b)))))
(define (outer-product a b) 
"Args: a b
Outer product of lists a and b"
(outer-map * a b))
;
; Vector by matrix multiplication
;
(define (mult-vec-tmat a B) (map (lambda (x) (dotprod a x)) B))
(define (vecmat a B) 
"Args: a B
multiply vector by conformant matrix"
  (mult-vec-tmat a (transpose B)))
;
; Matrix multiplication
;
(define (matmult A B)
"Args: A (number or matrix) B (matrix)
matrix multiplication A * B"
  (cond 
    ((number? A) (matrix-map (lambda (x) (* A x)) B))
    (else (let ((r (nrow A)) (c (ncol B)) (common (ncol A)))
            (if (= common (nrow B))
              (let ((Btr (transpose B)))
                (let loop ((W A) (res '()))
                   (if (pair? W)
                     (loop (cdr W) 
                           (append res (list (mult-vec-tmat (car W) Btr))))
                     res)))
              (error "Nonconformable arrays!"))))))
;
; Matrix addition and subtraction
;
(define (matadd A B) 
  "matrix addition"
  (cond ((number? A) (matrix-map (lambda (x) (+ A x)) B))
        (else (matrix-matrix-map + A B))))
(define (matsub A B) 
  "matrix substraction"
  (matrix-matrix-map - A B))
;
; Sweep operator and multiple regression
;
(define (sweep X k) 
"(sweep X k) Apply sweep operator kth pivot to matrix X"
  (let ((pivot (matrix-ref X k k)))
     (if (zero? pivot)
        #f
        (let* ((d (/ 1 pivot))
               (res (matrix-row-set! X k 
                       (map (lambda (x) (* x d)) (row X k))))
               (b (list-element-set! (col res k) k 0)))
          (matrix-set! (matrix-col-set! (matsub res (outer-product b (row res k))) k 
                          (map (lambda (x) (* x (- d))) b)) k k d)))))

(define (inverse A)
"(inverse A) Inverse of matrix A"
  (let ((nv (nrow A)))
    (let loop ((i 1) (A A))
      (if (> i nv) 
         A
         (loop (+ i 1) (sweep A i))))))
;
; diag
;
(define (diag A)
  "extract diagonal of a square matrix"
  (let ((nr (nrow A)) (nc (ncol A)))
    (if (= nr nc)
      (let loop ((i 1) (A A) (res '()))
        (if (<= i nr)
           (loop (+ 1 i) (cdr A) (cons (list-ref (car A) (- i 1)) res))
           (reverse res)))
      (error "Not a square matrix!"))))
;
; create diagonal matrix
;
(define (diag-matrix n . diag-values)
"diag-matrix n [<diag-values>]
create diagonal matrix with specified values, defaulting to identity"
  (let ((d (cond ((null? diag-values) 
              (repeat 1 n))
           ((number? (car diag-values)) 
              (repeat (car diag-values) n))
           ((>= (length (car diag-values)) n) 
              (list-select (car diag-values) (seq 0 (- n 1))))
           (else (error "Provide 0, 1, n values for diag")))))
   (let loop ((i 1) (d d) (res '()))
     (if (> i n) 
        res
        (loop (1+ i) (cdr d) 
              (append res (list (append (repeat 0 (- i 1))
                                         (list (car d))
                                         (repeat 0 (- n i))))))))))
    
;
; Print method for matrices
;
(define (longest-string x) 
  (apply max (map string-length x)))

(define (longest-num x) 
  (apply max (map (lambda (x) 
    (string-length (if (number? x) (number->string x) x))) (flatten x))))

(define (str-form x)
  (string-append "~" (number->string (longest-string x)) "A"))

(define (num-form x)
  (string-append "~" (number->string (longest-num x)) "D"))

(define (formatted-row-print w a) 
  (display "[ ")
  (for-each (lambda (x) (format w x)) a) 
  (display "]")
  (newline))

(define (row-print a) 
  (display "[ ")
  (for-each (lambda (x) (format "~10d " x)) a) 
  (display "]")
  (newline))

(define (matrix-print A . names) 
"Args: A . names
Pretty prints matrix A with optional row and column labels."
  (if (matrix? A)
    (let ((width (number->string (longest-num A))))
     (if (and (> (length names) 0) (= (length (car names)) (length A)))
       (let* ((rowfor (string-append (str-form (car names)) " "))
             (width (if (> (length names) 2) (number->string (caddr names)) width))
             (headfor (string-append "~" width "@A "))
             (colfor (string-append "~" width "D ")))
         (newline)
         (if (> (length names) 1)
           (begin
             (format rowfor " ") (display "  ")
             (for-each (lambda (x) (format headfor x)) (cadr names))
             (newline)))
         (let loop ((x (car A)) (rest (cdr A)) (title (caar names)) (ret (cdar names)))
           (format rowfor title)
           (formatted-row-print colfor x)
           (if (not (null? rest)) (loop (car rest) (cdr rest)  (car ret) (cdr ret))
                                  (newline))))
       (for-each (lambda (x) (formatted-row-print (string-append "~" width "D ") x)) A)))
     (display "Not a matrix!")))

(define (mprint A)
  (matrix-print (transpose A)))
;
; Hyperbolic functions
;
(define (1+ x) (+ 1 x))
(define (1- x) (- x 1))
(define (sinh x) (/ (- (exp x) (exp (- x))) 2.0))
(define (cosh x) (/ (+ (exp x) (exp (- x))) 2.0))
(define (tanh x) (/ (sinh x) (cosh x)))
(define (asinh x) (log (+ x (sqrt (+ 1.0 (* x x))))))
(define (acosh x) (log (+ x (* (1+ x) (sqrt (/ (1- x) (1+ x)))))))
(define (atanh x) (log (/ (1+ x) (sqrt (- 1.0 (* x x))))))
;
; Allow compare accuracy of logpnorm to high precision erfc
;  functions in programs like Mupad
;
(define (logerfc x) (+ (log10 2) (logpnorm (* x (sqrt 2)))))
;
; Stirling number of second kind
;
(define (stirling2 n k) 
  (define (s2 n k) 
    (if (or (= k 1) (= k n)) 
       1 
       (+ (s2 (- n 1) (- k 1)) (* k (s2 (- n 1) k))))) 
  (cond  
    ((negative? n) (error "n < 0")) 
    ((negative? k) (error "k < 0"))
    ((not (integer? n)) (error "n not an integer"))
    ((not (integer? k)) (error "k not an integer"))
    ((< n k) 0) 
    ((= n k 0) 1) 
    ((= k 0) 0) 
    (else (s2 n k))))
;
; a few useful functions
;
(define (square x) (* x x))
(define (sumsquare x) (apply + (map square x)))

(define (sum x)
"Args: (x)
Returns the sum of the elements x. List reducing."
  (apply + x))

(define (mean x) 
"Args: (x)
Returns the mean of the elements x. List reducing."
  (list-ref (stats x) 2))

(define (standard-deviation x)
"Args: (x)
Returns the standard deviation of the elements x. List reducing."
  (sqrt (list-ref (stats x) 3)))

(define (minimum x)
"Args: (x)
Returns the minimum of the elements of X."
  (list-ref (stats x) 4))

(define (maximum x)
"Args: (x)
Returns the minimum of the elements of X."
  (list-ref (stats x) 5))

(define (median x)
"Args: (x)
Returns the median of the elements of X."
  (quantile x 0.5))

(define (interquartile-range x)
"Args: (number-data)
Returns the interquartile range of the elements of X."
  (apply - (quantile x '(0.75 0.25))))

(define (fivnum x)
"Args: (number-data)
Returns the five number summary (min, 1st quartile, median, 3rd quartile,
max) of the elements X."
   (quantile x '(0 .25 .5 .75 1)))

(define (which-min x) 
"Args? (number-data)
Returns position of minimum value(s) in list"
  (let ((minx (list-ref (stats x) 4))) 
    (which (map (lambda (x) (and (number? x) (= x minx))) x))))

(define (which-max x) 
"Args? (number-data)
Returns position of maximum value(s) in list"
  (let ((maxx (list-ref (stats x) 5))) 
    (which (map (lambda (x) (= x maxx)) x))))

(define (rseq a b num)
"Args: (a b num)
Returns a list of NUM equally spaced points starting at A and ending at B."
  (map (lambda (x) (+ a (* (/ (- b a) num) x))) (seq (- num 1))))

(define (ppoints n)
"Args: number of points
Returns NUM evenly spaced points between 0 and 1 suitable for Q-Q diagnostic plots"
  (let* ((a (if (> n 10) (/ 3 8) (/ 1 2))))
    (map (lambda (x) (/ (- x a) (+ n 1 (* (- 2) a)))) (seq 1 n))))

(define (accumulate el tot) 
"(accumulate el tot) Add number to running total"
  (if (number? el) (+ el tot) tot))

(define (cumsum x)
"Args: (x)
Returns the cumulative sum of X."
  (cond ((list? x)
           (let loop ((el (car x)) (rest (cdr x)) 
                      (tot 0) (res '()))
              (if (null? rest)
                 (reverse res)
                 (let ((n (accumulate el tot)))
                   (loop (car rest) (cdr rest) n (cons n res))))))
         ((number? x) x)
         (else '())))
;
; proportions
;
(define (proportions x)
  "Proportions of total sum of a sequence"
  (let ((n (apply + x)))
    (map (lambda(x) (/ x n)) x)))
;
; factorial
;
(define (factorial n)
"Simple iterative factorial n"
  (let loop ((i 2) (res 1))
    (if (> i n) res (loop (1+ i) (* res i)))))
; 
; Fisher method for combining P-values
; 
(define (fisher-meta pvals) 
  "fisher-meta <p-values>: unweighted combination of P-values"
  (let* ((k (length pvals)) 
         (S (* -2 (apply + (map log pvals))))) 
     (pchisq S (* 2 k))))
;
; Lancaster's generalization of Fisher combination of P values
;
(define (lancaster-meta x n)
  "lancaster-meta <p-values> <sample-sizes>: weighted combination of P-values"
  (if (and (pair? x) (pair? n))
     (if (= (length x) (length n))
        (pchisq (apply + (map qchisq x n)) (apply + n))
        (error "x and n different lengths!"))
     (error "x and n have to have length > 1")))
;
; Subset a matrix by rows on values of index list
;
(define (select-rows X idx)
"select-rows X idx
Subset a matrix by rows on values of index list"
  (let ((nr (nrow X)) (nc (ncol X)) (n (length idx)))
    (if (= nr n)
      (list-select X (which idx))
      (error "Length of idx should match number of rows of X"))))

(define (select-cols X idx)
"select-cols X idx
Subset a matrix by rows on values of index list"
  (let ((nr (nrow X)) (nc (ncol X)) (n (length idx)))
    (if (= nc n)
      (transpose (list-select (transpose X) (which idx)))
      (error "Length of idx should match number of cols of X"))))

;
; Subset a list by values of index lists of same length 
; mapping function to each subset
;
(define (one-level x idx target f)
   (let ((w (which (map (lambda (y) (eqv? y target)) idx))))
    (if (matrix? x)
      (let ((selected (map (lambda (y) (list-select y w)) x)))
        (write selected) (display " matrix? ") (display (matrix? selected)) (newline)
        (if (matrix? selected) (map f selected) (apply f selected)))
      (f (list-select x w)))))

(define (by x idx f)
  "by <var> <strat> <function>: applies f to var for each level of strat"
  (cond ((not (pair? x))  
           (format "ERROR: ~s is not a list!~%" (quote x))
           (display "Usage: (by x idx f)") (newline))
        ((not (pair? idx)) 
           (format "ERROR: ~s is not a list!~%" (quote idx))
           (display "Usage: (by x idx f)") (newline))
        ((not (procedure? f)) 
           (format "ERROR: ~s is not a function!~%" (quote f)))
        ((and (matrix? x) (not (= (ncol x) (length idx))))
           (format "ERROR: col length ~s is ~d, but length ~s is ~d~%" 
              (quote x) (ncol x) (quote idx) (length idx))
           (display "Usage: (by x idx f)") (newline))
        ((and (not (matrix? x)) (not (= (length x) (length idx))))
           (format "ERROR: length ~s is ~d, but length ~s is ~d~%" 
              (quote x) (length x) (quote idx) (length idx))
           (display "Usage: (by x idx f)") (newline))
        (else
          (let ((levels (quicksort (unique idx))))
             (let loop ((target (car levels)) (rest (cdr levels)) (res '()))
                (if (not (null? rest))
                   (loop (car rest) (cdr rest) 
                         (cons (one-level x idx target f) res))
                   (list levels (reverse (cons (one-level x idx target f) res)))))))))

(define (by-subset byresult indx)
"by-subset <by-output> <indx>: 
Extracts value corresponding to index stratum."
  (car (list-select (cadr byresult) 
    (which (map (lambda (x) (string=? indx x)) (car m_c))))))

(define (summary x idx)
"Args: (number-data) (idx) 
Gives 6 number summary of x for each level of idx"
  (let ((tab (by x idx  stats)))
    (matrix-print (transpose (cadr tab)) (list "N" "Miss" "Mean" "Var" "Min" "Max") (car tab))))

(define (indicator target x)
"Args: target x
Return indicator vector for target = x"
  (map (lambda (y) (if (equal? target y) 1 0)) x))

;;
;; tabulate 
;;
(define (inc-table ht val . w)
"Args: hash-table val [sampling-weight]
Increment count in table for given entry"
  (let  ((count  (hash 'ref ht val))
	 (wt (if (null? w) 1 (car w))))
    (if (null? count)
      (hash 'set! ht val wt)
      (hash 'set! ht val (+ count wt)))))

(define (tabulate x . weights)
"Args: list of values
Returns hash table of unique values and counts"
  (let ((tab (hash 'make)) (var x))
    (if (and (not (null? weights)) 
             (= (length (car weights)) (length var)))
      (map (lambda (x y) (inc-table tab x y)) var (car weights))
      (map (lambda (x) (inc-table tab x)) var))
    tab))

(define (tprint ht)
"Args: hash table of values and their counts
Prints contingency table"
   (let* ((v (max 5 (longest-num (car ht))))
          (vwid (number->string v))
          (w (longest-num (cadr ht)))
          (cwid (number->string w))
          (fstr (string-append "~" vwid "A  ~" cwid "A~%")))
     (newline)
     (format fstr "Var" "Freq")
     (format fstr (make-string v "~") (make-string w "~"))
     (let loop ((levels (car ht)) (counts (cadr ht)) (total 0))
       (if (pair? levels)
	 (begin
	   (format fstr (car levels) (car counts))
	   (loop (cdr levels) (cdr counts) (+ total (car counts))))
         (begin
           (format fstr (make-string v "~") (make-string w "~"))
           (format fstr "Total" total)
           (format fstr (make-string v "~") (make-string w "~")))))))

(define (main-effect x . full)
"Args: x . <drop-first-level>
Return indicator matrix for factor level membership of vector"
  (let ((all-levels (unique x)))
    (if (not (null? all-levels))
       (let* ((nr (length x))
              (drop-level (if (null? full) #t (car full)))
              (levels (if drop-level (cdr all-levels) all-levels))
              (nlev (length levels)))
          (if (> nlev 0) 
            (let loop ((rest levels)
                       (res '()))
               (if (not (null? rest))
                  (loop (cdr rest) (cons (indicator (car rest) x) res))
                  res))
            (error "No variation in levels"))))))

(define (vecmat-el-mult x Y) 
  (map (lambda (Y) (map * x Y)) Y))

(define (interaction X Y)
"Args: X Y
Return indicator matrix for crossing of factor indicator matrices"
  (if (and (matrix? X) (matrix? Y))
    (if (= (ncol X) (ncol Y))
      (let ((nr (nrow X)))
        (let loop ((i 1) (X (cdr X)) (res (vecmat-el-mult (car X) Y)))
          (if (= i nr)
             res 
             (loop (+ 1 i) (cdr X) (append (vecmat-el-mult (car X) Y) res)))))
      (error "Nonconformant matrices"))
    (error "X and Y must be matrices")))
;
; Evaluate list of formula terms to return design matrix
; A B C vectors of category levels -> factor matrix A B C
; (* A B) factor crossing
;
(define (factor x) (main-effect x))
(define (cross . terms)
  (newline) (display "CROSS ") (display (length terms)) (display ":") (newline)
  (display terms) (newline) (newline)
  (if (list? terms)
    (if (> (length terms) 2)
      (interaction (main-effect (car terms)) (apply cross (cdr terms)))
      (interaction (main-effect (car terms)) (main-effect (cadr terms))))
    (error "Expected a list of model terms!")))
  
(define (eval-formula f)
"Args: <design formula>
Return appropriate design matrix"
  (if (list? f)
    (let loop ((i 1) (first (car f)) (res '()) (rest (cdr f)))
      (let ((op (if (and (list? first) (eqv? '* (car first))) 
                   (cons 'cross (cdr first)) (cons 'factor (list first)))))
        (if (not (null? rest))
         (loop (+ i 1) (car rest) (cons op res) (cdr rest))
         (cons 'append  (reverse (cons op res))))))
    (error "Expected a list representing a model formula")))

(define (drop-term f term)
  (list-select f (setdiff (seq 0 (- (length f) 1)) (list term))))

(define (add-inter inter i res)
    (if inter (cons i res) res))

(define (interaction-index f)
"Args: <design formula>
Return indices of interaction terms"
  (if (list? f)
    (let loop ((i 0) (first (car f)) (res '()) (rest (cdr f)))
      (let ((inter (and (list? first) (eqv? '* (car first)))))
        (if (not (null? rest))
         (loop (+ i 1) (car rest) (add-inter inter i res) (cdr rest))
         (reverse (add-inter inter i res)))))
    (error "Expected a list representing a model formula")))
      
(define (loglin y formula . flags)
"Args counts S-formula
Fit a log-linear model."
  (let* ((X (eval (eval-formula formula)))
         (ni (length y))
         (np (nrow X))
         (df (- ni np 1))
         (m1 (gllm y (flatten X) flags))
         (lrts (cadr m1)))
   (if (and (not (null? flags)) (eqv? (car flags) 'v))
      (begin
        (display "Observed     Expected    Z") (newline)
        (map (lambda (o e) (format " ~6D    ~8,1D   ~5,1D~%" 
                (inexact->exact o) e (/ (- o e)))) y (car m1))
        (newline) (display formula) (newline)
        (newline) (display "N=") (display (apply + (map inexact->exact y)))
        (newline) (display "LRTS=") (display lrts)
        (display "  df=") (display df) 
        (display " P=") (display (pchisq lrts df)) (newline)))
    (list formula lrts df)))

(define (loglin-anova y formula)
"Args counts S-formula
Fit a log-linear model and drop X2 for interaction terms."
  (let ((form-index (seq 0 (- (length formula) 1)))
        (iterms (interaction-index formula))
        (m0 (loglin y formula)))
    (let loop ((iterms iterms) 
               (res (list (list "Full" (cadr m0) (caddr m0) 
                             (pchisq (cadr m0) (caddr m0))))))
      (if (not (null? iterms))
        (let* ((curr (car iterms))
               (term (list-select formula curr))
               (f1 (drop-term formula curr))
               (m1 (loglin y f1))
               (x2 (- (cadr m1) (cadr m0)))
               (df (- (caddr m1) (caddr m0)))
               (pval (pchisq x2 df)))
           (loop (cdr iterms) 
                 (cons (list (paste term) x2 df pval) res)))
         (reverse res)))))

(define (lca design counts nclasses)
"Args: design_matrix count_array nclasses 
Latent Class Analysis using EM gllm."
  (let* ((C design) 
         (ni (nrow C))
         (Z (matrix 0 ni (ncol C)))
         (X (matrix (append (repeat 0 32) (repeat 1 32)) 1 64)) 
         (M (rbind X (rbind (cbind C Z) (cbind Z C)))) 
         (S (gl 32 1 64)) 
         (m1 (gllm counts (flatten M) 'scatter S 'maxit 1000))) 
  (matrix-print (transpose (rbind (matrix counts 1 32) 
    (rbind (matrix (cadr m1) 1 32) (matrix (car m1) 2 32)))) 
    (map number->string (seq 1 32)) '("Observed" "Expected" "LC1" "LC2")) 
    (newline) (display "LRTS=") (display (caddr m1)
    (display "  df=") (display (- ni nclasses (* 2 (ncol C)))))))

;
; Power of cluster randomized trial
;
(define (cluster-power k m p1 p2 r alpha) 
"Args: <nclust_arm> <clust_size> <p1> <p2> <icc> <alpha>  
Power of cluster randomized trial."
  (let* ((total (* 2 m k)) 
         (VIF (+ 1 (* (- m 1) r))) 
         (d (abs (- p1 p2))) 
         (s2 (* 0.5 (+ (* p1 (- 1 p1)) (* p2 (- 1 p2))))) 
         (v (/ (* s2 2 VIF) (* m k))) 
         (Z (/ d (sqrt v))) 
         (power (- 1 (pnorm (- Z (qnorm (/ alpha 2))))))) 
    (format "VIF   = ~6,3D~%" VIF) 
    (format "d     = ~6,3D~%" d) 
    (format "s2    = ~6,3D~%" s2) 
    (format "v     = ~6,3D~%" v) 
    (format "E(X2) = ~6,3D~%" (* Z Z)) 
    (format "alpha = ~6,3D~%" alpha) 
    (format "Power = ~6,3D~%" power) 
    (format "Tot N = ~6D~%" total) 
    power))
;
; Algorithm AS 319 Unconstrained variable metric function minimization without derivatives.
; Translated from the Fortran code by DLD Jul 2026
; Koval JJ. Variable Metric Function Minimization
; J R Statistl Soc. Series C (Applied Statistics)
;   Vol. 46, No. 4 (1997), pp. 515-521
;
(define (grad fun b f0 er)
"(grad fun b f0 g sa er) Estimate gradient of function"
  (define (perturb b i val)
  "(perturb b i val) Peturb ith value"
    (append (list-select b (seq 0 (- i 1)))
            (list val) (list-tail b (+ i 1))))
  (let* ((npar (length b)) (jcmax (- npar 2)) (ser (sqrt er)))
    (let loop ((i 0) (jc 0) (g '()))
      (cond 
        ((= i npar) (reverse g))
        (else 
          (let* ((this (list-ref b i))
                 (h (* (+ (abs this) ser) ser)) 
                 (f1 (fun (perturb b i (+ this h)))))
            (loop (+ 1 i) jc (cons (/ (- f1 f0) h) g))))))))

(define (varmet fun b gradtl toler maxfn0 plevel)
"(varmet fun b gradtl toler maxfn ifault plevel) Maximize function"
  (define (print-iter label ifn f1 b)
     (print label ": " ifn f1 b))
  (define (update-d d1 h c)
    (let ((n (nrow h)))
      (let loop1 ((i 0) (d2 0) (newd '()))
         (if (< i n)
            (let ((s (dotprod (row h (+ 1 i)) c)))
              (loop1 (+ 1 i) (+ d2 (* s (list-ref c i))) (cons s newd)))
             (list (+ 1 (/ d2 d1)) (reverse newd))))))

  (define (update-h d1 h t c)
  "(update-h d1 h t c) Update numerical estimate Hessian"
    (let* ((n (nrow h)) (n2 (* n n)) (newd (update-d d1 h c)) (d2 (car newd)) (d (cadr newd)))
      (let loop1 ((i 0) (j 0) (k 1) (ti (car t)) (di (car d)) 
                  (oldh (flatten h)) (newh '()))
        (let* ((tj (list-ref t j)) (dj (list-ref d j))
               (hterm (- (car oldh) (/ (- (+ (* di tj) (* ti dj)) 
                                        (* d2 ti tj)) d1))))
           (cond ((= k n2)  (matrix (reverse (cons hterm newh)) n n))
                 ((= j (- n 1)) (loop1 (+ 1 i) 0 (+ 1 k) (list-ref t (+ 1 i)) 
                                       (list-ref d (+ 1 i)) (cdr oldh) (cons hterm newh)))
                 (else (loop1 i (+ 1 j) (+ 1 k) ti di (cdr oldh) (cons hterm newh))))))))
      
  (define  ig 0)
  (define  ifn 1)
  (define  ler #f)
  (define  ifault 0)
  (define  npar (length b))
  (define  np (+ npar 1))
  (define maxfn (if (= maxfn0 0) 1000 maxfn0))
  (define w 0.2)
  (define icmax 20)
  (define  f0 (fun b))
  (if (> plevel 0)
    (begin
      (print ";; npar=" npar "gradtl=" gradtl "toler=" toler "maxfn=" maxfn)
      (print-iter ";; Init" ifn f0 b)))

  (define initgrad (grad fun b f0 gradtl))
  (set! ig (+ 1 ig))
  (set! ifn (+ ifn npar))
  (if (> plevel 0) (print ";; Initial g=" initgrad))

  (if (> ifn maxfn) 
    (error (paste "More (" ifn ") parameters than maximum allowed (" 
                  maxfn ") function evaluations!")))

;;;; label 10 reset h to identity matrix
  (let mainloop ((lik f0) (b b) (g initgrad) (h (diag-matrix npar)) (ilast ig))
;;;; outer iteration
      (let dloop ((i 0) (d1 0.0) (b b) (g g) (h h) (t '()))
;;;; (print "outer iteration t=" t "d1=" d1 "i=" i "g=" g "h=" h)
        (let ((d b) (c g))
          (if (< i npar) 
            (let ((s (- (dotprod (row h (+ i 1)) g))))
              (dloop (+ i 1) (- d1 (* s (list-ref g i))) b g h (append t (list s))))
            (if (<= d1 0.0) 
               (if (= ilast ig) 
                 (list fun b (fun b) ifault) ;; return
                 (mainloop lik b g (diag-matrix npar) ilast))
;;; label 90 inner iteration
               (let loop90 ((ck 1) (ic 0) (d d) (t t))
                  (let* ((b (map + d (map (lambda (x) (* ck x)) t)))
                         (icount (sum (map (lambda (x) (if x 1 0)) (map = b d)))))
;;;; (print "inner iteration" ic "icount=" icount "ig=" ig "ilast=" ilast " ck=" ck "t=" t "d=" d)
                   (if (>= icount npar)
                     (if (= ilast ig) 
                       (list fun b (fun b) ifault) ;; return
                       (mainloop lik b g (diag-matrix npar) ig))
                     (let ((f1 (fun b)))
                       (set! ifn (+ 1 ifn))
                       (if (> plevel 0) (print-iter ";; Iter" ifn f1 b))
                       (cond ((> ifn maxfn) (display "ifault=4 ") (list fun b f1 4)) ;; return
                             (ler (if (= ic icmax) 
                                     (list fun b f1 3) ;; return
                                     (loop90 (* w ck) (+ ic 1) d t)))
                             ((>= f1 (- f0 (* d1 ck toler)))
;;;; (print "Back to loop90 f1=" f1 "f0=" f0 " d1*ck*toler=" (* d1 ck toler))
                                 (loop90 (* w ck) (+ ic 1) d t))
                             (else
                               (set! f0 f1)
                               (set! g (grad fun b f0 gradtl))
;;;; (print "Proceeding g=" g " t=" t " ck=" ck " c=" c)
                               (set! ig (+ 1 ig))
                               (set! ifn (+ npar ifn))
                               (if (> ifn maxfn) 
                                 (list fun b f0 4) ;; return
                                 (let* ((t (map (lambda (x) (* ck x)) t))
                                        (c (map - g c))
                                        (d1 (apply + (map * t c))))
;;;; (print "new t=" t " c=" c " d1=" d1)
                                   (if (<= d1 0)
                                     (mainloop f0 b g (diag-matrix npar)  ig)
                                     (dloop 0 0.0 b g (update-h d1 h t c) '()))))))))))))))))  
;
; Linear regression
;
(define (reg y X) 
"(reg y X) Perform multiple linear regression via sweep operator"
  (let* ((XY (concat X y)) (nv (+ (ncol XY) 1)) (W (cbind (matrix 1 (nrow XY) 1) XY)))
    (let loop ((i 1) (A (matmult (transpose W) W))) 
      (if (= i nv) 
         (list (col A nv) (- nv 2) (/ (matrix-ref A nv nv) (- (nrow X) (ncol XY))) A)
         (loop (+ i 1) (sweep A i))))))

(define (mvnlik pars)
"(mvnlik pars)  MNV likelihood function for maximization pars: (fixef ranef vare)"
  (let* ((fixef (list-select pars (seq 1 nx)))
         (ranef (list-ref pars (+ 1 nr)))
         (vare (list-ref pars (+ 2 nr)))
         (d (matsub y (vecmat fixef IX)))
         (S (matadd (vecmat ranef ZZ) (matmult vare E)))
         (invS (inverse S))
         (logdet (log (sum (diag (inverse S)))))
         (qform (dotprod (vecmat d invS) d)))
    (* -0.5 (+ logdet qform))))

(define (lmm y X Z . args)
"(llm y X Z) Fit linear mixed model via direct maximization of the likelihood"
  (cond 
    ((not (matrix? X)) (error "X must be a matrix!"))
    ((not (matrix? Z)) (error "Z must be a matrix!"))
    ((not (= (ncol Z) (length y)) ) (error "Z nonconformant with y!"))
    (else
  (let* ((n (length y)) (vary (list-ref (stats y) 3)) (IX  (cbind (matrix 1 n 1) X)) 
         (nx (ncol IX)) (nr (ncol Z)) 
         (E (diag-matrix n)) (ZZ (matmult (transpose Z) Z)) (beta (head (car (reg y X)) nx))
         (plevel (count-items 'v args)) (gradtl ))
    (display "n=") (display n) (display " Var(y)=") (display vary)
    (display " nx=") (display nx) 
    (display " nr=") (display nr) (newline)
    (display "beta0= ") (display beta) (newline)
    (varmet mvnlik (append beta 0 vary) 1.0e-12 1.0d-7 0 plevel)))))