<commit_before>(ns bluebell.utils.ebmd
  "Example-based overloading"
  (:require [clojure.spec.alpha :as spec]
            [bluebell.utils.wip.specutils :as specutils]
            [clojure.core :as c]
            [bluebell.utils.wip.pareto :as pareto]
            [bluebell.utils.wip.core :as utils]
            [bluebell.utils.wip.check :refer [check-io]]
            [clojure.set :as cljset]
            [bluebell.utils.render-text :as render-text]))

(declare filter-positive)
(declare any-arg)
(declare arg-spec?)
(declare check-valid-arg-spec)
(declare render-overload-text)
(declare re-resolve-arg-spec)
(declare arg-spec-key)


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;;  Specs
;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(spec/def ::key (spec/or :complex (spec/cat :prefix keyword?
                                            :data (spec/* any?))
                         :simple keyword?))

(spec/def ::pos coll?)
(spec/def ::neg coll?)
(spec/def ::pred fn?)
(spec/def ::desc string?)
(spec/def ::spec #(or (spec/spec? %)
                        (keyword? %)))
(spec/def ::valid? boolean?)


(spec/def ::arg-spec (spec/keys :req-un [::pred
                                         ::pos
                                         ::neg
                                         ::key
                                         ::valid?
                                         ::desc]
                                :opt-un [::spec]))

(spec/def ::general-arg-spec (spec/or :key ::key
                                      :arg-spec ::arg-spec))

(spec/def ::input-arg-spec
  (spec/keys :opt-un [::pred ::pos ::neg
                      ::key ::desc ::spec]))



(spec/def ::arg-specs (spec/coll-of ::general-arg-spec))
(spec/def ::fn fn?)
(spec/def ::joint ::general-arg-spec)
(spec/def ::overload (spec/keys :req-un [::arg-specs ::fn]
                                :opt-un [::joint]))

(spec/def ::arg-binding (spec/cat :arg-spec any?
                                  :binding any?))

(spec/def ::joint-binding (spec/cat :prefix #{:joint}
                                    :arg-spec any?))

(spec/def ::def-poly-arg-list
  (spec/* (spec/alt :joint ::joint-binding
                    :arg-binding ::arg-binding)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;;  Impl
;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(def key? (partial spec/valid? ::key))

(def init-reg {::reg-counter 0})

(defonce arg-spec-registry (atom init-reg))

(defn reset-registry! []
  (reset! arg-spec-registry init-reg))

;; (reset-registry!)


(defn register-arg-spec
  ([k arg-spec]
   {:pre [(spec/valid? ::key k)
          (spec/valid? ::general-arg-spec arg-spec)]}
   (swap! arg-spec-registry
          (fn [reg]
            (-> reg
                (assoc k arg-spec)
                (update ::reg-counter inc))))
   arg-spec)
  ([arg-spec]
   {:pre [(spec/valid? ::arg-spec arg-spec)]}
   (register-arg-spec (:key arg-spec) arg-spec)))

(defn get-reg-counter []
  (-> arg-spec-registry
      deref
      ::reg-counter))

(defn resolve-arg-spec [init]
  (loop [x init]
    (if (arg-spec? x)
      x
      (if-let [ag (get (deref arg-spec-registry) x)]
        (recur ag)
        (throw (ex-info "No arg-spec with key"
                        {:key x
                         :init init}))))))

;;; Warning: Make sure that the arg specs don't differentiate
;;; between vectors and seqs, because internally, we store the
;;; samples in a set, meaning that vectors and seqs will be
;;; collapsed to a single element, which is possibly a seq.
;;; In short, use sequential? and avoid
;;; vector? or seq? to test args.

(defn- decorate-key-and-pred-from-spec [input-arg-spec]
  (if-let [s (:spec input-arg-spec)]
    (merge {:key [::spec s]
            :pred (specutils/pred s)} input-arg-spec)
    input-arg-spec))

(defn- validate-on-samples [arg-spec]
  (let [pred (:pred arg-spec)
        pos-failures (filter (complement pred)
                             (:pos arg-spec))
        neg-failures (filter pred (:neg arg-spec))]
    (if (and (empty? pos-failures)
               (empty? neg-failures))
      (assoc arg-spec :valid? true)
      (merge arg-spec
             {:valid? false
              :pos-failures pos-failures
              :neg-failures neg-failures}))))

(defn- decorate-desc [arg-spec]
  {:pre [(contains? arg-spec :key)]}
  (merge {:desc (str (:key arg-spec))} arg-spec))


(defn- init-overload-state [overload-sym init-samples]
  (check-io
   [:pre [(symbol? overload-sym)]]
   
   {:name overload-sym
    :init-samples init-samples
    :dirty? true
    :samples init-samples
    :arg-specs {}
    :overloads {}}))

(defn- state-arg-specs [state]
  (-> state
      :arg-specs

      ;; TODO: This is now a set!
      
      vals
      set))

(defn- mark-dirty [x]
  (assoc x :dirty? true))

(defn- unmark-dirty [x]
  (assoc x :dirty? false))

(defn- add-arg-spec [state arg-spec]
  (check-io
   [:pre [(map? state)
          ::general-arg-spec arg-spec]]
   (let [arg-spec (resolve-arg-spec arg-spec)]
     (-> state
         mark-dirty
         (update :arg-specs conj [(arg-spec-key arg-spec) nil])))))

(defn- add-arg-specs [state arg-specs]
  (reduce add-arg-spec state arg-specs))

(defn- extend-arg-specs-with-joint [overload]
  (conj (vec (:arg-specs overload))
        (or (:joint overload)
            any-arg)))

(defn- add-overload [state overload]
  (check-io
   [:pre [(map? state)
          ::overload overload]]
   (let [arg-specs (extend-arg-specs-with-joint overload)]
     (-> state
         mark-dirty
         (add-arg-specs arg-specs)
         (assoc-in [:overloads (count arg-specs)
                    (mapv arg-spec-key arg-specs)]
                   overload)))))

(defn- rebuild-arg-spec-samples [state]
  (let [samples (:samples state)]
    (assert (set? samples))
    (update state
            :arg-specs
            (fn [arg-specs]
              (transduce<commit_msg>Simplify ebmd<commit_after>(ns bluebell.utils.ebmd
  "Example-based overloading"
  (:require [clojure.spec.alpha :as spec]
            [bluebell.utils.wip.specutils :as specutils]
            [clojure.core :as c]
            [bluebell.utils.wip.pareto :as pareto]
            [bluebell.utils.wip.core :as utils]
            [bluebell.utils.wip.check :refer [check-io]]
            [clojure.set :as cljset]
            [bluebell.utils.render-text :as render-text]))

(declare filter-positive)
(declare any-arg)
(declare arg-spec?)
(declare check-valid-arg-spec)
(declare render-overload-text)
(declare re-resolve-arg-spec)
(declare arg-spec-key)


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;;  Specs
;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(spec/def ::key (spec/or :complex (spec/cat :prefix keyword?
                                            :data (spec/* any?))
                         :simple keyword?))

(spec/def ::pos coll?)
(spec/def ::neg coll?)
(spec/def ::pred fn?)
(spec/def ::desc string?)
(spec/def ::spec #(or (spec/spec? %)
                        (keyword? %)))
(spec/def ::valid? boolean?)


(spec/def ::arg-spec (spec/keys :req-un [::pred
                                         ::pos
                                         ::neg
                                         ::key
                                         ::valid?
                                         ::desc]
                                :opt-un [::spec]))

(spec/def ::general-arg-spec (spec/or :key ::key
                                      :arg-spec ::arg-spec))

(spec/def ::input-arg-spec
  (spec/keys :opt-un [::pred ::pos ::neg
                      ::key ::desc ::spec]))



(spec/def ::arg-specs (spec/coll-of ::general-arg-spec))
(spec/def ::fn fn?)
(spec/def ::joint ::general-arg-spec)
(spec/def ::overload (spec/keys :req-un [::arg-specs ::fn]
                                :opt-un [::joint]))

(spec/def ::arg-binding (spec/cat :arg-spec any?
                                  :binding any?))

(spec/def ::joint-binding (spec/cat :prefix #{:joint}
                                    :arg-spec any?))

(spec/def ::def-poly-arg-list
  (spec/* (spec/alt :joint ::joint-binding
                    :arg-binding ::arg-binding)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;;  Impl
;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(def key? (partial spec/valid? ::key))

(def init-reg {::reg-counter 0})

(defonce arg-spec-registry (atom init-reg))

(defn reset-registry! []
  (reset! arg-spec-registry init-reg))

;; (reset-registry!)


(defn register-arg-spec
  ([k arg-spec]
   {:pre [(spec/valid? ::key k)
          (spec/valid? ::general-arg-spec arg-spec)]}
   (swap! arg-spec-registry
          (fn [reg]
            (-> reg
                (assoc k arg-spec)
                (update ::reg-counter inc))))
   arg-spec)
  ([arg-spec]
   {:pre [(spec/valid? ::arg-spec arg-spec)]}
   (register-arg-spec (:key arg-spec) arg-spec)))

(defn get-reg-counter []
  (-> arg-spec-registry
      deref
      ::reg-counter))

(defn resolve-arg-spec [init]
  (loop [x init]
    (if (arg-spec? x)
      x
      (if-let [ag (get (deref arg-spec-registry) x)]
        (recur ag)
        (throw (ex-info "No arg-spec with key"
                        {:key x
                         :init init}))))))

;;; Warning: Make sure that the arg specs don't differentiate
;;; between vectors and seqs, because internally, we store the
;;; samples in a set, meaning that vectors and seqs will be
;;; collapsed to a single element, which is possibly a seq.
;;; In short, use sequential? and avoid
;;; vector? or seq? to test args.

(defn- decorate-key-and-pred-from-spec [input-arg-spec]
  (if-let [s (:spec input-arg-spec)]
    (merge {:key [::spec s]
            :pred (specutils/pred s)} input-arg-spec)
    input-arg-spec))

(defn- validate-on-samples [arg-spec]
  (let [pred (:pred arg-spec)
        pos-failures (filter (complement pred)
                             (:pos arg-spec))
        neg-failures (filter pred (:neg arg-spec))]
    (if (and (empty? pos-failures)
               (empty? neg-failures))
      (assoc arg-spec :valid? true)
      (merge arg-spec
             {:valid? false
              :pos-failures pos-failures
              :neg-failures neg-failures}))))

(defn- decorate-desc [arg-spec]
  {:pre [(contains? arg-spec :key)]}
  (merge {:desc (str (:key arg-spec))} arg-spec))


(defn- init-overload-state [overload-sym init-samples]
  (check-io
   [:pre [(symbol? overload-sym)]]
   
   {:name overload-sym
    :init-samples init-samples
    :dirty? true
    :samples init-samples
    :arg-specs {}
    :overloads {}}))

(defn- state-arg-specs [state]
  (-> state
      :arg-specs

      ;; TODO: This is now a set!
      
      vals
      set))

(defn- mark-dirty [x]
  (assoc x :dirty? true))

(defn- unmark-dirty [x]
  (assoc x :dirty? false))

(defn- add-arg-spec [state arg-spec]
  (check-io
   [:pre [(map? state)
          ::general-arg-spec arg-spec]]
   (-> state
       mark-dirty
       (update :arg-specs conj [(arg-spec-key arg-spec) nil]))))

(defn- add-arg-specs [state arg-specs]
  (reduce add-arg-spec state arg-specs))

(defn- extend-arg-specs-with-joint [overload]
  (conj (vec (:arg-specs overload))
        (or (:joint overload)
            any-arg)))

(defn- add-overload [state overload]
  (check-io
   [:pre [(map? state)
          ::overload overload]]
   (let [arg-specs (extend-arg-specs-with-joint overload)]
     (-> state
         mark-dirty
         (add-arg-specs arg-specs)
         (assoc-in [:overloads (count arg-specs)
                    (mapv arg-spec-key arg-specs)]
                   overload)))))

(defn- rebuild-arg-spec-samples [state]
  (let [samples (:samples state)]
    (assert (set? samples))
    (update state
            :arg-specs
            (fn [arg-specs]
              (transduce
{'fraction_non_alphanumeric': 0.19943304985826246, 'fraction_numerical': 0.00016675004168751042, 'mean_word_length': 1.9994998749687423, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 30, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2238811', 'n_tokens_mistral': 4179, 'n_tokens_neox': 3911, 'n_words': 944}