(use-modules (ice-9 format) (coasim batch))

(define no-leaves 20)
(define no-samples 1000)
(define rhos '(0.3 1 1.5 2)); 3 4))


(define (calc-scale-factor beta)
  (let ((no-growth-tree-length 
	 (let loop ((j 1) (sum 0))
	   (if (< j no-leaves) 
	       (loop (+ j 1) (+ sum (/ 1 j)))
	       sum)))
	(growht-tree-length 
	 (let* ((no-iterations 10000)
		(branch-sum
		 (fold no-iterations (lambda (val sum) (+ val sum)) 0
		       (let* ((m (snp-marker 0.5 0 1))
			      (arg (simulate (list m) no-leaves :beta beta))
			      (tree (car (local-trees arg)))
			      (branch-length (total-branch-length tree)))
			 branch-length))))
	   (/ branch-sum no-iterations))))
    (/ no-growth-tree-length growht-tree-length)))

(define rho-scale-factor
  ;; map from beta to scale factor
  ;;          beta scale-factor
  (list (cons   0  1)
	(cons  10  (calc-scale-factor 10))
	(cons 100  (calc-scale-factor 100))))

  

(define (simulate-rho-beta rho beta)
  (let* ((scaled-rho (* rho (assoc-ref rho-scale-factor beta)))
	 (reject (lambda (n1 n2 k) (throw 'reject)))
	 (terminal-branch-length
	  (lambda (arg)
	    (let ((f (lambda (n l)
		       (cond
			((coalescent-node? n)
			 (let* ((cs (children n))
				(child1 (car cs))
				(child2 (cadr cs)))
			   (+ l
			      (if (leaf-node? child1)
				  (- (event-time n) (event-time child1))
				  0)
			      (if (leaf-node? child2)
				  (- (event-time n) (event-time child1))
				  0))))
			((leaf-node? n) l)
			(else (throw 'unexpected-node))))))
	      (fold-nodes arg f 0))))
	 (try (lambda ()
		(let* ((arg (simulate '() no-leaves
				      :rho scaled-rho :beta beta
				      :recombination-callback reject
				      :keep-empty-intervals #t))
		       (tree (car (local-trees arg))))
		  (list (terminal-branch-length arg)
			(total-branch-length tree)
			(tree-height tree)))))
	 (except (lambda (key . args) #f))
	 (sim (lambda ()
		(catch 'reject try except))))
    
    (let loop ((n 1)
	       (reject 0))
      (if (< no-samples n)
	  reject
	  (let ((summary (sim)))
	    (if summary
		(begin 
		  (format #t "~f ~f ~f ~f ~f\n"
			  beta rho 
			  (car summary) (cadr summary) (caddr summary))
		  (loop (+ n 1) reject))
		(loop n (+ reject 1))))))))

(call-with-output-file "rejections.txt"
  (lambda (port)

    ;; beta 0
    (let* ((simulate-rho (lambda (rho) (simulate-rho-beta rho 0)))
	   (rejections (map simulate-rho rhos)))
      (for-each (lambda (rho r) 
		  (display "0 " port)
		  (display rho port )(display " " port)
		  (display r port)
		  (newline port)) 
		rhos rejections))

    ;; beta 100
    (let* ((simulate-rho (lambda (rho) (simulate-rho-beta rho 100)))
	   (rejections (map simulate-rho rhos)))
      (for-each (lambda (rho r) 
		  (display "100 " port)
		  (display rho port)(display " " port)
		  (display r port)
		  (newline port)) 
		rhos rejections))))