; @Integer(label="Num positive samples") num-positive-samples
; @Integer(label="Num negative samples") num-negative-samples
; @Dataset(label="Vasculature channel") vasculature
; @Dataset(label="Red blood cell channel") rbc
; @OUTPUT File segmentation-file

(ns zebrafish-vasculature.train-segmentation
  (:require [fun.imagej.img :as img]
            [fun.imagej.img.cursor :as cursor]
            [fun.imagej.img.shape :as shape]
            [fun.imagej.img.type :as tpe]
            [fun.imagej.core :as ij]
            [fun.imagej.ops :as ops]
            [fun.imagej.img.utils :as img-utils]
            [fun.imagej.mesh :as msh]
            [fun.imagej.segmentation.fast-segmentation :as fseg]
            [clojure.java.io :as io]
            [clojure.string :as string]
            [zebrafish-vasculature.training-utils :as tutils]))

(defn -main [& args]
  (let [;; First put everything into a map
        argmap (apply hash-map
                      (mapcat #(vector (read-string (first %)) (second %) #_(read-string (second %)))
                              (partition 2 args)))
        ;; Then read-string on *some* args, but ignore others
        argmap (apply hash-map
                      (apply concat
                             (for [[k v] argmap]
                               [k (cond (= k :output-directory) v
                                        (= k :cache-directory) v
                                        (= k :directory) v
                                        (= k :basename) v
                                        (= k :rbc-filename) v
                                        (= k :vasculature-filename) v
                                        :else (read-string v))])))
        arg-params (merge {:num-positive-samples 1000
                           :num-negative-samples 1000
                           :label-rbc true
                           :generate-dataset true
                           :solve-segmentation true
                           :save-segmentation-config true
                           ;:cache-directory "/projects/VirtualFish/kyle/"
                           ;:cache-directory "/Users/kharrington/Data/Daetwyler_Stephan/test_ISVs/cache/"
                           :segmentation-config-filename "segmentation_config.clj"
                           :write-segmentation true
                           :segmentation-filename "segmentationMap.tif"
                           ;:basename "161122_angle001_t0188_registered"
                           ;:basename "t00161_angle002_crop_001"
                           ;:directory "/Users/kharrington/Data/Daetwyler_Stephan/consecutiveData/t00161/"
                           :directory "/Users/kharrington/Data/Daetwyler_Stephan/movie_stack/"
                           :basename "head_vasculature_crop_003_t0"
                           :verbose true
                           ; consider adding a tag
                           }; default params
                          argmap)
        basename (str (:directory arg-params)
                      (:basename arg-params))
        weight-filename (str basename "_weights.csv")]
    (println args)
    (println argmap)
    (println arg-params)
    (println (:basename arg-params) (:directory arg-params) basename)

    (def full-image (fun.imagej.ops.convert/float32 (ij/open-img (str basename ".tif"))))
    (def vasculature (img/normalize (img/hyperslice full-image 2 0)))
    (def rbc (img/normalize (img/hyperslice full-image 2 1)))
    
    (def seg (atom (fseg/create-segmentation-meta (if (:load-segmentation-config arg-params)
                                                    (read-string (slurp (str basename (:segmentation-config-filename arg-params))))
                                                    {:num-positive-samples (:num-positive-samples arg-params)
                                                     :num-negative-samples (:num-negative-samples arg-params)
                                                     :basename (:basename arg-params)
                                                     :verbose (:verbose arg-params)
                                                     :cache-directory (if (:no-caching arg-params)
                                                                        nil
                                                                        (:cache-directory arg-params))
                                                     :segmentation-type :3D}))))

    (println "Adding feature maps")
    (reset! seg (tutils/setup-default-feature-maps @seg))

    (when (:label-rbc arg-params)
      (println "Generating sample points.")
      (reset! seg (if (:rbc-presegmented arg-params)
                    (fseg/generate-sample-points @seg rbc)
                    (let [; try cleaning by taking median along z

                          histogram (fun.imagej.ops.image/histogram rbc)
                          thresh-val (first (fun.imagej.ops.threshold/intermodes histogram))
                          rbc-label (fun.imagej.ops.threshold/apply rbc thresh-val)]
                      (fseg/generate-sample-points @seg rbc-label)))))

    (when (:generate-dataset arg-params)
      (println "Generating dataset.")
      (reset! seg (fseg/generate-dataset @seg vasculature)))

    (when (:solve-segmentation arg-params)
      (println "Solving segmentation")
      (reset! seg (fseg/solve-segmentation @seg)))

    (println "Weights:")
    (println (:weights @seg))
    (spit (str directory
               basename
               "_weights.csv")
          (string/join "," (:weights @seg)))

    (when (:save-segmentation-config arg-params)
      (println "Saving segmentation config")
      (fseg/save-segmentation-config @seg (:segmentation-config-filename arg-params)))))

;(-main)