-- | Utilities functions for manipulating tiers (sized lists of lists)
module Test.LeanCheck.Tiers
  (
  -- * Additional tiers constructors
    consFromList
  , consFromAscendingList
  , consFromStrictlyAscendingList
  , consFromSet
  , consFromNoDupList

  -- * Products of tiers
  , product3With
  , productMaybeWith

  -- * Tiers of lists
  , listsOf
  , ascendingListsOf
  , strictlyAscendingListsOf
  , setsOf
  , noDupListsOf
  , products
  , listsOfLength

  , deleteT
  , normalizeT

  -- * Tiers of choices
  , choices
  , ascendingChoices
  , strictlyAscendingChoices
  )
where

import Test.LeanCheck.Basic
import Data.Maybe (catMaybes)

-- | Given a constructor that takes a list,
--   return tiers of applications of this constructor.
--
--   This is equivalent to 'cons1'.
consFromList :: Listable a => ([a] -> b) -> [[b]]
consFromList = (`mapT` listsOf tiers)

-- | Given a constructor that takes a list with ascending elements,
--   return tiers of applications of this constructor.
--
-- For example, a 'Bag' represented as a list.
--
-- > consFromAscendingList Bag
consFromAscendingList :: Listable a => ([a] -> b) -> [[b]]
consFromAscendingList = (`mapT` ascendingListsOf tiers)

-- | Given a constructor that takes a list with ascending elements,
--   return tiers of applications of this constructor.
--
-- For example, a 'Set' represented as a list.
--
-- > consFromAscendingList Set
consFromStrictlyAscendingList :: Listable a => ([a] -> b) -> [[b]]
consFromStrictlyAscendingList = (`mapT` strictlyAscendingListsOf tiers)

-- | Given a constructor that takes a set of elements (as a list),
--   return tiers of applications of this constructor.
--
-- For example, a 'Set' represented as a list.
--
-- > consFromAscendingList Set
consFromSet :: Listable a => ([a] -> b) -> [[b]]
consFromSet = (`mapT` setsOf tiers)

-- | Given a constructor that takes a list with no duplicate elements,
--   return tiers of applications of this constructor.
consFromNoDupList :: Listable a => ([a] -> b) -> [[b]]
consFromNoDupList f = mapT f (noDupListsOf tiers)


-- | Like 'productWith', but over 3 lists of tiers.
product3With :: (a->b->c->d) -> [[a]] -> [[b]] -> [[c]] -> [[d]]
product3With f xss yss zss = productWith ($) (productWith f xss yss) zss

-- | Take the product of lists of tiers
--   by a function returning a 'Maybe' value
--   discarding 'Nothing' values.
productMaybeWith :: (a->b->Maybe c) -> [[a]] -> [[b]] -> [[c]]
productMaybeWith _ _ [] = []
productMaybeWith _ [] _ = []
productMaybeWith f (xs:xss) yss = map (xs **) yss
                               \/ productMaybeWith f xss yss `addWeight` 1
  where xs ** ys = catMaybes [ f x y | x <- xs, y <- ys ]


-- | Given tiers of values, returns tiers of lists of those values
--
-- > listsOf [[]] == [[[]]]
--
-- > listsOf [[x]] == [ [[]]
-- >                  , [[x]]
-- >                  , [[x,x]]
-- >                  , [[x,x,x]]
-- >                  , ...
-- >                  ]
--
-- > listsOf [[x],[y]] == [ [[]]
-- >                      , [[x]]
-- >                      , [[x,x],[y]]
-- >                      , [[x,x,x],[x,y],[y,x]]
-- >                      , ...
-- >                      ]
listsOf :: [[a]] -> [[[a]]]
listsOf xss = cons0 []
           \/ productWith (:) xss (listsOf xss) `addWeight` 1

-- | Generates several lists of the same size.
--
-- > products [ xss, yss, zss ] ==
--
-- Tiers of all lists combining elements of tiers: xss, yss and zss
products :: [ [[a]] ] -> [[ [a] ]]
products = foldr (productWith (:)) [[[]]]

-- | Delete the first occurence of an element in a tier.
--
-- For tiers without repetitions, the following holds:
--
-- > deleteT x = normalizeT . (`suchThat` (/= x))
deleteT :: Eq a => a -> [[a]] -> [[a]]
deleteT _ [] = []
deleteT y ([]:xss) = [] : deleteT y xss
deleteT y [[x]]        | x == y    = []
deleteT y ((x:xs):xss) | x == y    = xs:xss
                       | otherwise = [[x]] \/ deleteT y (xs:xss)

-- | Normalizes tiers by removing an empty tier from the end of a list of
--   tiers.
--
-- > normalizeT [xs0,xs1,...,xsN,[]] = [xs0,xs1,...,xsN]
--
--   Note this will only remove a single empty tier:
--
-- > normalizeT [xs0,xs1,...,xsN,[],[]] = [xs0,xs1,...,xsN,[]]
normalizeT :: [[a]] -> [[a]]
normalizeT [] = []
normalizeT [[]] = []
normalizeT (xs:xss) = xs:normalizeT xss

-- | Given tiers of values, returns tiers of lists with no repeated elements.
--
-- > noDupListsOf [[0],[1],[2],...] ==
-- >   [ [[]]
-- >   , [[0]]
-- >   , [[1]]
-- >   , [[0,1],[1,0],[2]]
-- >   , [[0,2],[2,0],[3]]
-- >   , ...
-- >   ]
noDupListsOf :: [[a]] -> [[[a]]]
noDupListsOf =
  ([[]]:) . concatT . choicesWith (\x xss -> mapT (x:) (noDupListsOf xss))

-- | Lists tiers of all choices of values from tiers.
-- Choices are pairs of values and tiers excluding that value.
--
-- > choices [[False,True]] == [[(False,[[True]]),(True,[[False]])]]
-- > choices [[1],[2],[3]]
-- >   == [ [(1,[[],[2],[3]])]
-- >      , [(2,[[1],[],[3]])]
-- >      , [(3,[[1],[2],[]])] ]
--
-- Each choice is sized by the extracted element.
choices :: [[a]] -> [[(a,[[a]])]]
choices = choicesWith (,)

-- | Like 'choices', but allows a custom function.
choicesWith :: (a -> [[a]] -> b) -> [[a]] -> [[b]]
choicesWith f []           = []
choicesWith f [[]]         = []
choicesWith f ([]:xss)     = [] : choicesWith (\y yss -> f y ([]:yss)) xss
choicesWith f ((x:xs):xss) = [[f x (xs:xss)]]
                          \/ choicesWith (\y (ys:yss) -> f y ((x:ys):yss)) (xs:xss)

-- | Given tiers of values,
--   returns tiers of lists of elements in ascending order
--                               (from tiered enumeration).
--
ascendingListsOf :: [[a]] -> [[[a]]]
ascendingListsOf =
  ([[]]:) . concatT . ascendingChoicesWith (\x xss -> mapT (x:) (ascendingListsOf xss))

-- > ascendingChoices [[False,True]] =
-- >   [ [(False,[[False,True]]), (True,[[True]])]
-- >   ]
--
-- > ascendingChoices [[1],[2],[3],...] =
-- >   [ [(1,[[1],[2],[3],...])]
-- >   , [(2,[[ ],[2],[3],...])]
-- >   , [(3,[[ ],[ ],[3],...])]
-- >   , ...
-- >   ]
ascendingChoices :: [[a]] -> [[(a,[[a]])]]
ascendingChoices = ascendingChoicesWith (,)

ascendingChoicesWith :: (a -> [[a]] -> b) -> [[a]] -> [[b]]
ascendingChoicesWith f []           = []
ascendingChoicesWith f [[]]         = []
ascendingChoicesWith f ([]:xss)     = [] : ascendingChoicesWith (\y yss -> f y ([]:yss)) xss
ascendingChoicesWith f ((x:xs):xss) = [[f x ((x:xs):xss)]]
                                   \/ ascendingChoicesWith f (xs:xss)

-- | Given tiers of values,
--   returns tiers of lists of elements in strictly ascending order
--                              (from tiered enumeration).
--   If you only care about whether elements are in returned lists,
--   this returns the tiers of all sets of values.
--
-- > strictlyAscendingListsOf [[0],[1],[2],...] ==
-- >   [ [[]]
-- >   , [[0]]
-- >   , [[1]]
-- >   , [[0,1],[2]]
-- >   , [[0,2],[3]]
-- >   , [[0,3],[1,2],[4]]
-- >   , [[0,1,2],[0,4],[1,3],[5]]
-- >   , ...
-- >   ]
strictlyAscendingListsOf :: [[a]] -> [[[a]]]
strictlyAscendingListsOf =
  ([[]]:) . concatT .
  strictlyAscendingChoicesWith
    (\x xss -> mapT (x:) (strictlyAscendingListsOf xss))

-- | Returns tiers of sets represented as lists of values (no repeated sets).
--   Shorthand for 'strictlyAscendingListsOf'.
setsOf :: [[a]] -> [[[a]]]
setsOf = strictlyAscendingListsOf

-- | Like 'choices', but paired tiers are always strictly ascending (in terms
--   of enumeration).
--
-- > strictlyAscendingChoices [[False,True]] == [[(False,[[True]]),(True,[[]])]]
-- > strictlyAscendingChoices [[1],[2],[3]]
-- >   == [ [(1,[[],[2],[3]])]
-- >      , [(2,[[],[],[3]])]
-- >      , [(3,[[],[],[]])]
-- >      ]
strictlyAscendingChoices :: [[a]] -> [[(a,[[a]])]]
strictlyAscendingChoices = strictlyAscendingChoicesWith (,)

-- | Like 'strictlyAscendingChoices' but customized by a function.
strictlyAscendingChoicesWith :: (a -> [[a]] -> b) -> [[a]] -> [[b]]
strictlyAscendingChoicesWith f []           = []
strictlyAscendingChoicesWith f [[]]         = []
strictlyAscendingChoicesWith f ([]:xss)     = [] : strictlyAscendingChoicesWith (\y yss -> f y ([]:yss)) xss
strictlyAscendingChoicesWith f ((x:xs):xss) = [[f x (xs:xss)]]
                                           \/ strictlyAscendingChoicesWith f (xs:xss)


-- | Given tiers, returns tiers of lists of a given length.
listsOfLength :: Int -> [[a]] -> [[[a]]]
listsOfLength n xss = products (replicate n xss)