; @Dataset(label="Vasculature channel") vasculature
; @Dataset(label="Red blood cell channel") rbc
; @File(label="Trained segmentation file") segmentation-file
; @OUTPUT Dataset segmentation-map

(ns zebrafish-vasculature.segment-image
  (: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
                                        :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")]
    (doseq [[k v] arg-params]
      (params/set-param k v))
    (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))

    (swap! seg assoc :weights (map read-string (string/split (slurp weight-filename) #",")))

    (println "Weights:")
    (println (:weights @seg))

    (when (:write-segmentation arg-params)
      (println "Writing segmentation map")
      (let [segmentation (fseg/segment-image @seg vasculature)]
        (ij/save-img segmentation
                     (str basename (params/get-param :segmentation-filename)))
        (let [background (fun.imagej.ops.filter/gauss (fun.imagej.ops.create/img segmentation)
                                                      segmentation
                                                      (double-array [25 25 0]))
              to-threshold (fun.imagej.ops.math/subtract segmentation background)
              thresholded (fun.imagej.ops.morphology/dilate
                            (fun.imagej.ops.morphology/erode (fun.imagej.ops.threshold/otsu to-threshold)
                                                             (shape/sphere 1))
                            (shape/sphere 2))]
          (ij/save-img (fun.imagej.ops.convert/uint8 thresholded)
                       (str basename "_segmentation_thresholded.tif"))
          (let [mesh (msh/marching-cubes thresholded)]
            (println "Number of facets in mesh:" (.size (.getFacets mesh)))
            (msh/write-mesh-as-stl mesh (str basename ".stl"))))))))

;(-main)