(ns com.manigfeald.surmise
  (:refer-clojure :exclude [int vector list hash-map map keyword
                            char boolean byte bytes sequence
                            shuffle not-empty symbol namespace
                            set sorted-set uuid double let short])
  (:require [clojure.test.check.generators :as g]
            [clojure.core :as cc]
            [clojure.test.check.rose-tree :as rose]
            [clojure.test.check.random :as rand]
            [clojure.test.check.properties :as props]))

(set! *warn-on-reflection* true)
;; (set! *unchecked-math* :warn-on-boxed)

(declare trace rose-tree-1 regenerate)

;; Monad combinators
;;
;; test.check is built on two monads, one for generators and one for
;; rose trees. this adds a 3rd alternative monad for creating
;; generators, which are valid test.check generators, but also contain
;; free monad values that describe the generator. free monad just
;; because I wasn't sure what I was going to end up with.
(defn pure [value]
  (assoc (g/return value) ::generator [:pure value]))

(defn make-gen
  "Given the free structure of a surmise generator, make a test.check generator"
  [g]
  (assoc (#'g/make-gen
          (fn [rnd size]
            (cc/let [[l v] (trace g rnd size)]
              (rose-tree-1 v l g size))))
         ::generator g))

(defn fmap [f m]
  (assert (fn? f))
  (assert (::generator m) m)
  (make-gen [:fmap f (::generator m)]))

(defn join [m]
  (assert (::generator m) m)
  (make-gen [:join (::generator m)]))

(defn bind [m f]
  (assert (::generator m))
  (join (fmap f m)))

(defn error [msg]
  (assoc (#'g/make-gen
          (fn [_rnd _size]
            (throw (ex-info msg {}))))
         ::generator [:error msg]))

(def large-integer
  "Generate a long between Long/MIN_VALUE and Long/MAX_VALUE inclusive, shrinks towards 0"
  (assoc g/large-integer ::generator [:draw]))

(defn sized [f]
  (assoc (g/sized f) ::generator [:get-size f]))

;; Generators
;;
;; mostly built on choose, trying to follow the names used in
;; test.check, but shrinking behavior may be a little different

(def ^Number n (-' Long/MAX_VALUE Long/MIN_VALUE))


(declare choose')

(defn choose-within-range-centered-on-zero ^long [^long lower ^long upper ^long i]
  (cc/let [scale (/ (-' upper lower) n)]
    (long (Math/round (cc/double (+' (*' (-' (long i) Long/MIN_VALUE) scale) lower))))))

(defn choose-within-offside-range ^long [^long lower ^long upper ^long n ^long size]
  (cc/let [offset (-' lower)
           olower (+' lower offset)
           oupper (+' upper offset)
           width (-' oupper olower)
           slide (long (Math/round (/ width 2.0)))
           nlower (-' olower slide)
           nupper (-' oupper slide)]
    (cc/let [i (choose' nlower nupper n size)]
      ;; zig zag encoding folds the negative side of the number
      ;; line into the positive, so the middle (0) that we
      ;; shrinking towards becomes the lower end
      (long
       (-
        (bit-xor (bit-shift-right i 63)
                 (bit-shift-left i 1))
        offset)))))

(defn choose-with-range-with-elements-inside ^long [^long lower ^long upper ^long n ^long size]
  (if (> upper 0 lower)
    (choose-within-range-centered-on-zero (max lower (- size)) (min upper size) n)
    (choose-within-offside-range lower upper n size)))

(defn choose-within-range-with-distinct-bounds ^long [^long lower ^long upper ^long n ^long size]
  (cc/let [diff (-' upper lower)]
    (if (= diff 1)
      (if (bit-test n 1) upper lower)
      (choose-with-range-with-elements-inside lower upper n size))))

;; #:com.manigfeald.surmise{:l [-9125064582277025061 -7798280036296055815], :s 3}
(defn choose' ^long [^long lower ^long upper ^long n ^long size]
  (if (= lower upper)
    lower
    (choose-within-range-with-distinct-bounds lower upper n size)))

(defn choose
  "Generates a long between lower and upper inclusive. If 0 is between
  lower and upper shrinks towards 0, otherwise thrinks towards lower"
  [^long lower ^long upper]
  (sized
   (fn [s]
     (fmap
      #(choose' lower upper % s)
      large-integer))))

(def byte
  "Generates a byte between Byte/MIN_VALUE and Byte/MAX_VALUE inclusive, shrinks towards 0"
  (fmap cc/byte (choose Byte/MIN_VALUE Byte/MAX_VALUE)))

(def short
  "Generates a short between Short/MIN_VALUE and Short/MAX_VALUE inclusive, shrinks towards 0"
  (fmap cc/short (choose Short/MIN_VALUE Short/MAX_VALUE)))

(def int
  "Generates a int between Integer/MIN_VALUE and Integer/MAX_VALUE inclusive, shrinks towards 0"
  (fmap cc/int (choose Integer/MIN_VALUE Integer/MAX_VALUE)))

(def boolean
  "Generates a boolean, shrinks to false"
  (fmap cc/boolean (choose 0 1)))

(def char
  "Generates random characters, shrinks towards \0"
  (fmap (fn [x] (cc/char (bit-and x 0xffff))) (choose 0 Short/MAX_VALUE)))

;; shrinking may be nuts
;; (def float (fmap (fn [i] (Float/intBitsToFloat i)) int))

(def double (fmap (fn [i] (Double/longBitsToDouble i)) large-integer))

(defn such-that [g p]
  ((fn such-that* [g p i]
     (if (> i 100)
       (error "such that failed")
       (bind g (fn [v] (if (p v) (pure v) (such-that* g p (inc i)))))))
   g
   p
   0))

(defn one-of [generators]
  (bind (choose 0 (dec (count generators)))
        (fn [i]
          (nth generators i))))

(def small-integer (choose Long/MIN_VALUE Long/MAX_VALUE))

(defn tuple
  ([]
   (pure []))
  ([a]
   (bind a (fn [v] (pure [v]))))
  ([a b]
   (bind a (fn [av] (bind b (fn [bv] (pure [av bv]))))))
  ([a b c & more]
   (bind (tuple a b)
         (fn [abt]
           (bind (apply tuple c more)
                 (fn [morev] (pure (into abt morev))))))))

(def nat (choose 0 Long/MAX_VALUE))

(defn vector [g n]
  (apply tuple (repeat n g)))

(defn set [g]
  (bind nat
        (fn [f]
          (fmap (partial into #{})
                (apply tuple (repeat f g))))))

(defn list [g]
  (bind nat
        (fn [n]
          (fmap
           (partial apply cc/list)
           (apply tuple (repeat n g))))))

(defn bytes [g]
  (fmap (fn [bs] (into-array Byte/TYPE bs))
        (list byte)))

(def uuid (fmap
           (fn [[least most]]
             (java.util.UUID. most least))
           (tuple large-integer large-integer)))

(def ratio (fmap
            (fn [[a b]] (/ a (inc b)))
            (tuple small-integer nat)))

(defn frequency [ps]
  (bind (choose 0 (apply + -1 (cc/map first ps)))
        (fn [i]
          (reduce
           (fn [a [p g]]
             (cc/let [a (+ a p)]
               (if (> a i)
                 (reduced g)
                 a)))
           0
           ps))))

(defmacro for-all
  "Creates a property for checking, like test.check's for-all. Unlike
  test.check's for-all binding is sequential."
  [bindings & body]
  (cc/let [bindings (partition-all 2 bindings)]
    `(props/for-all*
      ~[(reduce
         (fn [a [n g]]
           `(bind ~g (fn [y#] (fmap (fn [x#] (conj x# y#)) (cc/let [~n y#] ~a) ))))
         `(pure ())
         (reverse bindings))]
      (fn [[~@(cc/map first bindings)]]
        ~@body))))

(defmacro let
  [bindings & body]
  (cc/let [bindings (partition-all 2 bindings)]
    (reduce
     (fn [a [n g]]
       `(bind ~g (fn [~n] ~a)))
     `(cc/let [r# (do ~@body)]
        (if (g/generator? r#)
          r#
          (pure r#)))
     (reverse bindings))))

(defn elements [coll]
  (assert (seq coll))
  (cc/let [v (vec (seq coll))]
    (fmap
     #(nth v %)
     (choose 0 (dec (count v))))))

;; Interpreters
;;
;; Two interpreters that run the generator constructing free
;; monad. Trace runs the generator and returns a pair of a vector of
;; all the random longs used while generating and the value
;; generated. Regenerate takes a vector of longs and a generator and
;; generates a value using the supplied longs, returns [:just
;; the-value] or [:error msg].
(defn interpop [{:keys [l c k r size] :as m}]
  (case (first c)
    :fmap (cc/let [[_ f x] c]
            (assoc m
                   :c x
                   :k (fn [m']
                        (assoc m'
                               :c [:just (f (-> m' :c second))]
                               :k k))))
    :error (assoc m :halted true)
    :join (cc/let [[_ x] c]
            (assoc m
                   :c x
                   :k (fn [m']
                        (assoc m'
                               :c (-> m' :c second ::generator)
                               :k k))))
    :just (if k
            (k m)
            (assoc m :halted true))
    :draw (if-not (seq r)
            (assoc m
                   :halted true
                   :c [:error "no more random"])
            (assoc m
                   :l (conj l (first r))
                   :c [:just (first r)]
                   :r (rest r)))
    :pure (cc/let [[_ v] c]
            (assoc m
                   :c [:just v]))
    :get-size (cc/let [[_ f] c]
                (assoc m :c (::generator (f size))))))

(defn trace
  "Tracing interpreter, collects used random values"
  ([op rnd size]
   (cc/let [[rnd' rnd''] (rand/split rnd)]
     (loop [rnd rnd'
            m {:l (vector-of :long)
               :c op
               :k nil
               :r (cc/list (rand/rand-long rnd''))
               :size size}]
       (cc/let [m' (interpop m)]
         (cond (and (:halted m') (= :just (first (:c m'))))
               [(:l m') (second (:c m'))]
               (and (:halted m') (= :error (first (:c m'))))
               (throw (ex-info (second (:c m')) {}))
               (not (seq (:r m')))
               (cc/let [[rnd' rnd''] (rand/split rnd)]
                 (recur rnd' (assoc m' :r (cc/list (rand/rand-long rnd'')))))
               :else
               (recur rnd m')))))))

(defn regenerate
  "Execute using already generated list of random values"
  ([l op size]
   (assert op)
   (loop [m {:l []
             :c op
             :k nil
             :r l
             :size size}]
     (cc/let [m' (interpop m)]
       (cond (:halted m')
             (:c m')
             :else
             (recur m'))))))

;; Shrinking
;;
;; basically given a vector of longs, create a rose tree where the
;; children shrink each long towards 0, and shrink the vector as a
;; whole towards being empty. the exact stategies of shrinking the
;; longs toward 0 should be informed by how they are parsed by key
;; generators like choose.

;; (defprotocol ShrinkPass
;;   (init [s])
;;   (applicable? [s state value])
;;   (app [s state value])
;;   (next-state [s state])
;;   (cost [_]))

;; (deftype HalfShrink []
;;   ShrinkPass
;;   (init [_s] true)
;;   (applicable? [_s state value]
;;     (and state (or (> -1000 value) (> value 1000))))
;;   (app [_s _state value]
;;     (unchecked-subtract (long value) (quot (long value) 2)))
;;   (next-state [_s _state] false)
;;   (cost [_] 0.25))

;; (deftype AddShrink []
;;   ShrinkPass
;;   (init [_s] 1000)
;;   (applicable? [_s state value]
;;     (and (not (zero? state)) (neg? value)))
;;   (app [_s state value]
;;     (unchecked-add value state))
;;   (next-state [_s state]
;;     (dec state))
;;   (cost [_] 1.0))

;; (deftype SubtractShrink []
;;   ShrinkPass
;;   (init [_s] 1000)
;;   (applicable? [_s state value]
;;     (and (not (zero? state)) (pos? value)))
;;   (app [_s state value]
;;     (unchecked-subtract value state))
;;   (next-state [_s state]
;;     (dec state))
;;   (cost [_] 1.0))

;; (def shrink-passes
;;   [(->HalfShrink)
;;    (->AddShrink)
;;    (->SubtractShrink)])

;; (defn shrinky-dink [l i g size]
;;   (assert (vector? l))
;;   (loop [passes (seq shrink-passes)
;;          state (init (first passes))
;;          value (nth l i)
;;          l l
;;          success false
;;          v nil
;;          it 0]
;;     #_(prn (first passes) it value)
;;     (if (> it 10000)
;;       (do
;;         (println "too much shrinky-dink")
;;         (throw (ex-info "too much shrinky-dink"
;;                         {:passes passes
;;                          :state state
;;                          :value value
;;                          :l l
;;                          :success success
;;                          :v v
;;                          :it it})))
;;       (if (and success (= 1000 it))
;;         [l v]
;;         (if (and (zero? value) (not success))
;;           nil
;;           (if (applicable? (first passes) state value)
;;             (cc/let [new-value (app (first passes) state value)]
;;               (if (or (and (neg? value) (> 1 new-value value))
;;                       (and (pos? value) (> value new-value -1)))
;;                 (cc/let [new-l (assoc l i new-value)
;;                          [t new-v] (regenerate new-l g size)]
;;                   (if (= :just t)
;;                     [new-l new-v]
;;                     #_(recur (seq shrink-passes) (init (first shrink-passes)) new-value new-l true new-v (inc it))
;;                     (recur passes (next-state (first passes) state) value l success v (inc it))))
;;                 (recur passes (next-state (first passes) state) value l success v (inc it))))
;;             (cc/let [passes (rest passes)]
;;               (if (seq passes)
;;                 (recur passes (init (first passes)) value l success v (inc it))
;;                 (when success
;;                   [l v])))))))))


;; (defn shrinky-dink [l i g size]
;;   (assert (vector? l))
;;   (cc/let [v (nth l i)
;;            threshold (if (neg? v)
;;                        ...
;;                        (Long/highestOneBit v))]
;;     (loop [passes (seq shrink-passes)
;;            state (init (first passes))
;;            value (nth l i)
;;            l l
;;            success false
;;            v nil
;;            it 0]
;;       #_(prn (first passes) it value)
;;       (if (> it 10000)
;;         (do
;;           (println "too much shrinky-dink")
;;           (throw (ex-info "too much shrinky-dink"
;;                           {:passes passes
;;                            :state state
;;                            :value value
;;                            :l l
;;                            :success success
;;                            :v v
;;                            :it it})))
;;         (if (and (zero? value) (not success))
;;           nil
;;           (if (applicable? (first passes) state value)
;;             (cc/let [new-value (app (first passes) state value)]
;;               (if (or (and (neg? value) (> 1 new-value value))
;;                       (and (pos? value) (> value new-value -1)))
;;                 (cc/let [new-l (assoc l i new-value)
;;                          [t new-v] (regenerate new-l g size)]
;;                   (if (= :just t)
;;                     (recur (seq shrink-passes) (init (first shrink-passes)) new-value new-l true new-v (inc it))
;;                     (recur passes (next-state (first passes) state) value l success v (inc it))))
;;                 (recur passes (next-state (first passes) state) value l success v (inc it))))
;;             (cc/let [passes (rest passes)]
;;               (if (seq passes)
;;                 (recur passes (init (first passes)) value l success v (inc it))
;;                 (when success
;;                   [l v])))))))))


;; (defn rose-tree-2 [v l g size make-tree]
;;   (letfn [(popping [new-v new-l success?]
;;             (when (seq new-l)
;;               (cc/let [new-new-l (pop new-l)
;;                        [t new-new-v] (regenerate new-new-l g size)]
;;                 (if (= :just t)
;;                   (recur new-new-v new-new-l true)
;;                   (if success?
;;                     (cc/list (rose-tree-2 new-v new-l g size make-tree))
;;                     (parallel))))))
;;           (parallel []
;;             (when (seq l)
;;               (cc/for [i (range (count l))
;;                        :let [r (shrinky-dink l i g size)]
;;                        :when r
;;                        :let [[new-l new-v] r]]
;;                 (rose-tree-2 new-v new-l g size make-tree))))]
;;     (make-tree
;;      (if (instance? clojure.lang.IObj v)
;;        (vary-meta v assoc ::l l ::s size)
;;        v)
;;      (popping nil l false))))


(defn rose-tree-3 [l make-tree]
  (letfn [(make-rose [l make-tree k]
            (make-tree
             l
             #_(if (instance? clojure.lang.IObj v)
               (vary-meta v assoc ::l l ::s size)
               v)
             (wide 0 l k)))
          (wide [i l k]
            (narrow i l (partial (fn f [ii l]
                                   (when-not (= ii i)
                                     (narrow (mod ii (count l))
                                             l
                                             (partial f (inc ii)))))
                                 (inc i))))
          (narrow [i l k]
            (let [n (nth l i)]
              (if (zero? n)
                (k l)
                (concat (for [n' (cons 0
                                       (->> (if (neg? n) -1 1)
                                            (iterate #(* % 2))
                                            (cc/map #(bit-or % (bit-and n 1)))
                                            (cc/map #(bit-or % (bit-shift-left (unsigned-bit-shift-right n 63) 63)))
                                            (take-while
                                             (if (neg? n)
                                               #(> % n)
                                               #(> n %)))))]
                          (make-rose (assoc l i n') make-tree k)) (k l)))))]
    (make-rose l make-tree (fn [l] nil))))

(defn rose-tree-1 [v l g size]
  (assert (vector? g))
  (assert (keyword? (nth g 0)))
  (rose-tree-2 v l g size rose/make-rose))

(comment

  (load-file "/home/kevin/Sync/src/surmise/src/com/manigfeald/surmise.clj")
  (load-file "/home/kevin/Sync/src/surmise/test/com/manigfeald/surmise_test.clj")

  (in-ns 'com.manigfeald.surmise)

  (in-ns 'com.manigfeald.surmise-test)

  (run-tests)

  (def main-thread (Thread/currentThread))
  
  (.interrupt main-thread)

  (rose-tree-2
   nil
   [-8 4]
   (:com.manigfeald.surmise/generator (list large-integer))
   Long/MAX_VALUE
   cc/vector)

  (do
    (with-out-str
      (prn
       (rose-tree-2
        nil
        (vec (range -2 2))
        (:com.manigfeald.surmise/generator (list byte))
        Long/MAX_VALUE
        cc/vector)))
    nil)

  
  (prn (com.manigfeald.surmise/shrinky-dink [-1] 0 (:com.manigfeald.surmise/generator com.manigfeald.surmise/large-integer) Long/MAX_VALUE))
  (shrinky-dink [133154750] 0 (:com.manigfeald.surmise/generator large-integer) Long/MAX_VALUE)


  (prn (com.manigfeald.surmise/shrinky-dink [-9125064582277025061]
                                            0
                                            (:com.manigfeald.surmise/generator nat)
                                            3))
  
  (prn (com.manigfeald.surmise/shrinky-dink [-1]
                                            0
                                            (:com.manigfeald.surmise/generator nat)
                                            3))

  )