{-# LANGUAGE NoMonomorphismRestriction, FlexibleContexts #-}
module PTS.Core where

import Control.Monad
import Control.Monad.Trans

import Data.List
import Data.Maybe
import Data.Char
import Tools.Errors.Class

import Parametric.Error
import Parametric.Pretty

import PTS.AST
import PTS.Pretty
import PTS.Diff
import Parametric.Pretty
import PTS.Debug

-- implementation of pure type systems as described on p.98 in Barendregt's "Lambda Calculi with Types"
-- ftp://ftp.cs.kun.nl/pub/CompMath.Found/HBKJ.ps.Z
-- Restrictions: The PTS must be full and functional

-- some ideas regarding conversion to algorithmic style have been taken
-- from van Benthem et al's "Checking algorithms for Pure Type Systems"
-- http://www.springerlink.com/content/f510r63138068632/
--

import Debug.Trace

-- alpha equivalence
instance (Show c, Eq c) =>  Eq (Term c) where
  (Var x) == (Var y) = x == y
  (Nat i) == (Nat j) = i == j
  (NatOp i _ t1 t2) == (NatOp i' _ t1' t2') = i == i' && t1 == t1' && t2 == t2'
  (IfZero t1 t2 t3) == (IfZero t1' t2' t3') = t1==t1' && t2==t2' && t3==t3'
  (Const c) == (Const c') = c == c'
  (App t1 t2) == (App t1' t2') = t1 == t1' && t2 == t2'
  
  (Lam x1 q1 b1) == (Lam x2 q2 b2) = q1 == q2 && b1' == b2' where
    (_, b1', b2') = freshCommonVar x1 x2 b1 b2
  
  (Pi x1 q1 b1) == (Pi x2 q2 b2) = q1 == q2 && b1' == b2' where
    (_, b1', b2') = freshCommonVar x1 x2 b1 b2
  
  Pos _ t == Pos _ t' = t == t'
  _ == _ = False

instance (Pretty c, Show c, Eq c) => Ord (Term c) where
  compare a b = compare (show a) (show b)


-- Reduction and Evaluation of terms
isVal :: (Pretty c, Show c, Eq c) => Term c -> Bool
isVal (Const _) = True
isVal (Lam _ _ _ ) = True
isVal (Pi _ _ _ ) = True
isVal (Nat _) = True
isVal _ = False

reduce :: (Pretty c, Show c, Eq c) => Term c -> Term c
reduce (App t1 t2) | isVal t1 =
  let (Lam x _ t3) = t1 in
     subst t3 x t2
reduce (App t1 t2) = App (reduce t1) t2
reduce (NatOp _ f t1 t2) | (isVal t1) && (isVal t2) =
  let (Nat i) = t1
      (Nat j) = t2
  in
    Nat (f i j)
reduce (NatOp i f t1 t2) | (isVal t2) = NatOp i f (reduce t1) t2
reduce (NatOp i f t1 t2)  = NatOp i f t1 (reduce t2)

reduce (IfZero t1 t2 t3) | isVal t1 =
  let (Nat i) = t1 in
    if (i == 0) then t2 else t3
reduce (IfZero t1 t2 t3) = IfZero (reduce t1) t2 t3
--reduce App (Lam
reduce (Pos p t) = t
reduce t = error $ "Cannot reduce "++(show t)

eval t | isVal t = t
eval t = eval (reduce t)

-- This type defines a specific pure type system, see Barendregt
-- axioms, sorts, relations, nat
type PTS c = (c -> Maybe c, c -> Bool, c -> c -> Maybe c, c)

-- normalization and equivalence of terms
normalform :: (Show c, Eq c) => Term c -> Term c
normalform (App t1 t2) =
  case (normalform t1) of
   (Lam x _ t1) -> normalform $ subst t1 x t2
   _ -> App (normalform t1) (normalform t2)
normalform (Lam x t1 t2) = Lam x (normalform t1) (normalform t2)
normalform (Pi x t1 t2) = Pi x (normalform t1) (normalform t2)
normalform (NatOp i f t1 t2) =
  case ((normalform t1),(normalform t2)) of
    ((Nat i), (Nat j)) -> Nat (f i j)
    _ -> NatOp i f (normalform t1) (normalform t2)
normalform (IfZero t1 t2 t3) =
  case (normalform t1) of
    (Nat j) -> if (j==0) then (normalform t2) else (normalform t3)
    _ -> IfZero (normalform t1) (normalform t2) (normalform t3)
normalform (Pos p t) = normalform t
normalform t = t

equiv :: (Show c, Eq c) => Term c -> Term c -> Bool
equiv t1 t2 = (normalform t1) == (normalform t2)


check :: Monad m => Bool -> String -> m ()
check True _ = return ()
check False s = fail s

checkMaybe :: Monad m => Maybe a -> String -> m ()
checkMaybe (Just _) _ = return ()
checkMaybe Nothing s = fail s

ndots :: Int -> String
ndots n = replicate n '.'

mytrace d ctx c@(Const x) | (show c) == "*" = False
mytrace d ctx t =
  trace ((show d) ++ (ndots d) ++ (show ctx) ++ " |- "++(show t) ++ " : ???") False

-- safe bind
bind :: (Eq c, Show c) => [(String, Term c)] -> String -> Term c -> Term c -> (String, Term c, [(String, Term c)])
bind ctx x t b = case lookup x ctx of 
  Nothing -> (x, b, (x, t) : ctx)
  Just _ -> let nx = freshvarl (freevars b ++ map fst ctx) x in (nx, subst b x (Var nx), (nx, t) : ctx)

debug :: (Pretty c, Eq c, Show c, MonadIO m) => Int -> String -> [(String, Term c)] -> Term c -> m (Term c) -> m (Term c) 
debug d n ctx t result | DebugType `elem` options = do
  liftIO $ putStrLn $ 
    replicate (d+d) ' ' ++ 
    "--> " ++ n ++ ": " ++
    (let gamma = showCtx ctx in if null gamma then "" else gamma ++ " ") ++
    "|- " ++ show t ++ " : ???"
  x <- result
  liftIO $ putStrLn $ 
    replicate (d+d) ' ' ++ 
    "<-- " ++ n ++ ": " ++
    (let gamma = showCtx ctx in if null gamma then "" else gamma ++ " ") ++
    "|- " ++ show t ++ " : " ++ show x
  return x

debug d n ctx t result = result

typecheck :: (Pretty c, Eq c, Show c, MonadErrors Errors m, Functor m, MonadIO m) => PTS c -> Int -> [(String, Term c)] -> Term c -> m (Term c)
--typecheck p d ctx t | mytrace d ctx t = undefined
-- axioms



typecheck (axioms,_,_,_)  d  [] q@(Const c) = debug d "TypeConst" [] q $ do
  case (axioms c) of
    (Just t) -> return (Const t)
    _        -> fail $ "Unknown constant: " ++ show c


-- start
typecheck p@(_,sorts,_,_) d  ((x,a):ctx) q@(Var x') | (x == x') = debug d "TypeVar" ((x,a):ctx) q $ do 
  s <- (typecheck p (succ d) ctx a)
  case (normalform s) of
    (Const c) | sorts c -> return a
              | otherwise ->      fail $ "Not a valid sort: " ++ show c


-- var-not-found
typecheck p  d [] q@(Var a) = debug d "TypeVar" [] q $ 
  fail $ "Unbound identifier: " ++ a

-- product
typecheck p@(_,_,relations,_) d ctx q@(Pi x a b) = debug d "TypeFun" ctx q $
  do
    s1 <- typecheck p (succ d) ctx a
    let (newx, newb, newctx) = bind ctx x a b
    -- (newx, newb)  <- return $ case (lookup x ctx) of
    --                             Nothing -> (x :: String, b)
    --                             Just _  -> let nx = freshvarl (map fst ctx) x in (nx, subst b x (Var nx))

    s2 <- typecheck p (succ d) newctx newb
    check (case (normalform s1) of
                (Const _) -> True
		_ -> False)
	  ("Cannot normalize type "++(show s1)++" of arg "++(show a)++" to constant in "++(show q)++" in context "++(show ctx))
    check (case (normalform s2) of
                (Const _) -> True
		_ -> False)
	  ("Cannot normalize type "++(show s2)++" of body "++(show b)++" to constant in "++(show q)++" in context "++(show ctx))
    let (Const s1') =  (normalform s1)
    let Const s2' = normalform s2
    case relations s1' s2' of
      Just s -> return (Const s)
      Nothing -> fail $ "no relation " ++ show s1' ++ " : " ++ show s2'

-- application
typecheck p d ctx t@(App t1 t2) = debug d "TypeApp" ctx t $
  do
    tt1 <- typecheck p (succ d) ctx t1
    tt2 <- typecheck p (succ d) ctx t2
    r <- case normalform tt1 of
      (Pi x a b) -> if (equiv a tt2)
                       then return $ subst b x t2
                       else fail   $ let (expected, given) = showDiff 0 (diff a (normalform tt2)) in 
                                       multiLine 200 $ 
                                       text "Type error in application: Wrong argument type."
                                       $$ text "  expected:" <+> text expected
                                       $$ text "     given:" <+> text given
      _ -> fail $  "Cannot apply to term of type "++(show tt1)++ " in expression "++(show t)++" in context "++(show ctx)
    return r


-- abstraction
typecheck p@(_,_,relations,_) d ctx q@(Lam x a b) = debug d "TypeAbs" ctx q $
  do
    s1  <- typecheck p (succ d) ctx a
    let (newx, newb, newctx) = bind ctx x a b
    tb  <- typecheck p (succ d) newctx newb
    let tb' = normalform tb
    s2  <- typecheck p (succ d) newctx tb'
    let Const s1' = normalform s1
    let Const s2' = normalform s2
    maybe (fail $ "no relation " ++ show s1 ++ " : " ++ show s2) return (relations s1' s2')
    return (Pi newx a tb')


-- Nat
typecheck (_,_,_,nat) d ctx q@(Nat i)  =  debug d "TypeNat" ctx q $ return (Const nat)

-- NatOp
typecheck p@(_,_,_,nat) d ctx q@(NatOp i f t1 t2) = debug d "TypeNatOp" ctx q $
  do
    tt1 <- typecheck p (succ d) ctx t1
    check ((normalform tt1) == (Const nat)) $ "Expected: Nat, Found: "++(show (normalform tt1))++" in "++(show q)
    tt2 <- typecheck p (succ d) ctx t2
    check ((normalform tt2) == (Const nat)) $ "Expected: Nat, Found: "++(show (normalform tt2))++" in "++(show q)
    return $ Const nat

-- IfZero
typecheck p@(_,_,_,nat) d ctx q@(IfZero t1 t2 t3) = debug d "TypeIfZero" ctx q $
 do
  tt1 <- typecheck p (succ d) ctx t1
  check (normalform tt1 == (Const nat))  $ "Expected: Nat, Found: "++(show tt1)++" in "++(show q)
  tt2 <- typecheck p (succ d) ctx t2
  tt3 <- typecheck p (succ d) ctx t3
  check (normalform tt2 == normalform tt3) $ "Types in if0 do not match, found "++(show (normalform tt2))++" and "++(show (normalform tt3))++" in "++(show q)
  return (normalform tt2)

typecheck pts d ctx (Pos p t) =
  do
--    trace ("Start: "++(show ctx) ++ " |- " ++ (show t) ++ " : ???") (return ())
    x <- handlePos (typecheck pts d ctx) p t
--    trace ("End: "++(show ctx) ++ " |- " ++ (show t) ++ " : "++(show x)) (return ())
    return x

--- weakening
-- The check of the context has been disabled.
-- If we start with the empty context (which is well-formed), then this
-- check is not necessary, since everything that is added to the context
-- has already been checked earlier.
typecheck p@(_,sorts,_,_) d ((x,c):ctx) t = debug d "TypeWeakening" ((x,c):ctx) t $
--typecheck p@(_,sorts,_,_) d ((x,c):ctx) t =
--  do s <- (typecheck p ctx c)
--     case (normalform s) of
--         (Const c) -> if (sorts c) then (typecheck p ctx t) else fail $ "Not a valid sort: " ++ show c
  do
    check (lookup x ctx == Nothing) $ "Please rename variable "++x++"  - weakening error: Capture of variable "++ (show x) ++ " in context "++(show ctx)
    typecheck p (succ d) ctx t


-- some specific pure type systems


-- simply typed lambda calculus
simplytyped :: PTS String
simplytyped = ((flip lookup) [("Nat","*")],(flip elem) ["*"], rel, "Nat")
  where
    rel s1 s2 =
      case (s1,s2) of
        ("*","*") -> Just "*"
        _ -> Nothing

simplytypedConstants name = find (== name) ["Nat"]

fomega :: PTS String
fomega = ((flip lookup) [("Nat","*"), ("*", "**")], (flip elem) ["*","**"], rel, "Nat")
  where
    rel s1 s2 =
      case (s1,s2) of
        ("*","*") -> Just "*"
        ("*","**") -> Nothing
        ("**","*") -> Just "*"
        ("**","**") -> Just "**"
        _ -> Nothing

fomegaConstants name = find (== name) ["*", "**", "Nat"]

-- calculus of construction
coc :: PTS String
coc = ((flip lookup) [("Nat","*"), ("*", "**")], (flip elem) ["*","**"], rel, "Nat")
  where
    rel s1 s2 =
      case (s1,s2) of
        ("*","*") -> Just "*"
        ("*","**") -> Just "**"
        ("**","*") -> Just "*"
        ("**","**") -> Just "**"
        _ -> Nothing

cocConstants name = find (== name) ["*", "**", "Nat"]

-- lambda* with * : *
lambdastar :: PTS String
lambdastar = ((flip lookup) [("Nat","*"), ("*", "*")], (flip elem) ["*"], rel, "Nat")
  where
    rel s1 s2 =
      case (s1,s2) of
        ("*","*") -> Just "*"
        _ -> Nothing
        
lambdastarConstants name = find (== name) ["*", "Nat"]

-- \*--
fomegastar :: PTS String
fomegastar = (flip lookup axioms, flip elem constants, rel, "Nat")
  where
    constants = ["*", "**"]
    
    axioms = 
      [("Nat", "*"),  -- data types and types treated 
       ("*",   "**"), --    like in F omega
       ("**",  "**")] -- kinds treated like in lambda star
       
    rel "*"  "*"  = Just "*"  -- normal lambda
    rel "*"  "**" = Nothing   -- no dependent types!
    rel "**" "*"  = Just "*"  -- type lambda
    rel "**" "**" = Just "**" -- everything allowed for kinds 
    rel _    _    = Nothing
    
fomegastarConstants name = find (== name) ["*", "**", "Nat"]    
    
-- FOmega with an infinite hierarchy of universes
-- We call this language FOmegaOmega
data ConstFOmegaOmega = FOOStar Int | FOmegaOmegaNat
  deriving (Eq)

instance Show ConstFOmegaOmega where
  show (FOOStar n) = "*" ++ show n
  show FOmegaOmegaNat = "Nat"

instance Pretty ConstFOmegaOmega where
  pretty _ = text . show

fomegaomega :: PTS ConstFOmegaOmega
fomegaomega = (axioms,sorts,rels,FOmegaOmegaNat)
  where
   axioms FOmegaOmegaNat = Just $ FOOStar 0
   axioms (FOOStar i) = Just $ FOOStar (i+1)
   sorts (FOOStar _) = True
   sorts FOmegaOmegaNat = False
   rels (FOOStar i) (FOOStar j) | i>=j = Just $ FOOStar j
   rels _ _ = Nothing

fomegaomegaConstants name = msum
  [ do guard (name == "Nat")
       return FOmegaOmegaNat
  , do ('*' : rest) <- return name
       guard (all isDigit rest)
       return (FOOStar (read rest)) ]

-- FOmegaOmega extended with "wildcards" to enable universe polymorphism



--type PTS c = (c -> Maybe c, c -> Bool, c -> c -> Maybe c, c)

typecheck_simplytyped = typecheck simplytyped 0 []
typecheck_fomega = typecheck fomega 0 []
typecheck_coc = typecheck coc 0 []

-- some test cases
star = Const "*"
nat = Const "Nat"

test1 = App (Lam "x" nat (NatOp "add" (+) (Nat 7) (Var "x"))) (Nat 5)
-- \x:*. (\y:*.x)
test2 = (Lam "x" star (Lam "y" star (Var "x")))

test3 = typecheck coc 0 [("A", star)] $ Pi "x" (Var "A") star

-- this is the proof of the reflexivity of inclusion, see Barendregt p.94
test4 = typecheck coc 0 [("P", Pi "x" (Var "A") star),("A", star)]  (Lam "a" (Var "A") (Lam "x" (App (Var "P") (Var "a")) (Var "x")))

{-
quote :: (MonadErrors Errors m, Functor m) =>  [(String, Term String)] -> Term String -> m (Term String)
quote ctx (Nat i) = return $ App (Var "nat") (Nat i)
quote ctx (NatOp n f t1 t2) = 
  do
   t1' <- quote ctx t1
   t2' <- quote ctx t2 
   return $ App (App (Var n) t1') t2'  
quote ctx (IfZero t1 t2 t3) =
  do 
   t1' <- quote ctx t1
   ty1  <- typecheck fomegastar 0 ctx t1
   t2' <- quote ctx t2
   ty2  <- typecheck fomegastar 0 ctx t2
   t3' <- quote ctx t3
   return $ App (App (App (App (App (Var "ifzero") ty1) ty2) t1') t2') t3'
quote ctx (Var x) = return $ Var x
quote ctx (Const c) = return $ (Const c)
quote ctx (Pos p t) = 
  do
    t' <- quote ctx t
    return $ Pos p t'
quote ctx (Pi _ _ _) = fail "No Pi in quote"     
quote ctx (Lam x t1 t2) =
  do
    ty1 <- typecheck fomegastar 0 ctx t1
    tyt <- typecheck fomegastar 0 ((x,t1):ctx) t2
    ty2 <- typecheck fomegastar 0 ctx tyt
    t2' <- quote ((x,t1):ctx) t2
    if ((ty1 == (Const "*")) && (ty2 == (Const "*"))) then
       return $ App (App (App (Var "lam") t1) tyt) (Lam x (App (Var "R") t1) t2') 
      else (if ((ty1 == (Const "**")) && (ty2 == (Const "*"))) then
       return $ App (App (App (Var "tabs") t1) (Lam x t1 tyt)) (Lam x t1 t2')
      else (fail "invalid sort combination for quote"))   
quote ctx (App t1 t2) = 
  do
    ty1 <- typecheck fomegastar 0 ctx t1
    ty2 <- typecheck fomegastar 0 ctx t2
    tyty1 <- typecheck fomegastar 0 ctx ty1
    tyty2 <- typecheck fomegastar 0 ctx ty2
    ty1' <- quote ctx ty1
    ty2' <- quote ctx ty2
    if ((tyty1 == (Const "*")) && (tyty2 == (Const "*"))) then
      return $ App (App (App (App (Var "app") ty1) ty2) ty1') ty2'
     else (if (tyty2 == (Const "**")) then 
     (case ty1 of 
        (Pi x a b) -> return $ App (App (App (App (Var "tapp") ty2) (Lam x a b)) ty1') ty2'
        _ -> fail "unexpected error in quote") 
     else (fail "Sort error in quote"))
-}