{-# LANGUAGE DeriveAnyClass #-}

module Main where

import Control.Monad.State
import Data.Foldable (toList)
import Data.List (sort)
import qualified Data.Set as Set
import Test.Tasty
import Test.Tasty.HUnit

import Little.Earley
import Little.Earley.Examples (ArithN(..), CharT(..), arithG, matchCharT)
import Little.Earley.Internal.Core

-- Ambiguous grammar for arithmetic expressions
aArithRules :: ArithN -> [Rule ArithN CharT]
aArithRules n = case n of
  SUM ->
    [ [ N NUMBER ]
    , [ N SUM, T (OneOf ['+', '-']), N SUM ] ]
  NUMBER ->
    [ [ T Digit ]
    , [ T Digit, N NUMBER ] ]
  _ -> []

aArithG :: Grammar ArithN CharT Char
aArithG = mkGrammar aArithRules matchCharT

-- A ::= A | ""
-- Infinitely many possible parse trees for the empty string.
dumbG :: Grammar () () a
dumbG = Grammar
  { rules = \() -> [ [ N () ], [] ]
  , match = \() _ -> True
  , isNullable = \() -> True
  }

-- A ::= "" | B
-- B ::= A
-- Example from https://loup-vaillant.fr/tutorials/earley-parsing/empty-rules
pingpongG :: Grammar Bool () a
pingpongG = Grammar
  { rules = \b -> case b of
      True -> [ [], [ N False ] ]
      False -> [ [ N True ] ]
  , match = \_ _ -> True
  , isNullable = \_ -> True
  }

initialItems1 :: [Item ArithN CharT]
initialItems1 =
  [ Item (RuleId SUM 0) [ N PRODUCT ] 0
  , Item (RuleId SUM 1) [ N SUM, T (OneOf ['+', '-']), N PRODUCT ] 0 ]

state1 :: S ArithN CharT
state1 = emptyS { currItemSet = Set.fromList initialItems1 }

tree_12_3 :: TreeT f ArithN CharT Char
tree_12_3 =
  Brch (RuleId SUM 1) 0 5
    [ Brch (RuleId SUM 1) 0 3
        [ Brch (RuleId SUM 0) 0 1 [Brch (RuleId NUMBER 0) 0 1 [Leaf 0 Digit '1']]
        , Leaf 1 (OneOf "+-") '+'
        , Brch (RuleId SUM 0) 2 3 [Brch (RuleId NUMBER 0) 2 3 [Leaf 2 Digit '2']]
        ]
    , Leaf 3 (OneOf "+-") '+'
    , Brch (RuleId SUM 0) 4 5 [Brch (RuleId NUMBER 0) 4 5 [Leaf 4 Digit '3']]
    ]

tree_1_23 :: TreeT f ArithN CharT Char
tree_1_23 =
  Brch (RuleId SUM 1) 0 5
    [ Brch (RuleId SUM 0) 0 1 [Brch (RuleId NUMBER 0) 0 1 [Leaf 0 Digit '1']]
    , Leaf 1 (OneOf "+-") '+'
    , Brch (RuleId SUM 1) 2 5
        [ Brch (RuleId SUM 0) 2 3 [Brch (RuleId NUMBER 0) 2 3 [Leaf 2 Digit '2']]
        , Leaf 3 (OneOf "+-") '+'
        , Brch (RuleId SUM 0) 4 5 [Brch (RuleId NUMBER 0) 4 5 [Leaf 4 Digit '3']]
        ]
    ]

--

data E = E
  deriving (Eq, Ord, Enum, Bounded, Show, PrettyPrint)

iteRules :: E -> [Rule E String]
iteRules E =
  [ [ T "if", N E, T "then", N E ]
  , [ T "if", N E, T "then", N E, T "else", N E ]
  , [ T "T" ]
  ]

iteG :: Grammar E String String
iteG = mkGrammar iteRules (==)

trueLeaf :: Int -> TreeT f E String String
trueLeaf i = Brch (RuleId E 2) i (i+1) [Leaf i "T" "T"]

iteInput :: [String]
iteInput = ["if", "T", "then", "if", "T", "then", "T", "else", "T"]

tree_it_ite :: TreeT f E String String
tree_it_ite =
  Brch (RuleId E 0) 0 9
    [ Leaf 0 "if" "if"
    , trueLeaf 1
    , Leaf 2 "then" "then"
    , Brch (RuleId E 1) 3 9
        [ Leaf 3 "if" "if"
        , trueLeaf 4
        , Leaf 5 "then" "then"
        , trueLeaf 6
        , Leaf 7 "else" "else"
        , trueLeaf 8
        ]
    ]

tree_ite_it :: TreeT f E String String
tree_ite_it =
  Brch (RuleId E 1) 0 9
    [ Leaf 0 "if" "if"
    , trueLeaf 1
    , Leaf 2 "then" "then"
    , Brch (RuleId E 0) 3 7
        [ Leaf 3 "if" "if"
        , trueLeaf 4
        , Leaf 5 "then" "then"
        , trueLeaf 6
        ]
    , Leaf 7 "else" "else"
    , trueLeaf 8
    ]

tests :: TestTree
tests = testGroup "Tests"
  [ testCase "initial" $
      initialS arithG SUM @?= state1
  , testCase "predict" $ do
      let act0 = stepItem arithG Nothing (Item (RuleId SUM 0) [N PRODUCT] 0)
      sort (todo (execState act0 state1))
        @?= [ Item (RuleId PRODUCT 0) [ N FACTOR ] 0
            , Item (RuleId PRODUCT 1) [ N PRODUCT, T (OneOf ['*', '/']), N FACTOR ] 0 ]
      let act1 = stepItem arithG Nothing
                   (Item (RuleId SUM 1) [N SUM, T (OneOf ['+', '-']), N PRODUCT] 0)
      todo (execState act1 state1) @?= []
  , testCase "scan" $ do
      let act0 = stepItem arithG (Just '(')
                   (Item (RuleId FACTOR 1) [T (OneOf ['(']), N SUM, T (OneOf [')'])] 0)
      Set.toList (nextItemSet (execState act0 state1))
        @?= [Item (RuleId FACTOR 1) [N SUM, T (OneOf [')'])] 0]
  , testCase "complete" $ do
      let act0 = stepItem arithG Nothing (Item (RuleId PRODUCT 0) [] 0)
      todo (execState act0 state1)
        @?= [Item (RuleId SUM 0) [] 0]
  , testCase "inner loop" $ do
      let act0 = step arithG (Just '3')
      toList (allItemSets (execState act0 state1))
        @?= [ Set.fromList
              [ Item (RuleId SUM 0) [N PRODUCT] 0
              , Item (RuleId SUM 1) [N SUM,T (OneOf "+-"),N PRODUCT] 0
              , Item (RuleId PRODUCT 0) [N FACTOR] 0
              , Item (RuleId PRODUCT 1) [N PRODUCT,T (OneOf "*/"),N FACTOR] 0
              , Item (RuleId FACTOR 0) [N NUMBER] 0
              , Item (RuleId FACTOR 1) [T (OneOf "("),N SUM,T (OneOf ")")] 0
              , Item (RuleId NUMBER 0) [T Digit] 0
              , Item (RuleId NUMBER 1) [T Digit,N NUMBER] 0
              ]
            , Set.fromList
              [ Item (RuleId NUMBER 0) [] 0
              , Item (RuleId NUMBER 1) [N NUMBER] 0
              ]
            ]
  , testCase "accepts" $ do
      accepts arithG SUM "3" @?= True
      accepts arithG SUM "33" @?= True
      accepts arithG SUM "1+1" @?= True
      accepts arithG SUM "11+1" @?= True
      accepts arithG SUM "1*1" @?= True
      accepts arithG SUM "(1)" @?= True
      accepts arithG SUM "1+2*3" @?= True
      accepts arithG SUM "(1+2)*3" @?= True
      accepts arithG SUM "1+(2*3)" @?= True
      accepts arithG SUM "(" @?= False
      accepts arithG SUM "*" @?= False
      accepts aArithG SUM "1+2+3" @?= True
      accepts iteG E ["if", "T"] @?= False
      accepts iteG E ["if", "T", "then"] @?= False
      accepts iteG E ["if", "T", "then", "T"] @?= True
      accepts iteG E ["if", "T", "then", "T", "else"] @?= False
      accepts iteG E ["if", "T", "then", "T", "else", "T"] @?= True
      accepts iteG E ["if", "T", "then", "if", "T", "then", "T", "else", "T"] @?= True
  , testCase "parse" $ do
      parseTreeSet arithG SUM "3" @?= Just
        ( Brch (RuleId SUM 0) 0 1 . (:[])
        . Brch (RuleId PRODUCT 0) 0 1 . (:[])
        . Brch (RuleId FACTOR 0) 0 1 . (:[])
        . Brch (RuleId NUMBER 0) 0 1 . (:[])
        $ Leaf 0 Digit '3' )
      parseTreeSet arithG SUM "1+2*3" @?= Just
        ( Brch (RuleId SUM 1) 0 5
            [ Brch (RuleId SUM 0) 0 1
                [Brch (RuleId PRODUCT 0) 0 1
                  [Brch (RuleId FACTOR 0) 0 1 [Brch (RuleId NUMBER 0) 0 1 [Leaf 0 Digit '1']]]]
            , Leaf 1 (OneOf "+-") '+'
            , Brch (RuleId PRODUCT 1) 2 5
                [ Brch (RuleId PRODUCT 0) 2 3
                    [Brch (RuleId FACTOR 0) 2 3 [Brch (RuleId NUMBER 0) 2 3 [Leaf 2 Digit '2']]]
                , Leaf 3 (OneOf "*/") '*'
                , Brch (RuleId FACTOR 0) 4 5 [Brch (RuleId NUMBER 0) 4 5 [Leaf 4 Digit '3']]
                ]
            ])
      parseTreeSet aArithG SUM "1+2+3" @?= Just (tree_1_23 |: tree_12_3)
      (truncateTree 4 <$> parseTreeSet dumbG () "") @?= Just
        ( Brch (RuleId () 0) 0 0 [Brch (RuleId () 0) 0 0 [ellipsis] |: Brch (RuleId () 1) 0 0 []]
        |: Brch (RuleId () 1) 0 0 [] )
      (truncateTree 4 <$> parseTreeSet pingpongG True "") @?= Just
        ( Brch (RuleId True 0) 0 0 []
        |: Brch (RuleId True 1) 0 0 [Brch (RuleId False 0) 0 0 [ellipsis |: ellipsis]]
        )
      parseTreeSet iteG E iteInput @?= Just
        ( tree_it_ite |: tree_ite_it )
  , testCase "ambiguities" $ do
      ambiguities <$> parseTreeSet arithG SUM "1+2*3" @?= Just []
      ambiguities <$> parseTreeSet aArithG SUM "1+2+3" @?= Just
        [ (Range 5 0, Ambiguity tree_12_3 tree_1_23) ]
      ambiguities <$> parseTreeSet iteG E iteInput @?= Just
        [ (Range 9 0, Ambiguity tree_it_ite tree_ite_it) ]
  , testCase "draw" $ do
      fmap prettyTree
        (parseTreeSet arithG SUM "1+2*3" >>= fromSingleton) @?= Just
        [ "     +-----+--SUM #1---+           "
        , "     |     |           |           "
        , "  SUM #0   |      +PRODUCT #1-+    "
        , "     |     |      |     |     |    "
        , "PRODUCT #0 | PRODUCT #0 |     |    "
        , "     |     |      |     |     |    "
        , " FACTOR #0 |  FACTOR #0 | FACTOR #0"
        , "     |     |      |     |     |    "
        , " NUMBER #0 |  NUMBER #0 | NUMBER #0"
        , "     |     |      |     |     |    "
        , "-----------------------------------"
        , "     1     +      2     *     3    " ]
      fmap prettyTree
        (parseTreeSet arithG SUM "(1+2)*3" >>= fromSingleton) @?= Just
        [ "                SUM #0                 "
        , "                   |                   "
        , "             +PRODUCT #1----+-----+    "
        , "             |              |     |    "
        , "         PRODUCT #0         |     |    "
        , "             |              |     |    "
        , "+--------FACTOR #1--------+ |     |    "
        , "|            |            | |     |    "
        , "|      +---SUM #1---+     | |     |    "
        , "|      |     |      |     | |     |    "
        , "|   SUM #0   |      |     | |     |    "
        , "|      |     |      |     | |     |    "
        , "| PRODUCT #0 | PRODUCT #0 | |     |    "
        , "|      |     |      |     | |     |    "
        , "|  FACTOR #0 |  FACTOR #0 | | FACTOR #0"
        , "|      |     |      |     | |     |    "
        , "|  NUMBER #0 |  NUMBER #0 | | NUMBER #0"
        , "|      |     |      |     | |     |    "
        , "---------------------------------------"
        , "(      1     +      2     ) *     3    " ]
      prettyTree tree_it_ite @?=
        [ " +---+----+-------E #0----+             "
        , " |   |    |               |             "
        , " |   |    |   +---+----+E #1+----+---+  "
        , " |   |    |   |   |    |    |    |   |  "
        , " | E #2   |   | E #2   |  E #2   |  E #2"
        , " |   |    |   |   |    |    |    |   |  "
        , "----------------------------------------"
        , "if   T  then if   T  then   T  else  T  " ]
      prettyTree tree_ite_it @?=
        [ " +---+----+-------E #1-----------+---+  "
        , " |   |    |          |           |   |  "
        , " |   |    |   +---+-E #0----+    |   |  "
        , " |   |    |   |   |    |    |    |   |  "
        , " | E #2   |   | E #2   |  E #2   |  E #2"
        , " |   |    |   |   |    |    |    |   |  "
        , "----------------------------------------"
        , "if   T  then if   T  then   T  else  T  " ]
  ]

main :: IO ()
main = defaultMain tests