-- | Parse a LambdaPi program. Adapted from source code accompanying _A tutorial
--   implementation of a dependently typed lambda calculus_, by Löh, McBride and
--   Swierstra. See that paper for details.
module Parser (parseStr, ITerm(..), CTerm(..), Name(..), Program) where
  import Data.List (findIndex)
  import Data.Map (Map, fromList)
 
  import Text.PrettyPrint.HughesPJ hiding (parens)
  import qualified Text.PrettyPrint.HughesPJ as PP
  import Text.ParserCombinators.Parsec hiding (parse, State)
  import qualified Text.ParserCombinators.Parsec as P
  import Text.ParserCombinators.Parsec.Token
  import Text.ParserCombinators.Parsec.Language
  
  type Program = [(String,ITerm)]
  
  parseStr :: String -> String -> Either String Program
  parseStr moduleName str =
    case parseIO moduleName str of
        Left err -> Left $ show err
        Right prog -> Right prog

  parseIO f x = P.parse (whiteSpace lambdaPi >> (many $ parseDefinition) >>= \ x -> eof >> return x) f x
  
  data CTerm
     =  Inf  ITerm
     |  Lam  CTerm
     |  Zero
     |  Succ CTerm
     |  Nil CTerm
     |  Cons CTerm CTerm CTerm CTerm
     |  Refl CTerm CTerm
     |  FZero CTerm
     |  FSucc CTerm CTerm
     deriving (Show, Eq)
  
  data ITerm
     =  Ann CTerm CTerm
     |  Star
     |  Pi CTerm CTerm
     |  Bound  Int
     |  Free  Name
     |  ITerm :$: CTerm
     |  Nat
     |  NatElim CTerm CTerm CTerm CTerm
     |  Vec CTerm CTerm
     |  VecElim CTerm CTerm CTerm CTerm CTerm CTerm
     |  Eq CTerm CTerm CTerm
     |  EqElim CTerm CTerm CTerm CTerm CTerm CTerm
     |  Fin CTerm
     |  FinElim CTerm CTerm CTerm CTerm CTerm
     deriving (Show, Eq)
  data Type
     =  TFree  Name
     |  Fun    Type Type
    deriving (Show, Eq)
    
  data Name
       =  Global  String
       |  Local   Int
       |  Quote   Int
      deriving (Show, Eq)
     
  data Stmt i tinf = Let String i           --  let x = t
                   | Assume [(String,tinf)] --  assume x :: t, assume x :: *
                   | Eval i
                   | PutStrLn String        --  lhs2TeX hacking, allow to print "magic" string
                   | Out String             --  more lhs2TeX hacking, allow to print to files
                   deriving (Show)

  
  lambdaPi = makeTokenParser (haskellStyle { identStart = letter <|> P.char '_',
                                             reservedNames = ["forall", "let", "assume", "putStrLn", "out"] })

  parseDefinition = do
    reserved lambdaPi "let"
    x <- identifier lambdaPi
    reserved lambdaPi "="
    t <- parseITerm 0 []
    return (x,t)
  
  parseStmt :: [String] -> CharParser () (Stmt ITerm CTerm)
  parseStmt e =
        do
          reserved lambdaPi "let"
          x <- identifier lambdaPi
          reserved lambdaPi "="
          t <- parseITerm 0 e
          return (Let x t)
    <|> do
          reserved lambdaPi "assume"
          (xs, ts) <- parseBindings False [] 
          return (Assume (reverse (zip xs ts)))
    <|> do
          reserved lambdaPi "putStrLn"
          x <- stringLiteral lambdaPi
          return (PutStrLn x)
    <|> do
          reserved lambdaPi "out"
          x <- option "" (stringLiteral lambdaPi)
          return (Out x)
    <|> fmap Eval (parseITerm 0 e)

  parseBindings :: Bool -> [String] -> CharParser () ([String], [CTerm])
  parseBindings b e = 
                     (let rec :: [String] -> [CTerm] -> CharParser () ([String], [CTerm])
                          rec e ts =
                            do
                             (x,t) <- parens lambdaPi
                                        (do
                                           x <- identifier lambdaPi
                                           reserved lambdaPi "::"
                                           t <- parseCTerm 0 (if b then e else [])
                                           return (x,t))
                             (rec (x : e) (t : ts) <|> return (x : e, t : ts))
                      in rec e [])
                     <|>
                     do  x <- identifier lambdaPi
                         reserved lambdaPi "::"
                         t <- parseCTerm 0 e
                         return (x : e, [t])

  parseITerm :: Int -> [String] -> CharParser () ITerm
  parseITerm 0 e =
        do
          reserved lambdaPi "forall"
          (fe,t:ts) <- parseBindings True e
          reserved lambdaPi "."
          t' <- parseCTerm 0 fe
          return (foldl (\ p t -> Pi t (Inf p)) (Pi t t') ts)
    <|>
    try
       (do
          t <- parseITerm 1 e
          rest (Inf t) <|> return t)
    <|> do
          t <- parens lambdaPi (parseLam e)
          rest t
    where
      rest t =
        do
          reserved lambdaPi "->"
          t' <- parseCTerm 0 ([]:e)
          return (Pi t t')
  parseITerm 1 e =
    try
       (do
          t <- parseITerm 2 e
          rest (Inf t) <|> return t)
    <|> do
          t <- parens lambdaPi (parseLam e)
          rest t
    where
      rest t =
        do
          reserved lambdaPi "::"
          t' <- parseCTerm 0 e
          return (Ann t t')
  parseITerm 2 e =
        do
          t <- parseITerm 3 e
          ts <- many (parseCTerm 3 e)
          return (foldl (:$:) t ts)
  parseITerm 3 e =
        do
          reserved lambdaPi "*"
          return Star
    <|> do
          n <- natural lambdaPi
          return (toNat n)
    <|> do
          x <- identifier lambdaPi
          case findIndex (== x) e of
            Just n  -> return (Bound n)
            Nothing -> return (Free (Global x))
    <|> parens lambdaPi (parseITerm 0 e)
  
  parseCTerm :: Int -> [String] -> CharParser () CTerm
  parseCTerm 0 e =
        parseLam e
    <|> fmap Inf (parseITerm 0 e)
  parseCTerm p e =
        try (parens lambdaPi (parseLam e))
    <|> fmap Inf (parseITerm p e)
  
  parseLam :: [String] -> CharParser () CTerm
  parseLam e =
        do reservedOp lambdaPi "\\"
           xs <- many1 (identifier lambdaPi)
           reservedOp lambdaPi "->"
           t <- parseCTerm 0 (reverse xs ++ e)
           --  reserved lambdaPi "."
           return (iterate Lam t !! length xs)

  toNat :: Integer -> ITerm
  toNat n = Ann (toNat' n) (Inf Nat)

  toNat' :: Integer -> CTerm
  toNat' 0  =  Zero
  toNat' n  =  Succ (toNat' (n - 1))

  iPrint :: Int -> Int -> ITerm -> Doc
  iPrint p ii (Ann c ty)       =  parensIf (p > 1) (cPrint 2 ii c <> text " :: " <> cPrint 0 ii ty)
  iPrint p ii Star             =  text "*"
  iPrint p ii (Pi d (Inf (Pi d' r)))
                                 =  parensIf (p > 0) (nestedForall (ii + 2) [(ii + 1, d'), (ii, d)] r)
  iPrint p ii (Pi d r)         =  parensIf (p > 0) (sep [text "forall " <> text (vars !! ii) <> text " :: " <> cPrint 0 ii d <> text " .", cPrint 0 (ii + 1) r])
  iPrint p ii (Bound k)        =  text (vars !! (ii - k - 1))
  iPrint p ii (Free (Global s))=  text s
  iPrint p ii (i :$: c)         =  parensIf (p > 2) (sep [iPrint 2 ii i, nest 2 (cPrint 3 ii c)])
  iPrint p ii Nat              =  text "Nat"
  iPrint p ii (NatElim m z s n)=  iPrint p ii (Free (Global "natElim") :$: m :$: z :$: s :$: n)
  iPrint p ii (Vec a n)        =  iPrint p ii (Free (Global "Vec") :$: a :$: n)
  iPrint p ii (VecElim a m mn mc n xs)
                                 =  iPrint p ii (Free (Global "vecElim") :$: a :$: m :$: mn :$: mc :$: n :$: xs)
  iPrint p ii (Eq a x y)       =  iPrint p ii (Free (Global "Eq") :$: a :$: x :$: y)
  iPrint p ii (EqElim a m mr x y eq)
                                 =  iPrint p ii (Free (Global "eqElim") :$: a :$: m :$: mr :$: x :$: y :$: eq)
  iPrint p ii (Fin n)          =  iPrint p ii (Free (Global "Fin") :$: n)
  iPrint p ii (FinElim m mz ms n f)
                                 =  iPrint p ii (Free (Global "finElim") :$: m :$: mz :$: ms :$: n :$: f)
  iPrint p ii x                 =  text ("[" ++ show x ++ "]")

  cPrint :: Int -> Int -> CTerm -> Doc
  cPrint p ii (Inf i)    = iPrint p ii i
  cPrint p ii (Lam c)    = parensIf (p > 0) (text "\\ " <> text (vars !! ii) <> text " -> " <> cPrint 0 (ii + 1) c)
  cPrint p ii Zero       = fromNat 0 ii Zero     --  text "Zero"
  cPrint p ii (Succ n)   = fromNat 0 ii (Succ n) --  iPrint p ii (Free (Global "Succ") :$: n)
  cPrint p ii (Nil a)    = iPrint p ii (Free (Global "Nil") :$: a)
  cPrint p ii (Cons a n x xs) =
                             iPrint p ii (Free (Global "Cons") :$: a :$: n :$: x :$: xs)
  cPrint p ii (Refl a x) = iPrint p ii (Free (Global "Refl") :$: a :$: x)
  cPrint p ii (FZero n)  = iPrint p ii (Free (Global "FZero") :$: n)
  cPrint p ii (FSucc n f)= iPrint p ii (Free (Global "FSucc") :$: n :$: f)

  fromNat :: Int -> Int -> CTerm -> Doc
  fromNat n ii Zero = int n
  fromNat n ii (Succ k) = fromNat (n + 1) ii k
  fromNat n ii t = parensIf True (int n <> text " + " <> cPrint 0 ii t)

  nestedForall :: Int -> [(Int, CTerm)] -> CTerm -> Doc
  nestedForall ii ds (Inf (Pi d r)) = nestedForall (ii + 1) ((ii, d) : ds) r
  nestedForall ii ds x                = sep [text "forall " <> sep [parensIf True (text (vars !! n) <> text " :: " <> cPrint 0 n d) | (n,d) <- reverse ds] <> text " .", cPrint 0 ii x]

  parensIf :: Bool -> Doc -> Doc
  parensIf True  = PP.parens
  parensIf False = id
  
  vars :: [String]
  vars = [ c : n | n <- "" : map show [1..], c <- ['x','y','z'] ++ ['a'..'w'] ]