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

  ([id instance class-label set-type]
   (log/infof "loading instance (%s): %s => <%s>"
               id class-label
               (let [s (pr-str instance)]
                 (subs s 0 (min 80 (count s)))))
   (log/debugf "instance: %s" instance)
   (use-connection
     (with-context [instance-context]
       (let [doc (merge {:dataset {:class-label class-label
                                   :instance instance}}
                        (if set-type {:set-type set-type}))
             res (if id
                   (es/put-document id doc)

                                           :test (lazy-seq test)}})
            (persist-id-state @ids-inst)))))))

(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? 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)))))))
