{-# LANGUAGE RecursiveDo #-}
module Lambda where
import Control.Applicative
import Data.List as List
import Data.Foldable
import Test.Tasty
import Test.Tasty.HUnit as HU
import Test.Tasty.QuickCheck as QC

import Text.Earley

import qualified Arbitrary

tests :: TestTree
tests = testGroup "Lambda"
  [ HU.testCase "Generate exactly 0" $
      exactly 0 gen @?= []
  , HU.testCase "Generate upTo 0" $
      upTo 0 gen @?= []
  , HU.testCase "Generate exactly 4" $
      sort (snd <$> exactly 4 gen)
      @?=
      ["(a)a","(a)b","(aa)","(ab)","(b)a","(b)b","(ba)","(bb)"
      ,"\\a.a","\\a.b","\\b.a","\\b.b","a(a)","a(b)","a+aa","a+ab"
      ,"a+ba","a+bb","aa+a","aa+b","aaaa","aaab","aaba","aabb"
      ,"ab+a","ab+b","abaa","abab","abba","abbb","b(a)","b(b)"
      ,"b+aa","b+ab","b+ba","b+bb","ba+a","ba+b","baaa","baab"
      ,"baba","babb","bb+a","bb+b","bbaa","bbab","bbba","bbbb"
      ]
  , HU.testCase "upTo contains exactly" $ List.and (do
      m <- [0..5]
      let ys = snd <$> upTo m gen
      n <- [0..m]
      (_, x) <- upTo n gen
      return $ x `List.elem` ys)
    @? "exactly contains upTo"
  , HU.testCase "language contains upTo" $ do
    let ys = snd <$> language gen
    List.and (do
      n <- [0..5]
      (_, x) <- upTo n gen
      return $ x `List.elem` ys)
    @? "exactly contains upTo"
  , QC.testProperty "Arbitrary" $ do
    let p = parser grammar
    (e, s) <- Arbitrary.arbitrary $ generator grammar tokens
    return
      $ [e] === fst (fullParses p $ prettyExpr 0 e)
      .&&. [e] === fst (fullParses p s)
  ]
  where
    gen = generator grammar tokens

data Expr
  = Var Char
  | Lam String Expr
  | App Expr Expr
  | Add Expr Expr
  deriving (Eq, Ord, Show)

prettyExpr :: Int -> Expr -> String
prettyExpr _ (Var c) = [c]
prettyExpr d (Lam xs e) = prettyParens (d > 0) $ "\\" ++ xs ++ "." ++ prettyExpr d e
prettyExpr d (Add a b) = prettyParens (d > 1) $ prettyExpr 2 a ++ "+" ++ prettyExpr 1 b
prettyExpr d (App a b) = prettyParens (d > 3) $ prettyExpr 3 a ++ prettyExpr 4 b

prettyParens :: Bool -> String -> String
prettyParens True s  = "(" ++ s ++ ")"
prettyParens False s = s

tokens :: String
tokens = "(\\ab.+*)"

instance Arbitrary Expr where
  arbitrary = sized go
    where
      var = elements "ab"
      go 0 = Var <$> var
      go n = oneof
        [ Var <$> var
        , Lam <$> (take 2 <$> listOf1 var) <*> go'
        , App <$> go' <*> go'
        , Add <$> go' <*> go'
        ]
        where
          go' = go (n `div` 10)

  shrink (Var _) = []
  shrink (Lam xs e) = e : [Lam xs' e' | xs' <- shrink xs, not (null xs), e' <- shrink e]
  shrink (App a b) = a : b : [App a' b' | a' <- shrink a, b' <- shrink b]
  shrink (Add a b) = a : b : [Add a' b' | a' <- shrink a, b' <- shrink b]

grammar :: Grammar r (Prod r String Char Expr)
grammar = mdo
  let v = asum (token <$> "ab")
        <?> "variable"
  x1 <- rule
    $ Lam <$ token '\\' <*> some v <* token '.' <*> x1
    <|> x2
    <?> "lambda"
  x2 <- rule
    $ Add <$> x3 <* token '+' <*> x2
    <|> x3
    <?> "sum"
  x3 <- rule
    $ App <$> x3 <*> x4
    <|> x4
    <?> "application"
  let x4 = Var <$> v
        <|> token '(' *> x1 <* token ')'
  return x1