{-# LANGUAGE ConstraintKinds, TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RebindableSyntax #-}
module FORecursiveTypes.EliminateTypeArgumentOfCheck where

import Prelude hiding (Monad(..), (>=), (<=), lookup)
import GHC.Exts (Constraint)
import Data.List(find)
import Data.Map(Map)
import qualified Data.Map as Map

import FORecursiveTypes.Language
import Util.ErrorMessages

-- is needed because we use the RebindableSyntax extension
ifThenElse :: Bool -> a -> a -> a
ifThenElse True thn _ = thn
ifThenElse False _ els = els

type Error = [String]

data Infer a
  = Inferred a
  | NotInferred Error deriving (Eq, Show)

type Check = Infer ()

instance Functor Infer where
  fmap _ (NotInferred err) = NotInferred err
  fmap f (Inferred ty) = Inferred $ f ty

instance Applicative Infer where
  pure = Inferred
  (NotInferred err) <*> _ = NotInferred err
  (Inferred a) <*> something = fmap a something

class WithTop a where
  top :: a

instance WithTop Type where
  top = AnyType

instance WithTop () where
  top = ()

instance (WithTop a, WithTop b) => WithTop (a, b) where
  top = (top, top)

instance WithTop a => WithTop [a] where
  top = [top]

instance WithTop a => WithTop (Map Name a) where
  top = Map.empty

-- Had to define an own monad type class.
-- It is not possible otherwise to get the type constraint WithTop a.
-- We use the extension ConstraintKinds to support this.
-- Could not find a simpler solution for this problem.
-- The restricted monad problem is common.
class RMonad m where
  type RMonadCtx m a :: Constraint
  return :: RMonadCtx m a => a -> m a
  (>>=) :: (RMonadCtx m a, RMonadCtx m b) => m a -> (a -> m b) -> m b
  (>>) :: (RMonadCtx m a, RMonadCtx m b) => m a ->  m b -> m b
  m >> k = m >>= \_ -> k
  fail :: [String] -> m a

instance RMonad Infer where
  type RMonadCtx Infer a = WithTop a
  return = Inferred
  (Inferred ty) >>= f = f ty
  NotInferred err1 >>= f =
    -- we know that top is Inferred AnyType, (AnyType, AnyType) or () by definition
    case f top of
      Inferred _ -> fail err1
      NotInferred err2 -> fail $ err1 ++ err2
  fail = NotInferred

-- matching functions that extract the inner types if possible
-- one problem is that we do not get as good error messages, because term is not known in these functions
matchNat :: Type -> String -> Check
matchNat Nat _ = return ()
matchNat ty err = fail [natError ty err]

matchFun :: Type -> String -> Infer (Type, Type)
matchFun (Fun ty1 ty2) _ = return (ty1, ty2)
matchFun ty err = fail [funError ty err]

matchType :: Type -> Type -> String -> Check
matchType ty1 ty2 _
  | ty1 == ty2 = return ()
matchType ty1 ty2 err = fail [generalError (show ty1) ty2 err]

matchSum :: Type -> String -> Infer (Type, Type)
matchSum (Sum ty1 ty2) _ = return (ty1, ty2)
matchSum ty err = fail [sumError ty err]

matchVariant :: Type -> String -> Infer (Map.Map Name Type)
matchVariant (Variant types) _ = return types
matchVariant ty err = fail [variantError ty err]

lookupTypeVar :: TypeMap -> Name -> Infer Type
lookupTypeVar Empty x = fail ["Unbound type variable " ++ show x]
lookupTypeVar (Bind c x t) y
  | x == y = return t
  | otherwise = lookupTypeVar c y

matchTypeVar :: TypeMap -> Type -> Infer Type
matchTypeVar tymap (TypeVar x) = lookupTypeVar tymap x
matchTypeVar _ ty = return ty

liftMaybe :: WithTop a => Maybe a -> String -> Infer a
liftMaybe (Just a) _ = return a
liftMaybe Nothing err = fail [err]

lookup :: Ctx -> Name -> Infer Type
lookup Empty x = fail ["Unbound variable " ++ show x]
lookup (Bind c x t) y
  | x == y = return t
  | otherwise = lookup c y

inferType :: Ctx -> TypeMap -> Term -> Infer Type
inferType _ _ (Unit _) = return UnitT
inferType _ _ (Zero _) = return Nat
inferType ctx tymap (Succ t _) = do
  checkType ctx tymap t
  return Nat
inferType ctx tymap (Add t1 t2 _) = do
  checkType ctx tymap t1
  checkType ctx tymap t2
  return Nat
inferType ctx tymap (Mult t1 t2 _) = do
  checkType ctx tymap t1
  checkType ctx tymap t2
  return Nat
inferType ctx _ (Var name _) = lookup ctx name
inferType ctx tymap (Let name t body _) = do
  tyt <- inferType ctx tymap t
  inferType (Bind ctx name tyt) tymap body
inferType ctx tymap (Anno term ty _) = do
  checkType ctx tymap term
  return ty
inferType ctx tymap (App t1 t2 _) = do
  ty <- inferType ctx tymap t1
  rty <- matchTypeVar tymap ty
  (ty1, ty2) <- matchFun rty (show t1)
  checkType ctx tymap t2
  return ty2
inferType ctx tymap (LetType n ty t _) = inferType ctx (Bind tymap n ty) t
inferType _ _ t = fail ["Cannot infer type of term " ++ show t]

checkType :: Ctx -> TypeMap -> Term -> Check
checkType ctx tymap p@(Lam name t _) = do
  ty <- requiredType ctx tymap p
  rty <- matchTypeVar tymap ty
  (ty1, ty2) <- matchFun rty (show p)
  checkType (Bind ctx name ty1) tymap t
checkType ctx tymap p@(InL t _) = do
  ty <- requiredType ctx tymap p
  rty <- matchTypeVar tymap ty
  (ty1, ty2) <- matchSum rty (show p)
  checkType ctx tymap t
checkType ctx tymap p@(InR t _) = do
  ty <- requiredType ctx tymap p
  rty <- matchTypeVar tymap ty
  (ty1, ty2) <- matchSum rty (show p)
  checkType ctx tymap t
checkType ctx tymap p@(Case e n1 t1 n2 t2 _) = do
  ety <- inferType ctx tymap e
  rty <- matchTypeVar tymap ety
  (ty1, ty2) <- matchSum rty (show e)
  ty <- requiredType ctx tymap p
  checkType (Bind ctx n1 ty1) tymap t1
  checkType (Bind ctx n2 ty2) tymap t2
checkType ctx tymap p@(Tag n t _) = do
  ty <- requiredType ctx tymap p
  rty <- matchTypeVar tymap ty
  types <- matchVariant rty (show p)
  let lty = Map.lookup n types
  lty <- liftMaybe (Map.lookup n types) "Label not contained in Variant"
  checkType ctx tymap t
checkType ctx tymap p@(Match m cases _) = do
  ety <- inferType ctx tymap m
  rty <- matchTypeVar tymap ety
  types <- matchVariant rty (show m)
  ty <- requiredType ctx tymap p
  let subchecks =
        map (\(l, x, t) -> do
          lty <- liftMaybe (Map.lookup l types) "Could not find labeled type"
          checkType (Bind ctx x lty) tymap t
        ) cases
  foldl (>>) (return ()) subchecks

checkType ctx tymap t = do
  ty <- requiredType ctx tymap t
  ty' <- inferType ctx tymap t
  matchType ty ty' (show t)


requiredType :: Ctx -> TypeMap -> Term -> Infer Type
requiredType ctx tymap t = case parent t of
  Just (Succ t' _) | t == t' -> return Nat
  Just (Add t1 t2 _) | t == t1 ->
    return Nat
  Just (Add t1 t2 _) | t == t2 ->
    return Nat
  Just (Mult t1 t2 _) | t == t1 ->
    return Nat
  Just (Mult t1 t2 _) | t == t2 ->
    return Nat
  Just (Anno term ty _) | t == term -> return ty
  Just p@(App t1 t2 _) | t == t2 -> do
    ty <- inferType ctx tymap t1
    rty <- matchTypeVar tymap ty
    (ty1, ty2) <- matchFun rty (show t1)
    return ty1
  Just p@(Lam name term _) | t == term -> do
    ty <- requiredType ctx tymap p
    rty <- matchTypeVar tymap ty
    (ty1, ty2) <- matchFun rty (show p)
    return ty2
  Just p@(InL term _) | t == term -> do
    ty <- requiredType ctx tymap p
    rty <- matchTypeVar tymap ty
    (ty1, ty2) <- matchSum rty (show p)
    return ty1
  Just p@(InR term _) | t == term -> do
    ty <- requiredType ctx tymap p
    rty <- matchTypeVar tymap ty
    (ty1, ty2) <- matchSum rty (show p)
    return ty2
  Just p@(Case e n1 t1 n2 t2 _) | t == t1 -> do
    ety <- inferType ctx tymap e
    rty <- matchTypeVar tymap ety
    (ty1, ty2) <- matchSum rty (show e)
    ty <- requiredType ctx tymap p
    return ty
  Just p@(Case e n1 t1 n2 t2 _) | t == t2 -> do
    ety <- inferType ctx tymap e
    rty <- matchTypeVar tymap ety
    (ty1, ty2) <- matchSum rty (show e)
    ty <- requiredType ctx tymap p
    return ty
  Just p@(Tag n term _) | t == term -> do
    ty <- requiredType ctx tymap p
    rty <- matchTypeVar tymap ty
    types <- matchVariant rty (show p)
    let lty = Map.lookup n types
    liftMaybe (Map.lookup n types) "Label not contained in Variant"
  Just p@(Match m cases _) -> do
    ty <- requiredType ctx tymap p
    return ty
  _ -> fail ["Could not determine required type"]