{-# LANGUAGE PackageImports, TypeSynonymInstances, FlexibleInstances, GADTs #-}

-- Type checker loosely based on 
--
-- "On the Algebraic Foundation of Proof Assistants for Intuitionistic Type Theory", Abel, Coquand, Dybjer
--
-- some ideas from
--
-- "A Tutorial Implementation of a Dependently Typed Lambda Calculus", Löh, McBride, Swierstra
--
-- are also implemented.
--
-- The ideas related to parametricity and erasure are developed in
--
-- "Realizability and Parametricity in Pure Type Systems", Bernardy, Lasson
--

module TypeCheckerNF where

import Prelude hiding (length)
import Basics
import qualified Terms
import Terms (Term (Ann))
import Display
import Control.Monad.Error
import Data.Char
import Data.Maybe (isJust)
import Control.Monad.Trans.Error (ErrorT, runErrorT)
import Control.Monad.Trans.Writer
import Data.Functor.Identity
import Data.Sequence hiding (replicate)
import Data.Foldable (toList)
import Normal hiding (Term)
import Options

instance Error (Term,Doc) where
  strMsg s = (Terms.Hole dummyPosition "strMsg: panic!",text s)

type Result a = (ErrorT (Term,Doc)) -- term is used for position information
                (WriterT [Doc] Identity) a

report :: Doc -> Result ()
report x = lift $ tell [x]

runChecker :: Result a -> (Either (Term,Doc) a,[Doc])
runChecker x = runIdentity $ runWriterT $ runErrorT x

data Definition = Abstract -- ^ lambda, pi, sigma bound
                | Direct Value -- ^ pair bound

type Value    = NF
type Type     = Value
data Bind     = Bind {entryIdent :: Ident, 
                      entryValue :: Definition, -- ^ Value for identifier. 
                      entryType :: Type,    -- ^ Attention: context of the type does not contain the variable bound here.
                      entryRelevance :: Relevance
                     }
type Context  = Seq Bind

display :: Context -> NF -> Doc
display c = prettyTerm (fmap entryIdent c)

displayT :: Context -> Term -> Doc
displayT = Terms.prettyTerm . fmap entryIdent

dispContext :: Context -> Doc
dispContext ctx = case viewl ctx of
  EmptyL -> mempty
  Bind x val typ o :< ctx' -> let di = display ctx' in (case val of
    Abstract   ->             pretty x <+>                             ":" <+> di typ <+> ":" <+> pretty o
--    Direct (OfParam _ v) ->   "⟦"<>pretty x<>"⟧" <+> sep ["=" <+> parens (di v), "::" <+> di typ]
    Direct   v ->             pretty x <+> sep ["=" <+> parens (di v), ":" <+> di typ <+> ":" <+> pretty o]
    ) $$ dispContext ctx'

instance Lattice Sort where 
  s1@(Sort l1 r1) ⊔ s2@(Sort l2 r2) 
    | r1 /= r2 = s2
    | otherwise = case typeSystem options of
      CCω | l2 == 0 -> s2 -- The impredicative rule of CCω
      _ -> Sort (max l1 l2) r2

hole = Neu . Var . Hole

iType :: Context -> Term -> Result (Value,Type,Relevance)
iType g (Ann e tyt)
  =     do  (ty,o) <- iSort g tyt 
            v <- cType g e ty
            return (v,ty,sortRelevance o) -- annotations are removed
iType g t@(Terms.Star p s)
   =  return (Star s,Star $ above s, sortRelevance s)  
iType g (Terms.Pi ident tyt tyt')  
   =  do  (ty ,s1) <- iSort g tyt 
          let r1 = sortRelevance s1
          (ty',s2) <- iSort (Bind ident Abstract ty r1 <| g) tyt'
          let o = s1 ⊔ s2
          return (Pi r1 ident ty ty', Star o, sortRelevance o)
iType g (Terms.Sigma ident tyt tyt')  
   =  do  (ty,s1)  <- iSort g tyt 
          let r1 = sortRelevance s1
          (ty',s2) <- iSort (Bind ident Abstract ty r1 <| g) tyt'
          let o = s1 ⊔ s2
          return (Sigma r1 ident ty ty', Star o, sortRelevance o)
iType g e@(Terms.Bound _ x) = return $ (val $ entryValue e, wkn (x+1) $ entryType e,entryRelevance e)
  where val (Direct v) = wkn (x+1) v
        val _ = var x
        e = g `index` x
iType g (Terms.Hole p x) = do
  report $ hang (text ("context of " ++ x ++ " is")) 2 (dispContext g)
  return (hole x, hole ("type of " ++ x), 0)
iType g (e1 Terms.:$: e2)
  =     do  (v1,si,o') <- iType g e1
            case si of
              Pi o _ ty ty' -> do 
                   v2 <- cType g e2 ty
                   return (app o v1 v2, subst0 v2 ty',o') 
              _             ->  throwError (e1,"invalid application")
iType g (Terms.Proj e f) = do
  (v,t,o') <- iType g e
  search v t
 where search :: NF -> NF -> Result (Value,Type,Relevance)
       search v (Sigma o (Irr (Identifier (_,f'))) ty ty') 
              | f == f' = return (π1,ty,o)
              | otherwise = search π2 (subst0 π1 ty')
           where 
                 (π1,π2) = (case v of
                             Pair _ _ x y -> (x,y) -- substitution is useless if the pair is in normal form.
                             _ -> (proj o v True (Irr f'),proj o v False (Irr f'))  -- This could not happen if eta-expansion were done.
                             ) :: (NF,NF)
       search _ _ = throwError (e,"field not found")

iType g (Terms.Pair ident e1 e2) = do
  (v1,t1,o) <- iType g e1
  (v2,t2,o') <- iType (Bind ident (Direct v1) t1 o <| g) e2
  return $ (Pair o ident v1 (str v2),Sigma o ident t1 t2,o ⊔ o')
-- Note: the above does not infer a most general type: any potential dependency is discarded.

iType g t@(Terms.Lam x (Terms.Hole _ _) e) = throwError (t,"cannot infer type for" <+> displayT g t)
iType g (Terms.Lam x ty e) = do
    (vty,Sort _ o) <- iSort g ty
    (ve,t,o') <- iType (Bind x Abstract vty o <| g) e
    return $ (Lam o x vty ve, Pi o x vty t, o')

iType g (Terms.Param e) = do
  (v,t,o) <- iType g e
  return (param o v, app (next o) (param o t) (shift oneRel v), o)

iType g (Terms.Shift f e) = do
  (v,t,o) <- iType g e
  return (shift f v, shift f t, o + sortRelevance f)

iType g x@(Terms.Destroy d e) = do
  (v,t,o) <- iType g e  
  case d `destroys` o of
    True -> throwError (x,"total destruction is forbidden. destroyed term:" <+> display g v)
    False -> return (destroy d v,destroy d t, d-1) 

iSort :: Context -> Term -> Result (Type,Sort)
iSort g e = do
  (val,v,_) <- iType g e
  case v of 
    Star i -> return (val,i)
    (Neu (Var (Hole h))) -> do 
         report $ text h <+> "must be a type"
         return $ (hole h, Sort 1 0)
    _ -> throwError (e,displayT g e <+> "is not a type")

unify :: Context -> Term -> Type -> Type -> Result ()
unify g e q q' =
         do let ii = length g
                constraint = report $ vcat ["constraint: " <> display g q',
                                            "equal to  : " <> display g q]
            case (q,q') of
              ((Neu (Var (Hole _))), t) -> constraint
              (t, Neu (Var (Hole _))) -> constraint
              _ -> unless (q == q') 
                       (throwError (e,hang "type mismatch: " 2 $ vcat 
                                             ["inferred:" <+> display g q',
                                              "expected:" <+> display g q ,
                                              "for:" <+> displayT g e ,
                                              "context:" <+> dispContext g]))

-- Check the type and normalize
cType :: Context -> Term -> Type -> Result Value
cType g (Terms.Lam name (Terms.Hole _ _) e) (Pi o name' ty ty') = do
        e' <- cType (Bind name Abstract ty o <| g) e ty'
        return (Lam o name ty e') -- the type is filled in.

cType g (Terms.Lam name ty0 e) (Pi o name' ty ty')
  =     do (t,_o) <- iSort g ty0
           unify g (Terms.Hole (identPosition name) (show name)) t ty
           e' <- cType (Bind name Abstract ty o <| g) e ty'
           return (Lam o name ty e')

cType g (Terms.Pair name e1 e2) (Sigma o name' ty ty') = do
  -- note that names do not have to match.
  v1 <- cType g e1 ty           
  v2 <- cType (Bind name (Direct v1) ty o <| g) e2 (wk $ subst0 v1 ty') 
        -- The above weakening is there because:
        -- * the type contains no occurence of the bound variable after substitution, but
        -- * the context is extended anyway, to bind the name to its value.
  return $ Pair o name' v1 (str v2)
  -- note that the pair must use the name of the sigma for the
  -- field. (The context will use the field name provided by the type)

--  Γ ⊢ ⌊A⌋ : B
cType g (Terms.OfParam i e) t = do
  -- Γ ⊢ A ⌊A⌋ : ⟦B⟧ ⌊A⌋
  -- Γ ⊢ A x   : ⟦B⟧ x
  -- Γ ⊢ A     : (x : ⌊B⌋) → ⟦B⟧ x
  -- FIXME: here I just assume the relevance of t is the following.
  let theRelevance = 0 
      theType = Pi (next theRelevance) i (shift oneRel $ t) 
                   (zerInRel 0 t theRelevance)
  e' <- cType g e theType       
  return (Neu $ OfParam i e')

cType g (Terms.Shift f e) t = do
  shift f <$> cType g e (shift (negate f) t) 
  -- there might be negative sorts in there, but that should be fine;
  -- if they occur the type checker will simply reject the term
  -- because it's impossible to create an inhabitant of a negative
  -- sort.

cType g e v 
  =     do (e',v',_o) <- iType g e
           unify g e v v'
           return e'