(ns Path_Finder.src.drawing)

(import ij.measure.ResultsTable
        ij.plugin.filter.ParticleAnalyzer
        [ij.gui Roi PolygonRoi]
        ij.plugin.frame.RoiManager)

; ---- Utility ----

(defn make-16bit-img [title colorscheme w h]
  (let [img (ij.IJ/createImage title "16-bit black" w h 1)]
    (.show img)
    (ij.WindowManager/setTempCurrentImage img)
    (ij.IJ/run colorscheme)
    img))

; ---- PathFinder Pipeline ----

(defn label-vertices [vtx]
  (let [rt    (new ResultsTable)
        proc  (.. vtx getProcessor (convertToByte true))
        pa    (ParticleAnalyzer. ParticleAnalyzer/SHOW_ROI_MASKS 0 rt 0.0 200000.0 0.0 1.0)]
    (.invert proc)
    (.setProcessor vtx proc)
    (.. pa (analyze vtx))
    (.. pa getOutputImage)))

(defn qualityscore [path]
  (let [metricdist (.. path (get 0) distance)
        pixeldist  (.. path (get 0) pixelDistance)]
    (if (< 1.0 pixeldist) 
      (/ metricdist pixeldist)
      65536)))


; determines the pixels where paths in path_ids overlap
;   existing paths drawn to proc. Gets points where they 
;   hop on and hop off.
(defn get-overlap-points [pthfdr paths_and_scores proc min_path_len]
  (let [f1 first          ; [path score] → path
        f2 #(.get % 0)    ; get the last PathLink
        f3 #(< min_path_len (.pixelDistance %))
        pix_d (filter (comp f3 f2 f1) paths_and_scores)] ; short paths overlap too frequently
    (mapcat #(into [] (.compuNewVerts pthfdr (first %) proc)) pix_d)))

; ---- DATASTRUCTURES. Build Paths and Get Scores. ----

(defn get-path-edges [vert conn]
  (for [i (range 1 (.size vert))
        j (range i)
        :when (aget conn i j)]
      [i j (.get vert i) (.get vert j)]))

(defn get-pathinfo-from-edges [pthfdr edges mem vor distance_cutoff]
  (for [[i j start end] edges
        :let [endlink (.buildPath pthfdr start end mem vor distance_cutoff)
              path    (.getPath pthfdr endlink)
              score   (qualityscore path)]]
    (do (println i j score)
        [path i j [start end] score])))

; map each vertex to each path it starts OR ends
(defn paths-2-map [mapp [i j id score]]
  (let [mi  (mapp i []) ; list of paths mapped to by i
        mj  (mapp j []) ; ditto
        new_mi (conj mi [j id score])
        new_mj (conj mj [i id score])]
    (assoc mapp i new_mi j new_mj))) ; i ↔ j symmetry

(defn get-vtxmap [path_tuple]
  (let [scores       (map first path_tuple)
        vtx_id_score (map (fn [t id scr] [(t 1) (t 2) id (^number scr)]) path_tuple (range) scores)]
    (reduce paths-2-map {} vtx_id_score)))

(defn get-vtxmap-2 [scores edges]
  (let [vtx_id_score (map (fn [t id scr] [(first t) (second t) id scr]) edges (range) scores)]
    (reduce paths-2-map {} vtx_id_score)))

(defn make-good-path-list [vtxmap degree_cutoff]
  (let [third       (fn [^Number x] (second (rest x)))
        mv-vtx-end  (fn [[vtx tuples]] (map #(conj % vtx) tuples))
        good-score? (fn [x] (> (first x) ((second x) 2)))] ;1st=cutoff 2nd=tuple
    (into #{} ; remove duplicates
      (mapcat ; map across vertex_map
        (comp
          ; [j id score i] → id
          #(map second %)
          ; [scorecutoff tuple] → tuple
          #(map second %)
          ; remove paths with score > cutoff
          #(filter good-score? %)
          ; pair list with path-score cutoff
          #(map list degree_cutoff %)
          ; sort this list by it's path's score
          #(sort-by third %)
          ; add the start-vtx value to end of each list
          mv-vtx-end)
        vtxmap))))

; ---- DRAWING TOOLS ---- 

(defn draw-paths [paths_scores proc]
  (doseq [[path score] paths_scores
          pathlink  path
          :let [pixel (.pixel pathlink)
                x     (aget pixel 0)
                y     (aget pixel 1)
                val   (.get proc x y)]
          :when (or (= 0 val) (> val score))]
    (.set proc x y score)))

(defn redraw-paths [paths_scores proc]
  (doseq [[path score] paths_scores
          pathlink  path
          :let [pixel (.pixel pathlink)
                x     (aget pixel 0)
                y     (aget pixel 1)
                val   (.get proc x y)]]
    (.set proc x y score)))

; draw a finished set of paths to polygon rois (allows scaling)
;   path → xs, ys → PolygonRoi
(defn path-to-poly [path scale]
  (let [xs    (for [pl path :let [px (.pixel pl)]] (aget px 0))
        ys    (for [pl path :let [px (.pixel pl)]] (aget px 1))
        scaled #(map (partial * scale) %)
        x2    (conj (vec (take-nth 3 (take (- (count xs) 1) xs))) (last xs))
        y2    (conj (vec (take-nth 3 (take (- (count ys) 1) ys))) (last ys))
        x3    (int-array (scaled x2))
        y3    (int-array (scaled y2))
        pg    (PolygonRoi. x3 y3 (count x3) Roi/POLYLINE)]
    (.setStrokeWidth pg 2.0)
    (ij.IJ/log (str "I'm working! Stoke: " (.getStrokeWidth pg)))
    (.fitSpline pg)
    pg))
  ; (PolygonRoi. xs ys (count xs) Roi/POLYLINE)))
  ;; type ∈ Roi.POLYGON, Roi.FREEROI, Roi.TRACED_ROI, Roi.POLYLINE, Roi.FREELINE or Roi.ANGLE

; Fully segmented image
;   uses 4-conneced cells to work with 8-connected paths
(defn get-cell-img [path_img]
  (let [seg (.. path_img getProcessor (convertToByte true))
        x   (.threshold seg 0)
        _   (.setProcessor path_img seg)
        ; (exec [ImagePlus Name white-particles? 4-connected?])
        res (.exec (BinaryLabel_.) path_img "Cell Image" true true)]
    (aget res 1)))

; RoiManager, Paths (inplace rm change)
; paths = [path]
(defn add-paths-to-roimanager [rm paths scale]
  (doseq [path paths :let [poly (path-to-poly path scale)]]
    (.addRoi rm poly)))

(defn plot-paths-polyroi [paths imp scale]
  (let [rm (RoiManager.)]
    (add-paths-to-roimanager rm paths scale)
    (.runCommand rm imp "Draw")))

(defn get-scaled-up-img [orig mem paths_and_scores]
    (let [height (.getHeight orig)
          width  (.getWidth orig)
          scale  (/ height (.getHeight mem))
          img    (make-16bit-img "Paths" "blue orange icb" width height)
          good_paths (map first paths_and_scores)]
      (plot-paths-polyroi good_paths img scale)
      img))
