(;module: {#;doc "Pseudo-random number generation (PRNG) algorithms."}
  [lux #- list]
  (lux (control functor
                applicative
                monad
                hash)
       (data [bit]
             [char]
             [text "Text/" Monoid<Text>]
             text/format
             [product]
             [number]
             (coll [list "List/" Fold<List>]
                   ["A" array]
                   ["D" dict]
                   ["Q" queue]
                   ["S" set]
                   ["ST" stack]
                   ["V" vector]))
       (math ["r" ratio]
             ["c" complex])))

## [Exports]
(type: #export #rec PRNG
  {#;doc "An abstract way to represent any PRNG."}
  (-> Unit [PRNG Nat]))

(type: #export (Random a)
  {#;doc "A producer of random values based on a PRNG."}
  (-> PRNG [PRNG a]))

(struct: #export _ (Functor Random)
  (def: (map f fa)
    (lambda [state]
      (let [[state' a] (fa state)]
        [state' (f a)]))))

(struct: #export _ (Applicative Random)
  (def: functor Functor<Random>)

  (def: (wrap a)
    (lambda [state]
      [state a]))

  (def: (apply ff fa)
    (lambda [state]
      (let [[state' f] (ff state)
            [state'' a] (fa state')]
        [state'' (f a)]))))

(struct: #export _ (Monad Random)
  (def: applicative Applicative<Random>)

  (def: (join ffa)
    (lambda [state]
      (let [[state' fa] (ffa state)]
        (fa state')))))

(def: #export nat
  (Random Nat)
  (lambda [prng]
    (let [[prng left] (prng [])
          [prng right] (prng [])]
      [prng (n.+ (bit;<< +32 left)
                 right)])))

(def: #export int
  (Random Int)
  (lambda [prng]
    (let [[prng left] (prng [])
          [prng right] (prng [])]
      [prng (nat-to-int (n.+ (bit;<< +32 left)
                             right))])))

(def: #export bool
  (Random Bool)
  (lambda [prng]
    (let [[prng output] (prng [])]
      [prng (|> output (bit;& +1) (n.= +1))])))

(def: (bits n)
  (-> Nat (Random Nat))
  (lambda [prng]
    (let [[prng output] (prng [])]
      [prng (bit;>>> (n.- n +64) output)])))

(def: #export real
  (Random Real)
  (do Monad<Random>
    [left (bits +26)
     right (bits +27)]
    (wrap (|> right
              (n.+ (bit;<< +27 left))
              nat-to-int
              int-to-real
              (r./ (|> +1 (bit;<< +53) nat-to-int int-to-real))))))

(def: #export deg
  (Random Deg)
  (:: Monad<Random> map real-to-deg real))

(def: #export char
  (Random Char)
  (do Monad<Random>
    [base nat]
    (wrap (char;char base))))

(def: #export (text' char-gen size)
  (-> (Random Char) Nat (Random Text))
  (if (n.= +0 size)
    (:: Monad<Random> wrap "")
    (do Monad<Random>
      [x char-gen
       xs (text' char-gen (n.dec size))]
      (wrap (Text/append (char;as-text x) xs)))))

(def: #export (text size)
  (-> Nat (Random Text))
  (text' char size))

(do-template [<name> <type> <ctor> <gen>]
  [(def: #export <name>
     (Random <type>)
     (do Monad<Random>
       [left <gen>
        right <gen>]
       (wrap (<ctor> left right))))]

  [ratio   r;Ratio   r;ratio   nat]
  [complex c;Complex c;complex real]
  )

(def: #export (seq left right)
  {#;doc "Sequencing combinator."}
  (All [a b] (-> (Random a) (Random b) (Random [a b])))
  (do Monad<Random>
    [=left left
     =right right]
    (wrap [=left =right])))

(def: #export (alt left right)
  {#;doc "Heterogeneous alternative combinator."}
  (All [a b] (-> (Random a) (Random b) (Random (| a b))))
  (do Monad<Random>
    [? bool]
    (if ?
      (do @
        [=left left]
        (wrap (+0 =left)))
      (do @
        [=right right]
        (wrap (+1 =right))))))

(def: #export (either left right)
  {#;doc "Homogeneous alternative combinator."}
  (All [a] (-> (Random a) (Random a) (Random a)))
  (do Monad<Random>
    [? bool]
    (if ?
      left
      right)))

(def: #export (rec gen)
  {#;doc "A combinator for producing recursive random generators."}
  (All [a] (-> (-> (Random a) (Random a)) (Random a)))
  (lambda [state]
    (let [gen' (gen (rec gen))]
      (gen' state))))

(def: #export (filter pred gen)
  {#;doc "Retries the generator until the output satisfies a predicate."}
  (All [a] (-> (-> a Bool) (Random a) (Random a)))
  (do Monad<Random>
    [sample gen]
    (if (pred sample)
      (wrap sample)
      (filter pred gen))))

(def: #export (maybe value-gen)
  (All [a] (-> (Random a) (Random (Maybe a))))
  (do Monad<Random>
    [some? bool]
    (if some?
      (do @
        [value value-gen]
        (wrap (#;Some value)))
      (wrap #;None))))

(do-template [<name> <type> <zero> <plus>]
  [(def: #export (<name> size value-gen)
     (All [a] (-> Nat (Random a) (Random (<type> a))))
     (if (n.> +0 size)
       (do Monad<Random>
         [x value-gen
          xs (<name> (n.dec size) value-gen)]
         (wrap (<plus> x xs)))
       (:: Monad<Random> wrap <zero>)))]

  [list   List    (;list)  #;Cons]
  [vector V;Vector V;empty V;add]
  )

(do-template [<name> <type> <ctor>]
  [(def: #export (<name> size value-gen)
     (All [a] (-> Nat (Random a) (Random (<type> a))))
     (do Monad<Random>
       [values (list size value-gen)]
       (wrap (|> values <ctor>))))]

  [array A;Array  A;from-list]
  [queue Q;Queue  Q;from-list]
  [stack ST;Stack (List/fold ST;push ST;empty)]
  )

(def: #export (set Hash<a> size value-gen)
  (All [a] (-> (Hash a) Nat (Random a) (Random (S;Set a))))
  (if (n.> +0 size)
    (do Monad<Random>
      [xs (set Hash<a> (n.dec size) value-gen)]
      (loop [_ []]
        (do @
          [x value-gen
           #let [xs+ (S;add x xs)]]
          (if (n.= size (S;size xs+))
            (wrap xs+)
            (recur [])))))
    (:: Monad<Random> wrap (S;new Hash<a>))))

(def: #export (dict Hash<a> size key-gen value-gen)
  (All [k v] (-> (Hash k) Nat (Random k) (Random v) (Random (D;Dict k v))))
  (if (n.> +0 size)
    (do Monad<Random>
      [kv (dict Hash<a> (n.dec size) key-gen value-gen)]
      (loop [_ []]
        (do @
          [k key-gen
           v value-gen
           #let [kv+ (D;put k v kv)]]
          (if (n.= size (D;size kv+))
            (wrap kv+)
            (recur [])))))
    (:: Monad<Random> wrap (D;new Hash<a>))))

(def: #export (run prng calc)
  (All [a] (-> PRNG (Random a) [PRNG a]))
  (calc prng))

## [PRNGs]
## PCG32 http://www.pcg-random.org/
## Based on this Java implementation: https://github.com/alexeyr/pcg-java

(def: pcg-32-magic-mult Nat +6364136223846793005)

(def: #export (pcg-32 [inc seed])
  {#;doc "An implementation of the PCG32 algorithm.

          For more information, please see: http://www.pcg-random.org/"}
  (-> [Nat Nat] PRNG)
  (lambda [_]
    (let [seed' (|> seed (n.* pcg-32-magic-mult) (n.+ inc))
          xor-shifted (|> seed (bit;>>> +18) (bit;^ seed) (bit;>>> +27))
          rot (|> seed (bit;>>> +59))]
      [(pcg-32 [inc seed']) (bit;rotate-right rot xor-shifted)]
      )))

## Xoroshiro128+ http://xoroshiro.di.unimi.it/
(def: #export (xoroshiro-128+ [s0 s1])
  {#;doc "An implementation of the Xoroshiro128+ algorithm.

          For more information, please see: http://xoroshiro.di.unimi.it/"}
  (-> [Nat Nat] PRNG)
  (lambda [_]
    (let [result (n.+ s0 s1)
          s01 (bit;^ s0 s1)
          s0' (|> (bit;rotate-left +55 s0)
                  (bit;^ s01)
                  (bit;^ (bit;<< +14 s01)))
          s1' (bit;rotate-left +36 s01)]
      [(xoroshiro-128+ [s0' s1']) result])
    ))

## [Values]
(def: (swap from to vec)
  (All [a] (-> Nat Nat (V;Vector a) (V;Vector a)))
  (V;put to (default (undefined)
              (V;nth from vec))
         vec))

(def: #export (shuffle seed vector)
  {#;doc "Shuffle a vector randomly based on a seed value."}
  (All [a] (-> Nat (V;Vector a) (V;Vector a)))
  (let [_size (V;size vector)
        _shuffle (foldM Monad<Random>
                        (lambda [idx vec]
                          (do Monad<Random>
                            [rand nat]
                            (wrap (swap idx (n.% _size rand) vec))))
                        vector
                        (list;n.range +0 (n.dec _size)))]
    (|> _shuffle
        (run (pcg-32 [+123 seed]))
        product;right)))