module Cardano.Wallet.Primitive.Types.Coin.Gen
    ( chooseCoin
    , genCoin
    , genCoinPositive
    , shrinkCoin
    , shrinkCoinPositive
    , genCoinPartition
    ) where

import Prelude

import Cardano.Wallet.Primitive.Types.Coin
    ( Coin (..) )
import Control.Monad
    ( replicateM )
import Data.Coerce
    ( coerce )
import Data.List.NonEmpty
    ( NonEmpty )
import Test.QuickCheck
    ( Gen, choose, sized )
import Test.QuickCheck.Extra
    ( chooseNatural, shrinkNatural )

import qualified Cardano.Wallet.Primitive.Types.Coin as Coin
import qualified Data.List.NonEmpty as NE

--------------------------------------------------------------------------------
-- Choosing coins from a range.
--------------------------------------------------------------------------------

chooseCoin :: (Coin, Coin) -> Gen Coin
chooseCoin :: (Coin, Coin) -> Gen Coin
chooseCoin = ((Natural, Natural) -> Gen Natural) -> (Coin, Coin) -> Gen Coin
coerce (Natural, Natural) -> Gen Natural
chooseNatural

--------------------------------------------------------------------------------
-- Coins chosen according to the size parameter.
--------------------------------------------------------------------------------

genCoin :: Gen Coin
genCoin :: Gen Coin
genCoin = (Int -> Gen Coin) -> Gen Coin
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen Coin) -> Gen Coin) -> (Int -> Gen Coin) -> Gen Coin
forall a b. (a -> b) -> a -> b
$ \Int
n -> Natural -> Coin
Coin (Natural -> Coin) -> (Int -> Natural) -> Int -> Coin
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Coin) -> Gen Int -> Gen Coin
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
0, Int
n)

shrinkCoin :: Coin -> [Coin]
shrinkCoin :: Coin -> [Coin]
shrinkCoin (Coin Natural
c) = Natural -> Coin
Coin (Natural -> Coin) -> [Natural] -> [Coin]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Natural -> [Natural]
shrinkNatural Natural
c

--------------------------------------------------------------------------------
-- Coins chosen according to the size parameter, but strictly positive.
--------------------------------------------------------------------------------

genCoinPositive :: Gen Coin
genCoinPositive :: Gen Coin
genCoinPositive = (Int -> Gen Coin) -> Gen Coin
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen Coin) -> Gen Coin) -> (Int -> Gen Coin) -> Gen Coin
forall a b. (a -> b) -> a -> b
$ \Int
n -> Natural -> Coin
Coin (Natural -> Coin) -> (Int -> Natural) -> Int -> Coin
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Coin) -> Gen Int -> Gen Coin
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
1, Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
n)

shrinkCoinPositive :: Coin -> [Coin]
shrinkCoinPositive :: Coin -> [Coin]
shrinkCoinPositive (Coin Natural
c) = Natural -> Coin
Coin (Natural -> Coin) -> [Natural] -> [Coin]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Natural -> Bool) -> [Natural] -> [Natural]
forall a. (a -> Bool) -> [a] -> [a]
filter (Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Natural
0) (Natural -> [Natural]
shrinkNatural Natural
c)

--------------------------------------------------------------------------------
-- Partitioning coins
--------------------------------------------------------------------------------

-- | Partitions a coin randomly into a given number of parts.
--
-- Satisfies the following properties:
--
-- prop> forAll (genCoinPartition c i) $ (==       c) . fold
-- prop> forAll (genCoinPartition c i) $ (== max 1 i) . length
--
genCoinPartition :: Coin -> Int -> Gen (NonEmpty Coin)
genCoinPartition :: Coin -> Int -> Gen (NonEmpty Coin)
genCoinPartition Coin
c Int
i =
    Coin -> NonEmpty Coin -> NonEmpty Coin
Coin.partitionDefault Coin
c (NonEmpty Coin -> NonEmpty Coin)
-> Gen (NonEmpty Coin) -> Gen (NonEmpty Coin)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (NonEmpty Coin)
genWeights
  where
    genWeights :: Gen (NonEmpty Coin)
    genWeights :: Gen (NonEmpty Coin)
genWeights = [Coin] -> NonEmpty Coin
forall a. [a] -> NonEmpty a
NE.fromList ([Coin] -> NonEmpty Coin) -> Gen [Coin] -> Gen (NonEmpty Coin)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen Coin -> Gen [Coin]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
i)
        ((Coin, Coin) -> Gen Coin
chooseCoin (Natural -> Coin
Coin Natural
1, Coin -> Coin -> Coin
forall a. Ord a => a -> a -> a
max (Natural -> Coin
Coin Natural
1) Coin
c))