{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TupleSections    #-}

{- |
Module      : Language.Egison.ParserNonS
Copyright   : Satoshi Egi
Licence     : MIT

This module provide Egison parser.
-}

module Language.Egison.ParserNonS
       (
       -- * Parse a string
         readTopExprs
       , readTopExpr
       , readExprs
       , readExpr
       , parseTopExprs
       , parseTopExpr
       , parseExprs
       , parseExpr
       -- * Parse a file
       , loadLibraryFile
       , loadFile
       ) where

import           Control.Applicative     (pure, (*>), (<$>), (<*), (<*>))
import           Control.Monad.Except    (liftIO, throwError)
import           Control.Monad.Identity  (Identity, unless)
import           Prelude                 hiding (mapM)

import           System.Directory        (doesFileExist, getHomeDirectory)

import           Data.Char               (isLower, isUpper, toLower)
import           Data.Either
import           Data.Functor            (($>))
import           Data.List               (intercalate)
import           Data.List.Split         (split, splitOn, startsWithOneOf)
import           Data.Ratio
import           Data.Traversable        (mapM)

import           Text.Parsec
import           Text.Parsec.Expr
import           Text.Parsec.String
import qualified Text.Parsec.Token       as P

import qualified Data.Text               as T

import           Language.Egison.Desugar
import           Language.Egison.Types
import           Paths_egison            (getDataFileName)

readTopExprs :: String -> EgisonM [EgisonTopExpr]
readTopExprs = either throwError (mapM desugarTopExpr) . parseTopExprs

readTopExpr :: String -> EgisonM EgisonTopExpr
readTopExpr = either throwError desugarTopExpr . parseTopExpr

readExprs :: String -> EgisonM [EgisonExpr]
readExprs = liftEgisonM . runDesugarM . either throwError (mapM desugar) . parseExprs

readExpr :: String -> EgisonM EgisonExpr
readExpr = liftEgisonM . runDesugarM . either throwError desugar . parseExpr

parseTopExprs :: String -> Either EgisonError [EgisonTopExpr]
parseTopExprs = doParse $ do
  ret <- whiteSpace >> endBy topExpr whiteSpace
  eof
  return ret

parseTopExpr :: String -> Either EgisonError EgisonTopExpr
parseTopExpr = doParse $ do
  ret <- whiteSpace >> topExpr
  whiteSpace >> eof
  return ret

parseExprs :: String -> Either EgisonError [EgisonExpr]
parseExprs = doParse $ do
  ret <- whiteSpace >> endBy expr whiteSpace
  eof
  return ret

parseExpr :: String -> Either EgisonError EgisonExpr
parseExpr = doParse $ do
  ret <- whiteSpace >> expr
  whiteSpace >> eof
  return ret

-- |Load a libary file
loadLibraryFile :: FilePath -> EgisonM [EgisonTopExpr]
loadLibraryFile file = do
  homeDir <- liftIO getHomeDirectory
  doesExist <- liftIO $ doesFileExist $ homeDir ++ "/.egison/" ++ file
  if doesExist
    then loadFile $ homeDir ++ "/.egison/" ++ file
    else liftIO (getDataFileName file) >>= loadFile

-- |Load a file
loadFile :: FilePath -> EgisonM [EgisonTopExpr]
loadFile file = do
  doesExist <- liftIO $ doesFileExist file
  unless doesExist $ throwError $ Default ("file does not exist: " ++ file)
  input <- liftIO $ readUTF8File file
  exprs <- readTopExprs $ shebang input
  concat <$> mapM  recursiveLoad exprs
 where
  recursiveLoad (Load file)     = loadLibraryFile file
  recursiveLoad (LoadFile file) = loadFile file
  recursiveLoad expr            = return [expr]
  shebang :: String -> String
  shebang ('#':'!':cs) = ';':'#':'!':cs
  shebang cs           = cs

--
-- Parser
--

doParse :: Parser a -> String -> Either EgisonError a
doParse p input = either (throwError . fromParsecError) return $ parse p "egison" input
  where
    fromParsecError :: ParseError -> EgisonError
    fromParsecError = Parser . show

doParse' :: Parser a -> String -> a
doParse' p input = case doParse p input of
                     Right x -> x

--
-- Expressions
--

topExpr :: Parser EgisonTopExpr
topExpr = try defineExpr
          <|> try (Test <$> expr)
          <|> testExpr
          <|> loadFileExpr
          <|> loadExpr
          <?> "top-level expression"

defineExpr :: Parser EgisonTopExpr
defineExpr = try (Define <$ keywordDefine <*> identVar <*> (LambdaExpr <$> parens argNames' <* inSpaces (reservedOp "=") <* notFollowedBy (string "=") <*> expr))
             <|> try (Define <$> identVar <* inSpaces (reservedOp "=") <* notFollowedBy (string "=") <*> expr)
             <|> try (do (VarWithIndices name is) <- identVarWithIndices
                         inSpaces $ reservedOp "=" >> notFollowedBy (string "=")
                         Define (Var name (map f is)) . WithSymbolsExpr (map g is) . TransposeExpr (CollectionExpr (map (ElementExpr . VarExpr . stringToVar . g) is)) <$> expr)
 where
  argNames' :: Parser [Arg]
  argNames' = sepEndBy argName' comma
  argName' :: Parser Arg
  argName' = try (ScalarArg <$> ident)
        <|> try (InvertedScalarArg <$> (char '*' >> ident))
        <|> try (TensorArg <$> (char '%' >> ident))
  f (Superscript _)  = Superscript ()
  f (Subscript _)    = Subscript ()
  f (SupSubscript _) = SupSubscript ()
  g (Superscript i)  = i
  g (Subscript i)    = i
  g (SupSubscript i) = i

testExpr :: Parser EgisonTopExpr
testExpr = keywordTest >> Test <$> parens expr

loadFileExpr :: Parser EgisonTopExpr
loadFileExpr = keywordLoadFile >> LoadFile <$> parens stringLiteral

loadExpr :: Parser EgisonTopExpr
loadExpr = keywordLoad >> Load <$> parens stringLiteral

expr :: Parser EgisonExpr
expr = (try applyInfixExpr
        <|> try exprWithSymbol
        <|> try (buildExpressionParser table arg)
        <|> try ifExpr
        <|> try term)
        <?> "expression"
 where
  arg = (char '$' *> notFollowedBy varExpr *> (LambdaArgExpr <$> option "" index))
        <|> term
  index = (:) <$> satisfy (\c -> '1' <= c && c <= '9') <*> many digit
  table = [ [unary "not" AssocRight]
          , [binary "^" "**" AssocLeft]
          , [unary "-" AssocLeft]
          , [binary "*" "*" AssocLeft, binary "/" "/" AssocLeft, binary "." "." AssocLeft]
          , [binary "+" "+" AssocLeft, binary "-" "-" AssocLeft, binary "%" "remainder" AssocLeft]
          , [binary "==" "eq?" AssocLeft, binary "<=" "lte?" AssocLeft, binary "<" "lt?" AssocLeft, binary ">=" "gte?" AssocLeft, binary ">" "gt?" AssocLeft]
          , [binary ":" "cons" AssocLeft, binary ".." "between" AssocLeft]
          , [binary "&&" "and" AssocLeft, binary "||" "or" AssocLeft]
          , [binary "++" "join" AssocRight]
          ]
  unary "-" assoc = Prefix (try $ inSpaces (string "-") >> return (\x -> makeApply (stringToVarExpr "*") [IntegerExpr (-1), x]))
  unary op assoc = Prefix (try $ inSpaces (string op) >> return (\x -> makeApply (stringToVarExpr op) [x]))
  binary op name assoc
    | op == "/" = Infix (try $ (try (inSpaces1 $ string op) <|> (inSpaces (string op) >> notFollowedBy (string "m" <|> string "fn"))) >> return (\x y -> makeApply (stringToVarExpr name) [x, y])) assoc
    | op == "." || op == "%" = Infix (try $ inSpaces1 (string op) >> return (\x y -> makeApply (stringToVarExpr name) [x, y])) assoc
    | otherwise = Infix (try $ inSpaces (string op) >> return (\x y -> makeApply (stringToVarExpr name) [x, y])) assoc

inSpaces :: Parser a -> Parser ()
inSpaces p = skipMany (space <|> newline) >> p >> skipMany (space <|> newline)

inSpaces1 :: Parser a -> Parser ()
inSpaces1 p = skipMany (space <|> newline) >> p >> skipMany1 (space <|> newline)

exprWithSymbol :: Parser EgisonExpr
exprWithSymbol = (string "d/d" >> applyExpr'' (stringToVarExpr "d/d"))
                 <|> (string "V.*" >> applyExpr'' (stringToVarExpr "V.*"))
                 <|> (string "M.*" >> applyExpr'' (stringToVarExpr "M.*"))
                 <|> (lookAhead (string "let*") >> letStarExpr)

term :: Parser EgisonExpr
term = P.lexeme lexer
        (do term0 <- term'
            option term0 $ try (IndexedExpr False term0 <$ string "..." <*> parseindex
                           <|> IndexedExpr True term0 <$> parseindex))
 where
  parseindex :: Parser [Index EgisonExpr]
  parseindex = many1 $ try (MultiSubscript <$ char '_' <*> term' <* string "..._" <*> term')
                       <|> try (MultiSuperscript <$ char '~' <*> term' <* string "...~" <*> term')
                       <|> try (char '_' >> Subscript <$> term')
                       <|> try (char '~' >> Superscript <$> term')
                       <|> try (string "~_" >> SupSubscript <$> term')
                       <|> try (char '|' >> Userscript <$> term')

term' :: Parser EgisonExpr
term' = matchExpr
        <|> matchAllExpr
        <|> matchAllDFSExpr
        <|> matchLambdaExpr
        <|> matchAllLambdaExpr
        <|> matcherExpr
        <|> functionWithArgExpr
        <|> userrefsExpr
        <|> algebraicDataMatcherExpr
        <|> try applyExpr
        <|> cApplyExpr
        <|> try partialExpr
        <|> try partialVarExpr
        <|> try constantExpr
        <|> try freshVarExpr
        <|> try lambdaExpr
        <|> try cambdaExpr
        <|> try withSymbolsExpr
        <|> try varExpr
        <|> try vectorExpr
        <|> try tupleExpr
        <|> try hashExpr
        <|> try collectionExpr
        <|> inductiveDataExpr
        <|> try doExpr
        <|> generateTensorExpr
        <|> tensorExpr
        <|> letExpr
        <|> letRecExpr
        <|> letStarExpr
        <|> patternFunctionExpr
        <|> quoteExpr
        <|> quoteSymbolExpr
        <|> tensorContractExpr
        <|> subrefsExpr
        <|> suprefsExpr
        <|> macroExpr
        <|> ioExpr
        <|> seqExpr
        <|> memoizedLambdaExpr
        <|> procedureExpr
        <|> wedgeExpr
        <|> parens expr
        <?> "simple expression"

varExpr :: Parser EgisonExpr
varExpr = VarExpr <$> identVarWithoutIndex

freshVarExpr :: Parser EgisonExpr
freshVarExpr = char '#' >> return FreshVarExpr

inductiveDataExpr :: Parser EgisonExpr
inductiveDataExpr = angles $ InductiveDataExpr <$> upperName <*> sepEndBy term whiteSpace

tupleExpr :: Parser EgisonExpr
tupleExpr = parens $ TupleExpr <$> sepEndBy expr comma

collectionExpr :: Parser EgisonExpr
collectionExpr = brackets (CollectionExpr <$> sepEndBy innerExpr comma)
                 <|> braces (CollectionExpr <$> sepEndBy innerExpr comma)
 where
  innerExpr :: Parser InnerExpr
  innerExpr = (char '@' >> SubCollectionExpr <$> expr)
               <|> ElementExpr <$> expr

vectorExpr :: Parser EgisonExpr
vectorExpr = between lp rp $ VectorExpr <$> sepEndBy expr comma
  where
    lp = P.lexeme lexer (string "[|")
    rp = string "|]"

hashExpr :: Parser EgisonExpr
hashExpr = between lp rp $ HashExpr <$> sepEndBy pairExpr comma
  where
    lp = P.lexeme lexer (string "{|")
    rp = string "|}"
    pairExpr :: Parser (EgisonExpr, EgisonExpr)
    pairExpr = brackets $ (,) <$> expr <* comma <*> expr

quoteExpr :: Parser EgisonExpr
quoteExpr = char '\'' >> QuoteExpr <$> expr

wedgeExpr :: Parser EgisonExpr
wedgeExpr = char '!' >> WedgeExpr <$> expr

functionWithArgExpr :: Parser EgisonExpr
functionWithArgExpr = keywordFunction >> FunctionExpr <$> parens (sepEndBy expr comma)

quoteSymbolExpr :: Parser EgisonExpr
quoteSymbolExpr = char '`' >> QuoteSymbolExpr <$> expr

matchAllExpr :: Parser EgisonExpr
matchAllExpr = keywordMatchAll >> MatchAllExpr <$> expr <* keywordAs <*> expr <*> matchClauses

matchAllDFSExpr :: Parser EgisonExpr
matchAllDFSExpr = keywordMatchAllDFS >> MatchAllDFSExpr <$> expr <* keywordAs <*> expr <*> matchClauses

matchExpr :: Parser EgisonExpr
matchExpr = keywordMatch >> MatchExpr <$> expr <* keywordAs <*> expr <*> matchClauses

matchLambdaExpr :: Parser EgisonExpr
matchLambdaExpr = keywordMatchLambda >> MatchLambdaExpr <$ keywordAs <*> expr <*> matchClauses

matchAllLambdaExpr :: Parser EgisonExpr
matchAllLambdaExpr = keywordMatchAllLambda >> MatchAllLambdaExpr <$ keywordAs <*> expr <*> matchClauses

matchClauses :: Parser [MatchClause]
matchClauses = many1 matchClause

matchClause :: Parser MatchClause
matchClause = try $ inSpaces (string "|") >> (,) <$> pattern <* reservedOp "->" <*> expr

matcherExpr :: Parser EgisonExpr
matcherExpr = keywordMatcher >> MatcherExpr <$> ppMatchClauses

ppMatchClauses :: Parser MatcherInfo
ppMatchClauses = many1 ppMatchClause

ppMatchClause :: Parser (PrimitivePatPattern, EgisonExpr, [(PrimitiveDataPattern, EgisonExpr)])
ppMatchClause = inSpaces (string "|") >> (,,) <$> ppPattern <* keywordAs <*> expr <* reservedOp "->" <*> pdMatchClauses

pdMatchClauses :: Parser [(PrimitiveDataPattern, EgisonExpr)]
pdMatchClauses = many1 pdMatchClause

pdMatchClause :: Parser (PrimitiveDataPattern, EgisonExpr)
pdMatchClause = try $ inSpaces (string "|") >> (,) <$> pdPattern <* reservedOp "->" <*> expr

ppPattern :: Parser PrimitivePatPattern
ppPattern = P.lexeme lexer (ppWildCard
                        <|> try ppValuePat
                        <|> ppPatVar
                        <|> ppInductivePat
                        <?> "primitive-pattren-pattern")

ppWildCard :: Parser PrimitivePatPattern
ppWildCard = reservedOp "_" $> PPWildCard

ppPatVar :: Parser PrimitivePatPattern
ppPatVar = reservedOp "$" $> PPPatVar

ppValuePat :: Parser PrimitivePatPattern
ppValuePat = reservedOp "$" >> PPValuePat <$> ident

ppInductivePat :: Parser PrimitivePatPattern
ppInductivePat = angles (PPInductivePat <$> lowerName <*> sepEndBy ppPattern whiteSpace)

pdPattern :: Parser PrimitiveDataPattern
pdPattern = P.lexeme lexer pdPattern'

pdPattern' :: Parser PrimitiveDataPattern
pdPattern' = reservedOp "_" $> PDWildCard
                    <|> (char '$' >> PDPatVar <$> ident)
                    <|> brackets ((PDConsPat <$> pdPattern <* comma <*> (char '@' *> pdPattern))
                            <|> (PDSnocPat <$> (char '@' *> pdPattern) <* comma <*> pdPattern)
                            <|> pure PDEmptyPat)
                    <|> angles (PDInductivePat <$> upperName <*> sepEndBy pdPattern whiteSpace)
                    <|> parens (PDTuplePat <$> sepEndBy pdPattern comma)
                    <|> PDConstantPat <$> constantExpr
                    <?> "primitive-data-pattern"

ifExpr :: Parser EgisonExpr
ifExpr = keywordIf >> IfExpr <$> expr <* keywordThen <*> expr <* keywordElse <*> expr

lambdaExpr :: Parser EgisonExpr
lambdaExpr = LambdaExpr <$> argNames <* reservedOp "->" <*> expr

memoizedLambdaExpr :: Parser EgisonExpr
memoizedLambdaExpr = keywordMemoizedLambda >> MemoizedLambdaExpr <$> varNames <* reservedOp "->" <*> expr

memoizeFrame :: Parser [(EgisonExpr, EgisonExpr, EgisonExpr)]
memoizeFrame = braces $ sepEndBy memoizeBinding whiteSpace

memoizeBinding :: Parser (EgisonExpr, EgisonExpr, EgisonExpr)
memoizeBinding = brackets $ (,,) <$> expr <*> expr <*> expr

cambdaExpr :: Parser EgisonExpr
cambdaExpr = keywordCambda >> char '$' >> CambdaExpr <$> ident <* reservedOp "->" <*> expr

procedureExpr :: Parser EgisonExpr
procedureExpr = keywordProcedure >> ProcedureExpr <$> varNames <* reservedOp "->" <*> expr

macroExpr :: Parser EgisonExpr
macroExpr = keywordMacro >> MacroExpr <$> varNames <* reservedOp "->" <*> expr

patternFunctionExpr :: Parser EgisonExpr
patternFunctionExpr = keywordPatternFunction >> parens (PatternFunctionExpr <$> brackets (sepEndBy ident comma) <* comma <*> pattern)

letRecExpr :: Parser EgisonExpr
letRecExpr =  keywordLetRec >> LetRecExpr <$> bindings <* keywordLetIn <*> expr

letExpr :: Parser EgisonExpr
letExpr = keywordLet >> LetExpr <$> bindings <* keywordLetIn <*> expr

letStarExpr :: Parser EgisonExpr
letStarExpr = keywordLetStar >> LetStarExpr <$> bindings <* keywordLetIn <*> expr

withSymbolsExpr :: Parser EgisonExpr
withSymbolsExpr = keywordWithSymbols >> WithSymbolsExpr <$> braces (sepEndBy ident comma) <*> expr

doExpr :: Parser EgisonExpr
doExpr = keywordDo >> DoExpr <$> statements <*> option (ApplyExpr (stringToVarExpr "return") (TupleExpr [])) expr

statements :: Parser [BindingExpr]
statements = braces $ sepEndBy statement comma

statement :: Parser BindingExpr
statement = try binding
        <|> (([],) <$> expr)

bindings :: Parser [BindingExpr]
bindings = sepEndBy binding comma

binding :: Parser BindingExpr
binding = (,) <$> varNames' <* inSpaces (string "=") <*> expr

varNames :: Parser [String]
varNames = sepEndBy (char '$' >> ident) whiteSpace

varNames' :: Parser [Var]
varNames' = return <$> identVar
            <|> parens (sepEndBy identVar comma)

argNames :: Parser [Arg]
argNames = sepEndBy argName whiteSpace

argName :: Parser Arg
argName = try (ScalarArg <$> (char '$' >> ident))
      <|> try (InvertedScalarArg <$> (string "*$" >> ident))
      <|> try (TensorArg <$> (char '%' >> ident))

ioExpr :: Parser EgisonExpr
ioExpr = keywordIo >> parens (IoExpr <$> expr)

seqExpr :: Parser EgisonExpr
seqExpr = keywordSeq >> parens (SeqExpr <$> expr <* comma <*> expr)

cApplyExpr :: Parser EgisonExpr
cApplyExpr = keywordCApply >> parens (CApplyExpr <$> expr <* comma <*> expr)

applyExpr :: Parser EgisonExpr
applyExpr = (keywordApply >> parens (ApplyExpr <$> expr <* comma <*> expr))
            <|> try applyExpr'

applyExpr' :: Parser EgisonExpr
applyExpr' = do
  func <- try varExpr <|> try partialExpr <|> try partialVarExpr <|> parens expr
  applyExpr'' func

applyExpr'' :: EgisonExpr -> Parser EgisonExpr
applyExpr'' func = do
  argslist <- many1 $ parens args
  return $ foldl makeApply func argslist
 where
  args = sepEndBy arg comma
  arg = try expr
        <|> char '$' *> (LambdaArgExpr <$> option "" index)
  index = (:) <$> satisfy (\c -> '1' <= c && c <= '9') <*> many digit

applyInfixExpr :: Parser EgisonExpr
applyInfixExpr = do
  arg1 <- arg
  spaces
  func <- char '`' *> varExpr <* char '`'
  spaces
  arg2 <- arg
  return $ makeApply func [arg1, arg2]
 where
  arg = try term
        <|> char '$' *> (LambdaArgExpr <$> option "" index)
  index = (:) <$> satisfy (\c -> '1' <= c && c <= '9') <*> many digit

makeApply :: EgisonExpr -> [EgisonExpr] -> EgisonExpr
makeApply func xs = do
  let args = map (\x -> case x of
                          LambdaArgExpr s -> Left s
                          _               -> Right x) xs
  let vars = lefts args
  case vars of
    [] -> ApplyExpr func . TupleExpr $ rights args
    _ | all null vars ->
        let args' = rights args
            args'' = zipWith (curry f) args (annonVars 1 (length args))
            args''' = map (VarExpr . stringToVar . either id id) args''
        in ApplyExpr (LambdaExpr (map ScalarArg (rights args'')) (LambdaExpr (map ScalarArg (lefts args'')) $ ApplyExpr func $ TupleExpr args''')) $ TupleExpr args'
      | all (not . null) vars ->
        let n = length vars
            args' = rights args
            args'' = zipWith (curry g) args (annonVars (n + 1) (length args))
            args''' = map (VarExpr . stringToVar . either id id) args''
        in ApplyExpr (LambdaExpr (map ScalarArg (rights args'')) (LambdaExpr (map ScalarArg (annonVars 1 n)) $ ApplyExpr func $ TupleExpr args''')) $ TupleExpr args'
 where
  annonVars m n = take n $ map ((':':) . show) [m..]
  f (Left _, var)  = Left var
  f (Right _, var) = Right var
  g (Left arg, _)  = Left (':':arg)
  g (Right _, var) = Right var

partialExpr :: Parser EgisonExpr
partialExpr = (PartialExpr . read <$> index) <*> (char '#' >> (try (parens expr) <|> expr))
 where
  index = (:) <$> satisfy (\c -> '1' <= c && c <= '9') <*> many digit

partialVarExpr :: Parser EgisonExpr
partialVarExpr = char '%' >> PartialVarExpr <$> integerLiteral

algebraicDataMatcherExpr :: Parser EgisonExpr
algebraicDataMatcherExpr = keywordAlgebraicDataMatcher
                                >> AlgebraicDataMatcherExpr <$> parens (sepEndBy1 inductivePat' comma)
  where
    inductivePat' :: Parser (String, [EgisonExpr])
    inductivePat' = angles $ (,) <$> lowerName <*> sepEndBy expr whiteSpace

generateTensorExpr :: Parser EgisonExpr
generateTensorExpr = keywordGenerateTensor >> parens (GenerateTensorExpr <$> expr <* comma <*> expr)

tensorExpr :: Parser EgisonExpr
tensorExpr = keywordTensor >> parens (TensorExpr <$> expr <* comma <*> expr <*> option (CollectionExpr []) (comma *> expr) <*> option (CollectionExpr []) (comma *> expr))

tensorContractExpr :: Parser EgisonExpr
tensorContractExpr = keywordTensorContract >> parens (TensorContractExpr <$> expr <* comma <*> expr)

subrefsExpr :: Parser EgisonExpr
subrefsExpr = (keywordSubrefs >> parens (SubrefsExpr False <$> expr <* comma <*> expr))
               <|> (keywordSubrefsNew >> parens (SubrefsExpr True <$> expr <* comma <*> expr))

suprefsExpr :: Parser EgisonExpr
suprefsExpr = (keywordSuprefs >> SuprefsExpr False <$> expr <*> expr)
               <|> (keywordSuprefsNew >> SuprefsExpr True <$> expr <*> expr)

userrefsExpr :: Parser EgisonExpr
userrefsExpr = (do keywordUserrefs
                   xs <- parens $ sepEndBy expr comma
                   case xs of
                     [x, y] -> return $ UserrefsExpr False x y
                     _      -> unexpected "number of arguments (expected 2)")
                <|> (do keywordUserrefsNew
                        xs <- parens $ sepEndBy expr comma
                        case xs of
                          [x, y] -> return $ UserrefsExpr True x y
                          _ -> unexpected "number of arguments (expected 2)")

-- Patterns

pattern :: Parser EgisonPattern
pattern = P.lexeme lexer
            (try (buildExpressionParser table pattern')
             <|> try pattern'
             <?> "expression")

 where
  table = [ [unary "!" AssocRight, unary "not" AssocRight]
          , [binary'' "^" PowerPat AssocLeft]
          , [binary' "*" MultPat AssocRight, binary'' "/" DivPat AssocRight]
          , [binary' "+" PlusPat AssocRight]
          , [binary "<:>" "cons" AssocRight]
          , [binary' "and" AndPat AssocLeft, binary' "or" OrPat AssocLeft]
          , [binary "<++>" "join" AssocRight]
          ]
  unary op assoc = Prefix (try $ inSpaces (string op) >> return NotPat)
  binary op name = Infix (try $ inSpaces (string op) >> return (\x y -> InductivePat name [x, y]))
  binary' op epr = Infix (try $ inSpaces (string op) >> return (\x y -> epr [x, y]))
  binary'' op epr = Infix (try $ inSpaces (string op) >> return epr)

pattern' :: Parser EgisonPattern
pattern' = wildCard
           <|> contPat
           <|> try indexedPat
           <|> patVar
           <|> try loopPat
           <|> try pApplyPat
           <|> try dApplyPat
           <|> try varPat
           <|> valuePat
           <|> predPat
           <|> try tuplePat
           <|> inductivePat
           <|> letPat
           <|> parens pattern

pattern'' :: Parser EgisonPattern
pattern'' = wildCard
            <|> patVar
            <|> valuePat

wildCard :: Parser EgisonPattern
wildCard = reservedOp "_" >> pure WildCard

indexedPat :: Parser EgisonPattern
indexedPat = IndexedPat <$> (patVar <|> varPat) <*> many1 (try $ char '_' >> term')

patVar :: Parser EgisonPattern
patVar = char '$' >> PatVar <$> identVarWithoutIndex

varPat :: Parser EgisonPattern
varPat = char '\'' >> VarPat <$> ident

valuePat :: Parser EgisonPattern
valuePat = ValuePat <$> expr

predPat :: Parser EgisonPattern
predPat = char '?' >> PredPat <$> expr

letPat :: Parser EgisonPattern
letPat = keywordLet >> LetPat <$> bindings <* keywordLetIn <*> pattern

tuplePat :: Parser EgisonPattern
tuplePat = TuplePat <$> parens ((:) <$> pattern <* comma <*> sepEndBy1 pattern comma)

inductivePat :: Parser EgisonPattern
inductivePat = angles $ InductivePat <$> lowerName <*> sepEndBy pattern whiteSpace

contPat :: Parser EgisonPattern
contPat = keywordCont >> pure ContPat

pApplyPat :: Parser EgisonPattern
pApplyPat = PApplyPat <$> expr <*> brackets (sepEndBy pattern comma)

dApplyPat :: Parser EgisonPattern
dApplyPat = DApplyPat <$> pattern'' <*> parens (sepEndBy pattern comma)

loopPat :: Parser EgisonPattern
loopPat = keywordLoop >> parens (char '$' >> LoopPat <$> identVarWithoutIndex <*> (comma >> loopRange) <*> (comma >> pattern) <*> (comma >> option (NotPat WildCard) pattern))

loopRange :: Parser LoopRange
loopRange = parens (try (LoopRange <$> expr <* comma <*> expr <*> option WildCard (comma >> pattern))
                 <|> (do s <- expr
                         comma
                         ep <- option WildCard pattern
                         return (LoopRange s (ApplyExpr (stringToVarExpr "from") (ApplyExpr (stringToVarExpr "-'") (TupleExpr [s, IntegerExpr 1]))) ep)))

-- Constants

constantExpr :: Parser EgisonExpr
constantExpr = stringExpr
                 <|> boolExpr
                 <|> try charExpr
                 <|> try floatExpr
                 <|> try integerExpr
                 <|> (keywordSomething $> SomethingExpr)
                 <|> (keywordUndefined $> UndefinedExpr)
                 <?> "constant"

charExpr :: Parser EgisonExpr
charExpr = CharExpr <$> oneChar

stringExpr :: Parser EgisonExpr
stringExpr = StringExpr . T.pack <$> stringLiteral

boolExpr :: Parser EgisonExpr
boolExpr = BoolExpr <$> boolLiteral

floatExpr :: Parser EgisonExpr
floatExpr = do
  (x,y) <- try ((,) <$> floatLiteral <*> (sign' <*> positiveFloatLiteral) <* char 'i')
            <|> try ((,) 0 <$> floatLiteral <* char 'i')
            <|> try ((, 0) <$> floatLiteral)
  return $ FloatExpr x y

integerExpr :: Parser EgisonExpr
integerExpr = IntegerExpr <$> integerLiteral'

integerLiteral' :: Parser Integer
integerLiteral' = sign <*> positiveIntegerLiteral

positiveIntegerLiteral :: Parser Integer
positiveIntegerLiteral = read <$> many1 digit

positiveFloatLiteral :: Parser Double
positiveFloatLiteral = do
  n <- positiveIntegerLiteral
  char '.'
  mStr <- many1 digit
  let m = read mStr
  let l = m % (10 ^ fromIntegral (length mStr))
  return (fromRational (fromIntegral n + l) :: Double)

floatLiteral :: Parser Double
floatLiteral = sign <*> positiveFloatLiteral

--
-- Tokens
--

egisonDef :: P.GenLanguageDef String () Identity
egisonDef =
  P.LanguageDef { P.commentStart       = "#|"
                , P.commentEnd         = "|#"
                , P.commentLine        = ";"
                , P.identStart         = letter <|> symbol1
                , P.identLetter        = letter <|> digit <|> symbol0 <|> symbol2
                , P.opStart            = symbol1
                , P.opLetter           = symbol0 <|> symbol1
                , P.reservedNames      = reservedKeywords
                , P.reservedOpNames    = reservedOperators
                , P.nestedComments     = True
                , P.caseSensitive      = True }

symbol0 = oneOf "/."
symbol1' = oneOf "∂∇"
symbol1 = symbol1' <|> oneOf "+-"
symbol2 = symbol1' <|> oneOf "'!?"

lexer :: P.GenTokenParser String () Identity
lexer = P.makeTokenParser egisonDef

reservedKeywords :: [String]
reservedKeywords =
  [ "define"
  , "set!"
  , "test"
  , "loadFile"
  , "load"
  , "if"
  , "then"
  , "else"
  , "as"
  , "seq"
  , "apply"
  , "capply"
  , "lambda"
  , "memoizedLambda"
  , "cambda"
  , "procedure"
  , "macro"
  , "patternFunction"
  , "letrec"
  , "let"
  , "let*"
  , "in"
  , "withSymbols"
  , "loop"
  , "matchAll"
  , "matchAllDFS"
  , "matchAllLambda"
  , "match"
  , "matchLambda"
  , "matcher"
  , "do"
  , "io"
  , "something"
  , "undefined"
  , "algebraicDataMatcher"
  , "generateTensor"
  , "tensor"
  , "contract"
  , "subrefs"
  , "subrefs!"
  , "suprefs"
  , "suprefs!"
  , "userRefs"
  , "userRefs!"
  , "function"]

reservedOperators :: [String]
reservedOperators =
  [ "$"
  , "_"
  , "^"
  , "&"
  , "|*"
  , "("
  , ")"
  , "->"
  , "`"
  , "=="
  , "="
--  , "'"
--  , "~"
--  , "!"
--  , ","
--  , "@"
  , "..."]

reserved :: String -> Parser ()
reserved = P.reserved lexer

reservedOp :: String -> Parser ()
reservedOp = P.reservedOp lexer

keywordDefine               = reserved "def"
keywordSet                  = reserved "set!"
keywordTest                 = reserved "test"
keywordLoadFile             = reserved "loadFile"
keywordLoad                 = reserved "load"
keywordIf                   = reserved "if"
keywordThen                 = reserved "then"
keywordElse                 = reserved "else"
keywordAs                   = reserved "as"
keywordSeq                  = reserved "seq"
keywordApply                = reserved "apply"
keywordCApply               = reserved "capply"
keywordLambda               = reserved "lambda"
keywordMemoizedLambda       = reserved "memoizedLambda"
keywordCambda               = reserved "cambda"
keywordProcedure            = reserved "procedure"
keywordMacro                = reserved "macro"
keywordPatternFunction      = reserved "patternFunction"
keywordLetRec               = reserved "letrec"
keywordLet                  = reserved "let"
keywordLetStar              = reserved "let*"
keywordLetIn                = reserved "in"
keywordWithSymbols          = reserved "withSymbols"
keywordLoop                 = reserved "loop"
keywordCont                 = reserved "..."
keywordMatchAll             = reserved "matchAll"
keywordMatchAllDFS          = reserved "matchAllDFS"
keywordMatchAllLambda       = reserved "matchAllLambda"
keywordMatch                = reserved "match"
keywordMatchLambda          = reserved "matchLambda"
keywordMatcher              = reserved "matcher"
keywordDo                   = reserved "do"
keywordIo                   = reserved "io"
keywordSomething            = reserved "something"
keywordUndefined            = reserved "undefined"
keywordAlgebraicDataMatcher = reserved "algebraicDataMatcher"
keywordGenerateTensor       = reserved "generateTensor"
keywordTensor               = reserved "tensor"
keywordTensorContract       = reserved "contract"
keywordSubrefs              = reserved "subrefs"
keywordSubrefsNew           = reserved "subrefs!"
keywordSuprefs              = reserved "suprefs"
keywordSuprefsNew           = reserved "suprefs!"
keywordUserrefs             = reserved "userRefs"
keywordUserrefsNew          = reserved "userRefs!"
keywordFunction             = reserved "function"

sign :: Num a => Parser (a -> a)
sign = (char '-' >> return negate)
   <|> (char '+' >> return id)
   <|> return id

sign' :: Num a => Parser (a -> a)
sign' = (char '-' >> return negate)
    <|> (char '+' >> return id)

integerLiteral :: Parser Integer
integerLiteral = sign <*> P.natural lexer

stringLiteral :: Parser String
stringLiteral = P.stringLiteral lexer

charLiteral :: Parser Char
charLiteral = P.charLiteral lexer

oneChar :: Parser Char
oneChar = do
  string "c#"
  x <- (char '\\' >> anyChar >>= (\x -> return ['\\', x])) <|> (anyChar >>= (\x -> return [x]))
  return $ doParse' charLiteral $ "'" ++ x ++ "'"

boolLiteral :: Parser Bool
boolLiteral = char '#' >> (char 't' $> True <|> char 'f' $> False)

whiteSpace :: Parser ()
whiteSpace = P.whiteSpace lexer

parens :: Parser a -> Parser a
parens = P.parens lexer

brackets :: Parser a -> Parser a
brackets = P.brackets lexer

braces :: Parser a -> Parser a
braces = P.braces lexer

angles :: Parser a -> Parser a
angles = P.angles lexer

colon :: Parser String
colon = P.colon lexer

comma :: Parser String
comma = P.comma lexer

dot :: Parser String
dot = P.dot lexer

ident :: Parser String
ident = do
  idt <- P.identifier lexer
  let (f, s) = splitLast idt '.'
  case s of
    [] -> return f
    x:xs | isLower x -> return $ f ++ map toLower (intercalate "-" $ split (startsWithOneOf ['A'..'Z']) s)
         | otherwise -> return $ f ++ [x] ++ map toLower (intercalate "-" $ split (startsWithOneOf ['A'..'Z']) xs)
 where
  splitLast list elem = let (f, s) = span (/= elem) $ reverse list
                          in (reverse s, reverse f)

identVar :: Parser Var
identVar = P.lexeme lexer (do
  name <- ident
  is <- many indexType
  return $ Var (splitOn "." name) is)

identVarWithoutIndex :: Parser Var
identVarWithoutIndex = stringToVar <$> ident

identVarWithIndices :: Parser VarWithIndices
identVarWithIndices = P.lexeme lexer (do
  name <- ident
  is <- many indexForVar
  return $ VarWithIndices (splitOn "." name) is)

indexForVar :: Parser (Index String)
indexForVar = try (char '~' >> Superscript <$> ident)
        <|> try (char '_' >> Subscript <$> ident)

indexType :: Parser (Index ())
indexType = try (char '~' >> return (Superscript ()))
        <|> try (char '_' >> return (Subscript ()))

upperName :: Parser String
upperName = P.lexeme lexer upperName'

upperName' :: Parser String
upperName' = (:) <$> upper <*> option "" ident
 where
  upper :: Parser Char
  upper = satisfy isUpper

lowerName :: Parser String
lowerName = P.lexeme lexer lowerName'

lowerName' :: Parser String
lowerName' = (:) <$> lower <*> option "" ident
 where
  lower :: Parser Char
  lower = satisfy isLower