                                            :test (lazy-seq test)}})
             (persist-id-state @ids-inst)))))))
 
-(defn write-dataset-to-excel
-  "Write the data set to an excel file.
+(defn distribution
+  "Return maps representing the data set distribution by class label.  Each
+  element of the returned sequence has the following keys:
+
+  * **:class-label** the class-label fo the instances
+  * **:count** the number of instances for **:class-label**"
+  []
+  (use-connection
+    (with-context [instance-context]
+      (->> (es/buckets es-class-label-key)
+           (map #(rename-keys % {:name class-label-key}))))))
+
+(defn instances-by-class-label
+  "Return a map with class-labels for keys and corresponding instances for that
+  class-label.
 
   Keys
   ----
-
+  * **:max-instances** the maximum number of instances per class
+  * **:type* the data to grab, which is one of the following symbols:
+    * **ids** returns only IDs (not the document data)
+    * **document** returns the document data (`:class-label` and `:instance` keys)
+  * **:seed** if given, seed the random number generator, otherwise don't
+  return random documents"
+  [& {:keys [max-instances type seed]
+      :or {max-instances Integer/MAX_VALUE
+           type 'document}}]
+  (use-connection
+    (with-context [instance-context]
+      (letfn ([query [class-label]
+               (merge {:query
+                       (if seed
+                         {:function_score
+                          {:filter {:term {es-class-label-key class-label}}
+                           :functions [{:random_score
+                                        {:seed seed}}]}}
+                         {:term {es-class-label-key class-label}})}
+                      (if (= type 'ids) {:fields []}))])
+        (->> (distribution)
+             (map class-label-key)
+             (map (fn [class-label]
+                    (->> (query class-label)
+                         es/search
+                         (take max-instances)
+                         (map (case type
+                                document #(-> % :doc dataset-key)
+                                ids #(-> % :id Integer/parseInt)))
+                         (array-map class-label))))
+             (into {})
+             doall)))))
+
+(defn- divide-by-uneven-distribution-set
+  "Divide the dataset into a test and training *buckets*.
+
+  * **train-ratio** this is the percentage of data in the train bucket.
+
+  Keys
+  ----
+  * **:shuffle?** if `true` then shuffle the set before partitioning, otherwise
+  just update the *demarcation* boundary"
+  [train-ratio & {:keys [shuffle? ids
+                         max-instances]
+                  :or {shuffle? true
+                       max-instances Integer/MAX_VALUE}}]
+  (use-connection
+    (let [ids (->> (or ids (db-ids))
+                   ((if shuffle? shuffle identity))
+                   (take max-instances))
+          sz (count ids)
+          train-count (* train-ratio sz)
+          id-data {:train-test {:train (take train-count ids)
+                                :test (drop train-count ids)}}]
+      (reset! ids-inst id-data)
+      (log/infof "divided by set: %s" (stats))
+      (persist-id-state id-data)
+      (stats))))
+
+(defn- divide-by-class-distribution-set
+  "Just like [[divide-by-uneven-distribution-set]] but each test and training set
+  has an even distribution by class label.
+
+  The **keys** parameter(s) are described in [[divide-by-set]]
+  and [[instances-by-class-label]]."
+  [train-ratio & opts]
+  (use-connection
+    (with-context [instance-context]
+      (let [opts (concat opts [:type 'ids])
+            {:keys [shuffle?] :or {shuffle? true}} (apply hash-map opts)]
+        (->> opts
+             (apply instances-by-class-label)
+             vals
+             (reduce (fn [{:keys [train test]} ids]
+                       (let [sz (count ids)
+                             train-count (* train-ratio sz)]
+                         {:train (concat train (take train-count ids))
+                          :test (concat test (drop train-count ids))}))
+                     {})
+             ((if shuffle?
+                (fn [{:keys [test train]}]
+                  {:test (shuffle test)
+                   :train (shuffle train)})
+                identity))
+             (array-map :train-test)
+             (reset! ids-inst)
+             persist-id-state))
+      (let [stats (stats)]
+        (log/infof "divided by set with even distribution by class: %s" stats)
+        stats))))
+
+(defn divide-by-set
+  "Divide the dataset into a test and training *buckets*.
+
+  * **train-ratio** this is the percentage of data in the train bucket, which
+  defaults to `0.5`
+
+  Keys
+  ----
+  * **:dist-type** one of the following symbols:
+      *even* each test/training set has an even distribution by class label
+      *uneven* each test/training set has an uneven distribution by class label
+  * **:shuffle?** if `true` then shuffle the set before partitioning, otherwise
+  just update the *demarcation* boundary
+  * **:max-instances** the maximum number of instances per class
+  * **:seed** if given, seed the random number generator, otherwise don't
+  return random documents"
+  ([]
+   (divide-by-set 0.5))
+  ([train-ratio & {:keys [dist-type shuffle? max-instances seed]
+                    :as opts
+                    :or {shuffle? true
+                         dist-type 'uneven}}]
+   (-> (case dist-type
+         even divide-by-class-distribution-set
+         uneven divide-by-uneven-distribution-set)
+       (apply (apply concat (list train-ratio) opts)))))
+
+(defn dataset-file []
+  (use-connection
+    (->> (format "%s-dataset.xls" index-name)
+         (res/resource-path :analysis-report))))
+
+(defn write-dataset
+  "Write the data set to a spreadsheet.  If the file name ends with a `.csv` a
+  CSV file is written, otherwise an Excel file is written.
+
+  Keys
+  ----
   * **:output-file** where to write the file and defaults to
   [[res/resource-path]] `:analysis-report`
   * **:single?** if `true` then create a single sheet, otherwise the training
   and testing *buckets* are split between sheets"
-  [& {:keys [output-file single?]}]
-  (use-connection
-    (let [output-file (or output-file
-                          (->> (format "%s.xls" index-name)
-                               (res/resource-path :analysis-report)))]
-      (letfn [(data-set [set-type fields col-names header?]
-                (->> (instances :set-type set-type)
-                     (map (fn [{:keys [class-label instance]}]
-                            (concat fields [class-label instance])))
-                     ((if header?
-                        #(cons (concat col-names ["Label" "Instance"]) %)
-                        identity))
-                     ((if header? ss/headerize identity))))]
-        (-> (excel/build-workbook
-             (excel/workbook-hssf)
-             (if single?
-               {"Train and Test" (concat (data-set :train ["train"] ["Train"] true)
-                                         (data-set :train ["test"] ["Test"] false))}
-               {"Train" (data-set :train nil nil true)
-                "Test" (data-set :test nil nil true)}))
-            (ss/autosize-columns)
-            (excel/save output-file))))))
-
-(defn- main
-  "In REPL testing"
-  [& actions]
-  (use-connection
-    (->> actions
-         (map (fn [action]
-                (case action
-                  -2 (set-default-connection)
-                  -1 (clear :wipe-persistent? true)
-                  0 (instances-load)
-                  4 (create-id-list)
-                  7 (id-info)
-                  8 (with-context [stats-context]
-                      (pprint (es/describe)))
-                  9 (set-population-use 0.7)
-                  12 (stats))))
-         doall)))
+  [& {:keys [output-file single? instance-fn columns-fn]
+      :or {instance-fn identity
+           columns-fn (constantly ["Instance"])}}]
+  (use-connection
+    (let [output-file (or output-file (dataset-file))
+          csv? (re-matches #".*\.csv$" (.toString output-file))]
+      (letfn [(data-set [set-type fields col-names header? headerize?]
+                (let [insts (instances :set-type set-type :include-ids? true)
+                      header-inst (first insts)]
+                  (->> insts
+                       (map (fn [{:keys [class-label id instance]}]
+                              (->> (instance-fn instance)
+                                   (concat fields [class-label id]))))
+                       ((if header?
+                          #(cons (concat col-names
+                                         ["Label" "Id"]
+                                         (columns-fn header-inst))
+                                 %)
+                          identity))
+                       ((if headerize? ss/headerize identity)))))]
+        (if csv?
+          (with-open [writer (io/writer output-file)]
+            (->> (concat (data-set :train ["train"] ["Set Type"] true false)
+                         (data-set :train ["test"] nil false false))
+                 (csv/write-csv writer)))
+          (-> (excel/build-workbook
+               (excel/workbook-hssf)
+               (if single?
+                 {"Train and Test"
+                  (concat (data-set :train ["train"] ["Set Type"] true true)
+                          (data-set :train ["test"] ["Test"] false false))}
+                 {"Train" (data-set :train nil nil true true)
+                  "Test" (data-set :test nil nil true true)}))
+              (ss/autosize-columns)
+              (excel/save output-file)))))))

  ([id]
   (use-connection
     (with-context [instance-context]
       (:dataset (es/document-by-id id))))))

(defn clear
  "Clear the in memory instance data.  If key `:wipe-persistent?` is `true` all

                      :url url
                      :settings {"index.mapping.ignore_malformed" true}
                      :mapping-type-defs
                      {dataset-key {:properties mapping-type-def}})
   :stats-context (es/create-context
                   index-name "stats"
                   :url url
