(define-module (gn runner gemma)
  #:use-module (json)
  #:use-module (gn data genotype)
  #:use-module (gn util convert)
  #:use-module (ice-9 match)
  #:use-module (ice-9 format)
  #:use-module (ice-9 iconv)
  #:use-module (ice-9 receive)
  #:use-module (ice-9 string-fun)
  #:use-module (rnrs base)

  #:export (
            write-pheno-file
            invoke-gemma-wrapper-loco
            run-gemma
            ))

(define (write-pheno-file fn traits)
  (define bxd-inds (geno-inds-bxd "BXD.json"))
  (assert (= 235 (length bxd-inds)))
  (display bxd-inds)
  (call-with-output-file fn
    (lambda (port)
      (for-each
       (lambda (ind)
         (begin
           (let* [(value (assoc-ref traits ind))
                  (outvalue (if value
                                value
                                "NA"))]
             (if value
                 (begin
                   (format #t "~s ~s" ind outvalue)
                   (newline)))
             (display outvalue port)
             (newline port))))
       bxd-inds)
      (close port)
      ))
)

(define (invoke-gemma-wrapper-loco name trait-name pheno-fn)
  "Create a tmpdir and invoke gemma-wrapper using parallel LOCO. Note that at this point we use a number of defaults for BXD"
  (let* [(population "BXD")
         (sys-tmpdir (getenv "TMPDIR"))
         (tmpdir (mkdtemp (string-append sys-tmpdir "/run-gemma-XXXXXX")))
         (k-json-fn (string-append tmpdir "/K.json"))
         (gwa-json-fn (string-append tmpdir "/GWA.json"))]
    ;; --- First we compute K - control output goes to K.json
    (let [(err (system (string-append "/gemma-wrapper/bin/gemma-wrapper --verbose --population \"" population "\" --name \"" name "\" --trait \"" trait-name "\" --verbose --loco --json --parallel -- -gk -g BXD.8_geno.txt.gz -p " pheno-fn " -a BXD.8_snps.txt > " k-json-fn  )))]
      (if (not (= err 0))
          (exit err)))
    (let [(err (system (string-append "/gemma-wrapper/bin/gemma-wrapper --population \"" population "\" --name \"" name "\" --id \"" trait-name "\" --trait \"" trait-name "\" --verbose --loco --json --input " k-json-fn " -- -g BXD.8_geno.txt.gz -p " pheno-fn " -a BXD.8_snps.txt -lmm 9 -maf 0.1 > " gwa-json-fn)))]
      (if (not (= err 0))
          (exit err)))
    ;; (delete-file pheno-fn)
    ;; (delete-file gwa-json-fn)
    ;; (delete-file k-json-fn)
    ;; (rmdir tmpdir)
    )
  )

(define (run-gemma population data-id probeset-id probesetfreeze-id name trait-name traits)
  "Run gemma-wrapper to compute GRM and GWA. On failure the run will stop(!)"
  (define ids (string-append (int-to-string data-id) ","
                             (int-to-string probeset-id) ","
                             (int-to-string probesetfreeze-id)))
  (if name
      (display (string-append "WE HAVE OUR " population " DATASET " name " and trait " trait-name " for precompute!\n")))
  (display ids)
  (display traits)
  (newline)
  (let* [(sys-tmpdir (getenv "TMPDIR"))
         (tmpdir (mkdtemp (string-append sys-tmpdir "/run-gemma-XXXXXX")))
         (pheno-fn (string-append tmpdir "/pheno.txt"))
         (k-json-fn (string-append tmpdir "/K.json"))
         (gwa-json-fn (string-append tmpdir "/GWA.json"))]
    ;; ---- to start GEMMA precompute inside container
    ;; env TMPDIR=. LD_LIBRARY_PATH=$GUIX_ENVIRONMENT/lib/ guile -L . -s ./scripts/precompute/precompute-hits.scm
    ;; --- First we compute K - control output goes to K.json
    (let [(err (system (string-append "env GEMMA_COMMAND=gemma /gemma-wrapper/bin/gemma-wrapper --population \"" population "\" --name \"" name "\" --trait \"" trait-name "\" --verbose --loco --json --parallel -- -gk -g BXD.8_geno.txt.gz -p " pheno-fn " -a BXD.8_snps.txt > " k-json-fn  )))]
      (if (not (= err 0))
          (exit err)))
    (let [(err (system (string-append "env GEMMA_COMMAND=gemma /gemma-wrapper/bin/gemma-wrapper --keep --population \"" population "\" --name \"" name "\" --id \"" ids "\" --trait \"" trait-name "\" --verbose --loco --json --parallel --input " k-json-fn " -- -g BXD.8_geno.txt.gz -p " pheno-fn " -a BXD.8_snps.txt -lmm 9 -maf 0.1 > " gwa-json-fn)))]
      (if (not (= err 0))
          (exit err)))
    (delete-file pheno-fn)
    (delete-file gwa-json-fn)
    (delete-file k-json-fn)
    (rmdir tmpdir)
    )
  #t
  )