<commit_before>(ns cljam.sorter
  (:refer-clojure :exclude [sort sorted?])
  (:require [clojure.java.io :refer [file]]
            [clojure.string :as str]
            (cljam [sam :as sam]
                   [bam :as bam]
                   [common :refer [version]]
                   [util :as util]
                   [io :as io])
            [cljam.util.sam-util :as sam-util]))

(def ^:private chunk-size 1500000)

(def order-unknown :unknown)
(def order-coordinate :coordinate)
(def order-queryname :queryname)

;;
;; coordinate sorter
;;

(defn- replace-header [hdr vn so]
  (conj hdr {:HD {:VN vn, :SO so}}))

(defn- compkey-pos [ref-map b]
  [(get ref-map (:rname b)) (:pos b)])

(defn- sort-alignments-by-pos [rdr]
  (let [ref-map (apply merge
                       (map-indexed (fn [i val] {val i})
                                    (map :name (sam-util/make-refs (io/read-header rdr)))))]
    (sort-by (partial compkey-pos ref-map) (io/read-coordinate-blocks rdr))))

;;
;; queryname sorter
;;

;; (defn- compkey-qname [hdr aln]
;;   [(.indexOf ^List (map :name (sam-util/make-refs hdr)) (:rname aln))
;;    (:qname aln)
;;    (bit-and (:flag aln) 0xc0)])

;; FIXME
(defn- sort-alignments-by-qname [rdr]
  ;; (->> (sort-by (partial compkey-qname (:header sam)) (:alignments sam))
  ;;      (assoc sam :alignments))
  )

;;
;; split/merge
;;

(defn gen-cache-filename
  [prefix i]
  (format "%s/%s_%05d.cache" util/temp-dir prefix i))

(defn gen-sorted-cache-filename
  [prefix i]
  (format "%s/%s_%05d.sorted.cache" util/temp-dir prefix i))

(defn- find-head
  ([keyfn x] x)
  ([keyfn x y] (case (compare (keyfn x) (keyfn y))
                 0 x
                 1 y
                 -1 x))
  ([keyfn x y & more]
     (reduce #(find-head keyfn %1 %2) (find-head keyfn x y) more)))

(defn- head
  [blocks ref-map]
  (apply find-head
         (fn [[idx b]] (compkey-pos ref-map b))
         (filter
          (fn [[_ b]] (not (nil? b)))
          (map-indexed vector blocks))))

(defn split-sam
  [rdr name-fn]
  (doall
   (map
    (fn [[i blks]]
      (let [f (name-fn i)
            hdr (io/read-header rdr)]
        (with-open [wtr (bam/writer f)]
          (io/write-header wtr hdr)
          (io/write-refs wtr hdr)
          (io/write-blocks wtr blks))
        f))
    (map-indexed vector (partition-all chunk-size (io/read-blocks rdr))))))

(defn merge-sam
  [wtr header name-fn n]
  (let [rdrs (map #(bam/reader (name-fn %)) (range n))
        ref-map (apply merge
                       (map-indexed (fn [i val] {val i})
                                    (map :name (sam-util/make-refs header))))]
    (with-open [wtr wtr]
      (io/write-header wtr header)
      (io/write-refs wtr header)
      (loop [blocks-list (map #(io/read-coordinate-blocks %) rdrs)]
        (let [candidates (map first blocks-list)]
          (when-not (every? nil? candidates)
            (let [[i b] (head candidates ref-map)]
              (io/write-blocks wtr [b])
              (recur (map
                      (fn [[idx blocks]]
                        (if (= idx i)
                          (drop 1 blocks)
                          blocks))
                      (map-indexed vector blocks-list))))))))))

(defn clean-split-merge-cache
  [name-fn sorted-name-fn count]
  (let [clean (fn [path]
                (let [f (file path)]
                  (when (.exists f)
                    (.delete f))))]
    (doseq [i (range count)]
      (clean (name-fn i))
      (clean (sorted-name-fn i)))))

;;
;; sorter
;;

(defn sort-by-pos [rdr wtr]
  (let [filename (.getName (file (io/reader-path rdr)))
        basename (first (str/split filename #"\.(?=[^\.]+$)"))
        cache-name-fn (partial gen-cache-filename basename)
        sorted-cache-name-fn (partial gen-sorted-cache-filename basename)
        splited-files (split-sam rdr cache-name-fn)
        num-splited (count splited-files)
        hdr (replace-header (io/read-header rdr)
                            version
                            (name order-coordinate))]<commit_msg>fix sort order<commit_after>(ns cljam.sorter
  (:refer-clojure :exclude [sort sorted?])
  (:require [clojure.java.io :refer [file]]
            [clojure.string :as str]
            (cljam [sam :as sam]
                   [bam :as bam]
                   [common :refer [version]]
                   [util :as util]
                   [io :as io])
            [cljam.util.sam-util :as sam-util]))

(def ^:private chunk-size 1500000)

(def order-unknown :unknown)
(def order-coordinate :coordinate)
(def order-queryname :queryname)

;;
;; coordinate sorter
;;

(defn- replace-header [hdr vn so]
  (conj hdr {:HD {:VN vn, :SO so}}))

(defn- compare-key-pos [ref-map b1 b2]
  (if (= (:rname b1) (:rname b2))
    (compare (:pos b1) (:pos b2))
    (cond
     (= (:rname b1) "*") 1
     (= (:rname b2) "*") -1
     :else (compare (:rname b1) (:rname b2)))))

(defn- sort-alignments-by-pos [rdr]
  (let [ref-map (apply merge
                       (map-indexed (fn [i val] {val i})
                                    (map :name (sam-util/make-refs (io/read-header rdr)))))]
    (sort (compare-key-pos ref-map) (io/read-coordinate-blocks rdr))))

;;
;; queryname sorter
;;

;; (defn- compkey-qname [hdr aln]
;;   [(.indexOf ^List (map :name (sam-util/make-refs hdr)) (:rname aln))
;;    (:qname aln)
;;    (bit-and (:flag aln) 0xc0)])

;; FIXME
(defn- sort-alignments-by-qname [rdr]
  ;; (->> (sort-by (partial compkey-qname (:header sam)) (:alignments sam))
  ;;      (assoc sam :alignments))
  )

;;
;; split/merge
;;

(defn gen-cache-filename
  [prefix i]
  (format "%s/%s_%05d.cache" util/temp-dir prefix i))

(defn gen-sorted-cache-filename
  [prefix i]
  (format "%s/%s_%05d.sorted.cache" util/temp-dir prefix i))

(defn- find-head
  ([keyfn l]
     (reduce (fn [e1 e2]
               (condp = (compare e1 e2)
                 0 e1
                 1 e2
                 -1 e1))
             l)))

(defn- head
  [blocks ref-map]
  (find-head
         (fn [[idx1 b1] [idx2 b2]]
           (compare-key-pos ref-map b1 b2))
         (filter
          (fn [[_ b]] (not (nil? b)))
          (map-indexed vector blocks))))

(defn split-sam
  [rdr name-fn]
  (doall
   (map
    (fn [[i blks]]
      (let [f (name-fn i)
            hdr (io/read-header rdr)]
        (with-open [wtr (bam/writer f)]
          (io/write-header wtr hdr)
          (io/write-refs wtr hdr)
          (io/write-blocks wtr blks))
        f))
    (map-indexed vector (partition-all chunk-size (io/read-blocks rdr))))))

(defn merge-sam
  [wtr header name-fn n]
  (let [rdrs (map #(bam/reader (name-fn %)) (range n))
        ref-map (apply merge
                       (map-indexed (fn [i val] {val i})
                                    (map :name (sam-util/make-refs header))))]
    (with-open [wtr wtr]
      (io/write-header wtr header)
      (io/write-refs wtr header)
      (loop [blocks-list (map #(io/read-coordinate-blocks %) rdrs)]
        (let [candidates (map first blocks-list)]
          (when-not (every? nil? candidates)
            (let [[i b] (head candidates ref-map)]
              (io/write-coordinate-blocks wtr [b])
              (recur (map
                      (fn [[idx blocks]]
                        (if (= idx i)
                          (drop 1 blocks)
                          blocks))
                      (map-indexed vector blocks-list))))))))))

(defn clean-split-merge-cache
  [name-fn sorted-name-fn count]
  (let [clean (fn [path]
                (let [f (file path)]
                  (when (.exists f)
                    (.delete f))))]
    (doseq [i (range count)]
      (clean (name-fn i))
      (clean (sorted-name-fn i)))))

;;
;; sorter
;;

(defn sort-by-pos [rdr wtr]
  (let [filename (.getName (file (io/reader-path rdr)))
        basename (first (str/split filename #"\.(?=[^\.]+$)"))
        cache-name-fn (partial gen-cache-filename basename)
        sorted-cache-name-fn (partial gen-sorted-cache-filename basename)
        splited-files (split-sam rdr cache-name-fn)
        num-splited (count splited-files)
        hdr (replace-header (io/read-header rdr)
                            version
                            (name order-coordinate))]
{'fraction_non_alphanumeric': 0.16330007301046484, 'fraction_numerical': 0.007422730591384765, 'mean_word_length': 1.8118371536093054, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '638824', 'n_tokens_mistral': 2988, 'n_tokens_neox': 2859, 'n_words': 755}