content: {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE TupleSections #-}
  
  {-# OPTIONS_GHC -Wno-orphans #-}
  {-# OPTIONS_GHC -Wno-name-shadowing #-}
  
  module Arbitrary
    ( subexprs
    , Env(..)
    , NoAnfEnv(..)
    , AnfSymbol(..)
    , FlatAnfEnv(..)
    , ChainedAnfEnv(..)
    ) where
  
  import qualified Data.Text                 as Text
  import qualified Data.HashMap.Strict       as M
  import Test.Tasty.QuickCheck
  import GHC.Generics
  
  import Language.Fixpoint.Types.Refinements as R
  import Language.Fixpoint.Parse             (isNotReserved)
  import Language.Fixpoint.Types             as T hiding (Result)
  import Language.Fixpoint.Types.Spans       as Spans
  import Data.Traversable                    (for)
  
  {-
  
  -- An example of how the Arbitrary Expr instance can be used.
  -- Note that `expr == rr (showpp expr)` is *not* something we expect to hold.
  -- See https://github.com/ucsd-progsys/liquid-fixpoint/issues/46
  
  quickCheckTests :: TestTree
  quickCheckTests
    = testGroup "Properties"
        [ testProperty "prop_pprint_parse_inv_expr" prop_pprint_parse_inv_expr
        ]
  
  prop_pprint_parse_inv_expr :: Expr -> Bool
  prop_pprint_parse_inv_expr expr = expr == rr (showpp expr)
  
  -}
  
  -- NOTE: `shrink _ = mempty` is identical to the default (implicit) shrink implementation.
  -- We prefer to make it explicit.
  
  instance Arbitrary Expr where
    arbitrary = sized arbitraryExpr
    shrink x = filter valid $ genericShrink x
      where
        valid e@(T.PAnd es)     = length es >= 2 && all valid (subexprs e)
        valid e@(T.POr es)      = length es >= 2 && all valid (subexprs e)
        valid e@(T.PAll es _)   = length es >= 2 && all valid (subexprs e)
        valid e@(T.PExist es _) = length es >= 2 && all valid (subexprs e)
        valid e                 = all valid (subexprs e)
  
  subexprs :: Expr -> [Expr]
  subexprs (ESym _)        = []
  subexprs (ECon _)        = []
  subexprs (EVar _)        = []
  subexprs (EApp e0 e1)    = [e0, e1]
  subexprs (ENeg e)        = [e]
  subexprs (EBin _ e0 e1)  = [e0, e1]
  subexprs (EIte e0 e1 e2) = [e0, e1, e2]
  subexprs (ECst e _)      = [e]
  subexprs (ELam _ e)      = [e]
  subexprs (ETApp e _)     = [e]
  subexprs (ETAbs e _)     = [e]
  subexprs (T.PAnd es)     = es
  subexprs (T.POr es)      = es
  subexprs (T.PNot e)      = [e]
  subexprs (PImp e0 e1)    = [e0, e1]
  subexprs (PIff e0 e1)    = [e0, e1]
  subexprs (PAtom _ e0 e1) = [e0, e1]
  subexprs (PKVar _ _)     = []
  subexprs (PAll _ e)      = [e]
  subexprs (PExist _ e)    = [e]
  subexprs (PGrad _ _ _ e) = [e]
  subexprs (ECoerc _ _ e)  = [e]
  subexprs (ELet _ e1 e2)  = [e1, e2]
  
  -- TODO: Adjust frequencies
  -- | To ensure this reliably terminates we require that `zeroExprGen` generates
  -- atomic expressions.
  arbitraryFiniteExpr
    :: Gen Expr -- ^ called when the Int is zero.
    -> Int
    -> Gen Expr
  arbitraryFiniteExpr zeroExprGen 0 = zeroExprGen
  arbitraryFiniteExpr zeroExprGen n = frequency
    [ (1, EApp <$> arbitraryExpr' <*> arbitraryExpr')
    , (1, ENeg <$> arbitraryExpr')
    , (1, EBin <$> arbitrary <*> arbitraryExpr' <*> arbitraryExpr')
    , (1, EIte <$> arbitraryExpr' <*> arbitraryExpr' <*> arbitraryExpr')
    , (1, ECst <$> arbitraryExpr' <*> arbitrary)
    , (1, ELam <$> arbitrary <*> arbitraryExpr')
    , (1, ETApp <$> arbitraryExpr' <*> arbitrary)
    , (1, ETAbs <$> arbitraryExpr' <*> arbitrary)
    , (1, T.PAnd <$> arbitraryExprList)
    , (1, T.POr <$> arbitraryExprList)
    , (1, T.PNot <$> arbitraryExpr')
    , (1, PImp <$> arbitraryExpr' <*> arbitraryExpr')
    , (1, PIff <$> arbitraryExpr' <*> arbitraryExpr')
    , (1, PAtom <$> arbitrary <*> arbitraryExpr' <*> arbitraryExpr')
    , (1, PKVar <$> arbitrary <*> arbitrary)
    , (1, PAll <$> arbitraryList arbitrary <*> arbitraryExpr')
    , (1, PExist <$> arbitraryList arbitrary <*> arbitraryExpr')
    , (1, PGrad <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitraryExpr')
    , (1, ECoerc <$> arbitrary <*> arbitrary <*> arbitraryExpr')
    , (1, ELet <$> arbitrary <*> arbitraryExpr' <*> arbitraryExpr')
    ]
    where
      arbitraryExpr' = arbitraryFiniteExpr zeroExprGen (n `div` 2)
      arbitraryList :: Gen a -> Gen [a]
      arbitraryList gen = choose (2, 3) >>= (`vectorOf` gen)
      arbitraryExprList = arbitraryList arbitraryExpr'
  
  -- | Generates a finite expression, with the logarithm of the Int given
  -- suggesting the depth of the expression tree.
  arbitraryExpr :: Int -> Gen Expr
  arbitraryExpr = arbitraryFiniteExpr arbitraryAtomicExpr
  
  arbitraryAtomicExpr :: Gen Expr
  arbitraryAtomicExpr = oneof [ESym <$> arbitrary, ECon <$> arbitrary, EVar <$> arbitrary]
  
  arbitraryEqualityConstraint :: Gen Brel
  arbitraryEqualityConstraint = oneof [pure Eq, pure Ueq]
  
  -- | Generates an expression that involves the given Symbol in an Eq or Ueq
  -- PAtom with an arbitrary non-lq_anf$ EVar.
  arbitraryExprInvolving :: Symbol -> Int -> Gen Expr
  arbitraryExprInvolving sym = arbitraryFiniteExpr . pure $ EVar sym
  
  instance Arbitrary KVar where
    arbitrary = KV <$> arbitrary
  
  -- NOTE: This dummy Arbitrary instance returns a constant GradInfo.
  instance Arbitrary GradInfo where
    arbitrary = pure $ GradInfo (SS pos pos) Nothing
      where pos = Spans.dummyPos "<unknown>"
    shrink _ = mempty
  
  instance Arbitrary Subst where
    arbitrary = do
      n <- choose (0, 3)
      l <- vectorOf n arbitrary
      return $ Su $ M.fromList l
    shrink _ = mempty
  
  -- | This instance only creates `FVar` when they would be in scope from an
  -- enclosing `FAbs`, and does not create `FObj`s
  instance Arbitrary Sort where
    arbitrary = sized arbitrarySort
    shrink = genericShrink
  
  -- | Create an arbitrary well-formed sort that does not contain `FObj`s.
  --
  -- The sort is \"well-formed\" in the sense that all `FVar`s have an enclosing
  -- `FAbs` bringing them into scope.
  arbitrarySort :: Int -> Gen Sort
  arbitrarySort = arbitrarySortPossiblyInvolving []
  
  -- | Create an arbitrary sort, possibly involving the variables represented by
  -- the list of Ints. Can possibly create a `FAbs` that will also possibly
  -- reference the new variable in an `FVar`, even when the first argument is
  -- [].
  arbitrarySortPossiblyInvolving :: [Int] -> Int -> Gen Sort
  arbitrarySortPossiblyInvolving [] n = frequency
    [ (4, arbitrarySortNoAbs n)
    , (1, newAbs [] n) ]
  arbitrarySortPossiblyInvolving vars n = do
    let fvar = oneof $ pure . FVar <$> vars
    frequency
      [ (1, arbitrarySortNoAbs n)
      , (1, FFunc <$> fvar <*> arbitrarySortPossiblyInvolving vars (n `div` 2))
      , (1, FFunc <$> arbitrarySortPossiblyInvolving vars (n `div` 2) <*> fvar)
      , (1, FApp <$> fvar <*> arbitrarySortPossiblyInvolving vars (n `div` 2))
      , (1, FApp <$> arbitrarySortPossiblyInvolving vars (n `div` 2) <*> fvar)
      , (1, fvar)
      , (1, newAbs vars n)
      ]
  
  -- | Create a new FAbs sort whose body might involve the newly created variable.
  -- First argument is the variables already in scope.
  newAbs :: [Int] -> Int -> Gen Sort
  newAbs vars n = do
    v <- arbitrary
    if v `elem` vars
      then discard v
      else FAbs v <$> arbitrarySortPossiblyInvolving (v:vars) (n `div` 2)
  
  -- | Does not create FObj, FAbs, or FVar
  arbitrarySortNoAbs :: Int -> Gen Sort
  arbitrarySortNoAbs 0 = oneof (pure <$> [FInt, FReal, FNum, FFrac])
  arbitrarySortNoAbs n = frequency
    [ (1, FFunc <$> arbitrarySortNoAbs' <*> arbitrarySortNoAbs')
    , (1, FTC <$> arbitrary)
    , (1, FApp <$> arbitrarySortNoAbs' <*> arbitrarySortNoAbs')
    ]
    where
      arbitrarySortNoAbs' = arbitrarySortNoAbs (n `div` 2)
  
  instance Arbitrary Brel where
    arbitrary = oneof (map return [Eq, Ne, Gt, Ge, Lt, Le, Ueq, Une])
    shrink _ = mempty
  
  instance Arbitrary Bop where
    arbitrary = oneof (map return [Plus, Minus, Times, Div, Mod])
    shrink _ = mempty
  
  instance Arbitrary SymConst where
    arbitrary = SL . unShortLowercaseAlphabeticText <$> arbitrary
  
  -- | Note that we rely below on the property that the Arbitrary instance for
  -- Symbol cannot create lq_anf$ vars.
  instance Arbitrary Symbol where
    arbitrary = (symbol :: Text.Text -> Symbol) . unShortLowercaseAlphabeticText <$> arbitrary
    shrink _ = mempty
  
  newtype ShortLowercaseAlphabeticText = ShortLowercaseAlphabeticText { unShortLowercaseAlphabeticText :: Text.Text }
    deriving (Eq, Show, Generic)
  
  instance Arbitrary ShortLowercaseAlphabeticText where
    arbitrary = ShortLowercaseAlphabeticText <$> (choose (5, 12) >>= \n -> Text.pack <$> (vectorOf n char `suchThat` valid))
      where
        char = elements ['a'..'z']
        valid x = isNotReserved x && not (isFixKey (Text.pack x))
    shrink _ = mempty
  
  instance Arbitrary FTycon where
    arbitrary = do
      c <- elements ['A'..'Z']
      t <- unShortLowercaseAlphabeticText <$> arbitrary
      return $ symbolFTycon $ dummyLoc $ symbol $ c `Text.cons` t
    shrink _ = mempty
  
  instance Arbitrary Constant where
    arbitrary = oneof [ I <$> arbitrary `suchThat` (>= 0) -- Negative values use `ENeg`
                      , R <$> arbitrary `suchThat` (>= 0) -- Negative values use `ENeg`
                      , L . unShortLowercaseAlphabeticText <$> arbitrary <*> arbitrary
                      ]
    shrink (I x) = I <$> shrink x
    shrink (R x) = R <$> shrink x
    shrink (L x y) = L x <$> shrink y
  
  -- | Used in UndoANFTests.
  newtype AnfSymbol = AnfSymbol { unAnfSymbol :: Symbol }
    deriving (Eq, Show, Generic)
  instance Arbitrary AnfSymbol where
    arbitrary = AnfSymbol . mappendSym anfPrefix <$> arbitrary
    shrink = mempty
  
  -- | This instance does **not** create Refts with anf symbols.
  instance Arbitrary Reft where
    arbitrary = reft <$> arbitrary <*> arbitrary
    shrink = genericShrink
  
  -- | This instance does **not** create SortedRefts with anf symbols.
  instance Arbitrary SortedReft where
    arbitrary = sized $ arbitrarySortedReft (const arbitrary) (const arbitrary)
    shrink = genericShrink
  
  arbitrarySortedReft :: (Int -> Gen Sort) -> (Int -> Gen Symbol) -> Int -> Gen SortedReft
  arbitrarySortedReft sortGen symGen n = do
    sort <- sortGen n
    eq <- arbitraryEqualityConstraint
    sym <- symGen n
    RR sort . reft sym . PAtom eq (EVar sym) <$> arbitrary
  
  newtype IntSortedReft = IntSortedReft { unIntSortedReft :: SortedReft }
    deriving (Eq, Show)
  
  instance Arbitrary IntSortedReft where
    arbitrary = sized $ fmap IntSortedReft . arbitrarySortedReft (const . pure $ FInt) (const arbitrary)
  
  -- | Base environment with no declared properties; do not add an Arbitrary
  -- instance to this and instead use newtypes.
  newtype Env = Env { unEnv :: [(Symbol, SortedReft)] }
    deriving (Eq, Show, Generic)
  
  shrinkEnv :: Env -> [Env]
  shrinkEnv = fmap Env . traverse (traverse shrink) . unEnv
  
  -- | Env without anf vars.
  newtype NoAnfEnv = NoAnfEnv { unNoAnfEnv :: Env }
    deriving (Eq, Show, Generic)
  instance Arbitrary NoAnfEnv where
    arbitrary = sized (fmap NoAnfEnv . arbitraryEnv gen)
      where
        -- | Note that this relies on the property that the Arbitrary instance for
        -- Symbol cannot create lq_anf$ vars.
        gen n = vectorOf n ((\a b -> (a, unIntSortedReft b)) <$> arbitrary <*> arbitrary)
    shrink = fmap NoAnfEnv . shrinkEnv . unNoAnfEnv
  
  -- | Env with anf vars that do not reference further anf vars.
  newtype FlatAnfEnv = FlatAnfEnv { unFlatAnfEnv :: Env }
    deriving (Eq, Show, Generic)
  instance Arbitrary FlatAnfEnv where
    arbitrary = sized (fmap FlatAnfEnv . arbitraryEnv gen)
      where
        anfsGen n = vectorOf n ((\a b -> (unAnfSymbol a, unIntSortedReft b)) <$> arbitrary <*> arbitrary)
        gen = finalAnfGen anfsGen finalFlatGen
        finalFlatGen :: [(Symbol, SortedReft)] -> Gen (Symbol, SortedReft)
        -- This creates a final symbol which is either the conjunction or
        -- disjunction of all the anf symbols.
        finalFlatGen anfs = do
          conjOrDisj <- oneof $ pure <$> [T.PAnd, T.POr]
          let ultimateAnfExpr = conjOrDisj $ EVar . fst <$> anfs
          sym <- arbitrary
          ultimateAnfSym <- arbitrary
          pure (sym, RR FInt (reft ultimateAnfSym (PAtom Eq (EVar ultimateAnfSym) ultimateAnfExpr)))
    -- TODO
    shrink (FlatAnfEnv (Env (_x : xs))) = pure . FlatAnfEnv . Env $ xs
    shrink _ = mempty
  
  -- | Given a generator for a bunch of (`Symbol`, `SortedReft`) pairs which bind
  -- lq_anf$ vars, and another generator that takes those pairs and binds a
  -- non-lq_anf$ var to some subset of them, this function generates those pairs
  -- plus the \"final\" non-lq_anf$ expression, which represents the \"original\"
  -- expression brought to ANF.
  finalAnfGen :: (Int -> Gen [(Symbol, SortedReft)]) -> ([(Symbol, SortedReft)] -> Gen (Symbol, SortedReft)) -> Int -> Gen [(Symbol, SortedReft)]
  finalAnfGen anfsGen finalGen n = do
    anfs <- anfsGen n
    ultimateAnf <- finalGen anfs
    pure $ ultimateAnf : anfs
  
  -- | Create an arbitrary env up to size k with the given generator for Symbols
  -- and SortedRefts
  arbitraryEnv :: (Int -> Gen [(Symbol, SortedReft)]) -> Int -> Gen Env
  
  arbitraryEnv gen k = Env <$> (choose (0, k) >>= gen)
  
  -- | Env with anf vars that form a list of references.
  newtype ChainedAnfEnv = ChainedAnfEnv { unChainedAnfEnv :: Env }
    deriving (Eq, Show, Generic)
  instance Arbitrary ChainedAnfEnv where
    arbitrary = sized (fmap ChainedAnfEnv . arbitraryEnv gen)
      where
        gen = finalAnfGen (chainedAnfGen anfSymNGen) finalChainedGen
        finalChainedGen :: [(Symbol, SortedReft)] -> Gen (Symbol, SortedReft)
        finalChainedGen anfs =
          case anfs of
            -- No ANFs, so just an arbitrary int sorted expression will do
            [] -> fmap unIntSortedReft <$> arbitrary
            ((penultimateSym, _):_) -> do
              sym <- arbitrary
              let sreft = RR FInt (reft sym (PAtom Eq (EVar sym) (EVar penultimateSym)))
              (, sreft) <$> arbitrary
    -- TODO
    shrink (ChainedAnfEnv (Env (_x : xs))) = pure . ChainedAnfEnv . Env $ xs
    shrink _ = mempty
  
  -- | Creates a "chain" of referencing `lq_anf$` var Symbols of length `n` such
  -- that the first symbol references the second which references the third, and
  -- so on.  The last symbol is bound to an arbitrary non-lq_anf$ var.
  --
  -- This list is in an acceptable form to be passed to `finalAnfGen` to close the
  -- loop.
  chainedAnfGen :: (Int -> Gen AnfSymbol) -> Int -> Gen [(Symbol, SortedReft)]
  chainedAnfGen _ 0 = pure []
  chainedAnfGen symGen n = do
    syms <- fmap unAnfSymbol <$> for [1..n+1] symGen
    finalSym <- arbitrary
    let symPairs :: [(Symbol, Symbol)]
        symPairs = pairs (syms ++ [finalSym])
    for symPairs $ \(sym, prevSym) -> do
      otherSym <- arbitrary
      prevSymExpr <- arbitraryExprInvolving prevSym n
      pure (sym, RR FInt (reft otherSym (PAtom Eq (EVar otherSym) prevSymExpr)))
    where
      pairs xs = zip xs (drop 1 xs)
  
  -- This is not random, but is simplified so that you can make chains more
  -- easily.
  anfSymNGen :: Int -> Gen AnfSymbol
  anfSymNGen i = pure . AnfSymbol . mappendSym anfPrefix . symbol . show $ i
  