(use-modules (srfi srfi-1))

;;; SETUP OF SIMULATION PARAMETERS ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define rhos   '(500 1000 2000 3000 3500 4000 4500 5000 10000))
(define betas  '(10 100))
(define thetas '(1.5 3))

;; scaled values calculated with 'LD-study-tree-length.scm'
(define scaled-thetas '(((   0  1.5) . 1.5)
			((   0  3)   . 3)
			((  10  1.5) . 2.56)
			((  10  3)   . 5.12)
			(( 100  1.5) . 4.52)
			(( 100  3)   . 9.04)
			((1000  1.5) . 11)
			((1000  3)   . 22)))

(define K 50)

(define no-genotypes   1300)
(define no-haplotypes  (* 2 no-genotypes))
(define no-markers     15)
(define no-simulations 1000)

(define (product . list-of-lists)
  (letrec ((map-constant (lambda (c lst)
			   (let* ((box (lambda (e) (if (list? e) e (list e)))))
			     (map (lambda (e) (cons c (box e))) lst))))

	   (p (lambda ( . list-of-lists)
		(cond ((null? (cdr list-of-lists)) (car list-of-lists))
		      (else (let* ((first (car list-of-lists))
				   (second (apply p (cdr list-of-lists)))
				   (f (lambda (x) (map-constant x second)))
				   (prod (map f first)))
			      (concatenate prod)))))))
    (apply p list-of-lists)))


(define parameters (product rhos betas thetas))




;;; I/O CODE ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; some of this is now found in modules, but at the time of the study
;; we had to write it ourselves.

(define (combine-genotypes lst)
  (letrec ((combine-haplotypes
	    (lambda (h1 h2) (map (lambda (a1 a2) (cons a1 a2)) h1 h2)))
	   (first 
	    (lambda (lst acc) 
		    (cond ((null? lst) acc)
			  (else (second (car lst) (cdr lst) acc)))))
	   (second
	    (lambda (f lst acc)
	      (first (cdr lst)
		     (cons (combine-haplotypes f (car lst)) acc)))))
    (first lst '())))


(define (print-positions positions)
  (lambda (port)
    (for-each (lambda (pos) (display pos port)(display " " port)) positions)
    (newline port)))

(define (print-genotypes genotypes)
  (lambda (port)
    (let ((print-genotype 
	   (lambda (genotype)
	     (for-each (lambda (g)
			 (let ((a1 (car g)) (a2 (cdr g)))
			   (display a1 port)(display "\t" port)
			   (display a2 port)(display "\t" port)))
		       genotype)
	     (newline port))))
      (for-each print-genotype genotypes))))



;;; SIMULATION ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define (run-simulation rho beta theta dir sim-no)
  (let* ((positions (list-tabulate (+ no-markers 1)
				   (lambda (i) (* i (/ 1.0 (+ 1 no-markers))))))
	 (make-marker (lambda (pos) (ms-marker pos theta K)))
	 
 	 (markers (map make-marker positions))
 	 (haplotypes (simulate-sequences markers no-haplotypes
					 :rho rho :beta beta))
 	 (genotypes (combine-genotypes haplotypes)))

    (call-with-output-file
	(string-append dir "/positions-" (number->string sim-no))
      (print-positions positions))
    (call-with-output-file
	(string-append dir "/genotypes-" (number->string sim-no))
      (print-genotypes genotypes))))


;; run simulations
(define (run-parameter-set rho beta theta)
  (let ((scaled-theta (assoc-ref scaled-thetas (list beta theta)))
	(dir (string-append "theta-" (number->string theta) "--"
			    "rho-"   (number->string rho)   "--"
			    "beta-"  (number->string beta))))

    (display "rho: ")(display rho)(display " ")
    (display "beta: ")(display beta)(display " ")
    (display "theta: ")(display theta)(display " ")
    (display "scaled to: ")(display scaled-theta)(newline)

    (mkdir dir)
    (do ((i 1 (+ i 1))) 
	((> i no-simulations))
      (run-simulation rho beta scaled-theta dir i))))

(for-each (lambda (p) (apply run-parameter-set p)) parameters)