{-# LANGUAGE TypeInType #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
module Language.Symantic.Typing.Quantification where

import Data.Map.Strict (Map)
import Data.Set (Set)
import qualified Control.Monad.Trans.State.Strict as SS
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

import Language.Symantic.Parsing
import Language.Symantic.Typing.Constant
import Language.Symantic.Typing.Kind
import Language.Symantic.Typing.Type

-- | /type beta-reduction/: substitute the abstraction
-- of a 'TyQuant' by using the given argument, when possible.
unTyQuant :: Type src cs (a::ka) -> Type src cs b -> Maybe (KT src cs ka)
unTyQuant (TyQuant _src _ kv f) v
 | Just Refl <- eqKi kv (kindTyOS v)
 = Just (f v)
unTyQuant _f _v = Nothing
infixl 2 `unTyQuant`

-- | Monomorphize given 'Type':
-- introduce new 'TyVar's to remove rank-1 'TyQuant's.
monoTy :: Monad m
 => Type src cs (h::k)
 -> SS.StateT Var m (KT src cs k)
monoTy t@TyConst{}          = return $ KT t
monoTy t@(TyApp _src _f _a) = return $ KT t
monoTy (TyQuant src _n kv f) = do
	v <- SS.get
	SS.put $ v + 1
	case f $ TyVar src kv v of
	 KT t -> monoTy t
monoTy t@TyVar{} = return $ KT t
monoTy (Term te x) =
	monoTy x >>= \case
	 KT t | Just Refl <- eqTy x t -> return $ KT $ Term te t
	 t -> return t

-- * Type 'MapTyVar'
-- | /type variable substitution/
type MapTyVar src cs = Map Var (EType src cs)

-- | Substitute the 'TyQuant's of the given 'Type',
-- according to the given 'MapTyVar',
-- where their associated 'Var' counts from the given 'Var'.
substTyQuant :: forall src cs k h.
 Var -> MapTyVar src cs
 -> Type src cs (h::k)
 -> KT src cs k
substTyQuant var vs = go var
	where
	go :: forall kx x. Var -> Type src cs (x::kx) -> KT src cs kx
	go _v t@TyConst{} = KT t
	go v (TyApp src f a) =
		bind2KT (\f' a' -> KT $ TyApp src f' a')
		 (go v f)
		 (go v a)
	go v (TyQuant src n kv f) =
		case Map.lookup v vs of
		 Just (EType TyVar{}) -> KT $ TyQuant src n kv $ bindKT (go (v + 1)) . f
		 Just (EType r) | Just Refl <- eqKi kv (kindTyOS r) -> f r
		 Just{} -> error "BUG: substTyQuant: expected and actual kinds mismatch"
		 Nothing -> KT $ TyQuant src n kv $ bindKT (go (v + 1)) . f
	go _v t@TyVar{}   = KT t
	go v (Term te x) =
		case go v x of
		 KT t | Just Refl <- eqTy x t -> KT $ Term te t -- FIXME: likely useless
		 t -> t

-- | Substitute the 'TyVar's of the given 'Type',
-- according to the given 'MapTyVar'.
substTyVar :: forall src cs k h. MapTyVar src cs -> Type src cs (h::k) -> KT src cs k
substTyVar vs = go
	where
	go :: forall kx x. Type src cs (x::kx) -> KT src cs kx
	go (TyQuant src n kv f) = KT $ TyQuant src n kv $ bindKT go . f
	go (TyApp src f a)      = bind2KT (\f' a' -> KT $ TyApp src f' a') (go f) (go a)
	go t@TyConst{} = KT t
	go t@(TyVar _src kv v) =
		case Map.lookup v vs of
		 Nothing -> KT t
		 Just (EType x) | Just Refl <- eqKi kv (kindTyOS x) -> KT x
		 Just{} -> error "BUG: substTyVar: expected and actual kinds mismatch"
	go (Term te x) =
		case go x of
		 KT t | Just Refl <- eqTy x t -> KT $ Term te t -- FIXME: likely useless
		 t -> t

-- | Change the 'TyVar's of the given 'Type' matching the first 'Var', to the second.
--
-- NOTE: This is special case of 'substTyVar'
-- in which the indexed type can be preserved.
changeTyVar :: Var -> Var -> Type src cs (h::k) -> Type src cs (h::k)
changeTyVar ov nv = go
	where
	go :: Type src cs (h::k) -> Type src cs (h::k)
	go t@TyConst{} = t
	go (TyApp src f a) = TyApp src (go f) (go a)
	go (TyQuant src n kv f) = TyQuant src n kv $ mapKT go . f
	go t@(TyVar src kv v) | v == ov = TyVar src kv nv -- matching Var
	                      | otherwise = t
	go (Term te x) = Term te $ go x

-- | Unify two 'MapTyVar'.
--
-- NOTE: the union is left-biased: in case of duplicate 'Var',
-- it keeps those from the first 'MapTyVar' given.
--
-- NOTE: the first 'MapTyVar' given is applied to the second (with 'substTyVar'),
-- this way each 'Var' directly maps to an expanded 'Type',
-- so that, when using the resulting 'MapTyVar',
-- there is no need to apply it multiple times
-- until there is no more substitution to be done.
unionSubst :: MapTyVar src cs -> MapTyVar src cs -> MapTyVar src cs
unionSubst x y = x `Map.union` ((\(EType t) -> EType `ofKT` substTyVar x t) <$> y)

-- | Return the 'Var' of the 'TyVar's in given 'Type'.
freeTyVars :: Type src cs (h::k) -> Set Var
freeTyVars = go
	where
	go :: Type src cs h -> Set Var
	go TyConst{} = Set.empty
	go (TyApp _src f a) = go f `Set.union` go a
	go (TyQuant  src _n kv f) = go `ofKT` f (TyVar src kv (-1))
	go (TyVar _src _kv v) | v < 0     {- bound Var -} = Set.empty
	                      | otherwise {- free  Var -} = Set.singleton v
	go (Term _te x) = go x

-- | Like 'freeTyVars', but also return the 'Kind' of each 'Var'.
freeKiTyVars :: Type src cs (h::k) -> Map Var (EKind src)
freeKiTyVars = go
	where
	go :: Type src cs h -> Map Var (EKind src)
	go TyConst{} = Map.empty
	go (TyApp _src f a) = go f `uni` go a
	go (TyQuant  src _n kv f) = go `ofKT` f (TyVar src kv (-1))
	go (TyVar _src kv v) | v < 0     {- bound Var -} = Map.empty
	                     | otherwise {- free  Var -} = Map.singleton v (EKind kv)
	go (Term _te x) = go x
	
	uni :: Map Var (EKind src) -> Map Var (EKind src) -> Map Var (EKind src)
	uni = Map.unionWith $ \new@(EKind kx) (EKind ky) ->
		case eqKi kx ky of
		 Just{} -> new
		 Nothing -> error "BUG: freeKiTyVars: inconsistent kinds"

{-
-- | Remove 'TyQuant's binding a type variable which is not used.
cleanTyQuants :: Type src cs (h::k) -> Type src cs (h::k)
cleanTyQuants t = go (freeTyVar t) t
	where
	go :: Var -> Type src cs (h::k) -> Type src cs (h::k)
	go v (TyQuant src n kv f) =
		let x = f $ TyVar src kv v in
		if v `elem` freeTyVars x
		then TyQuant src n kv $ go (v + 1) . f
		else x
	go _v x@TyConst{}     = x
	go _v x@TyVar{}       = x
	go v (TyApp src f a)  = TyApp src (go v f) (go v a)
	go v (Term te x)    = Term te $ go v x
-}

-- ** Type 'Error_MGU'
-- | Reasons why two 'Type's cannot be unified.
data Error_MGU src cs
 =   Error_MGU_TyVar_Loop Var (EType src cs)
     -- ^ /occurence check/: a 'TyVar' is unified with a 'Type'
     --   which contains this same 'TyVar'.
 |   Error_MGU_TyConst_mismatch (EType src cs) (EType src cs)
     -- ^ Two 'TyConst's should be the same, but are different.
 |   Error_MGU_Kind_mismatch (EKind src) (EKind src)
     -- ^ Two 'Kind's should be the same, but are different.
 |   Error_MGU_Kind (Con_Kind src)
     -- ^ Two 'Kind's mismatch.
 |   Error_MGU_cannot_unify (EType src cs) (EType src cs)
     -- ^ Cannot unify those two 'Type's.
deriving instance
 ( Show src
 , Show_TyConst cs
 , Proj_TyConst cs (->)
 , Proj_TyConst cs (#>)
 ) => Show (Error_MGU src cs)

instance Inj_Error (Error_MGU src cs) (Error_MGU src cs) where
	inj_Error = id
instance Inj_Error (Con_Kind src) (Error_MGU src cs) where
	inj_Error = Error_MGU_Kind

-- | Return the /most general unification/ of two 'Type's, when it exists.
-- 
-- Support:
--
--  * 'TyConst' : OK
--  * 'TyApp'   : OK
--  * 'TyQuant' : Error_MGU_cannot_unify
--  * 'TyVar'   : OK
--  * 'Term'    : Transparent
mguTy
 :: forall src cs ki a b.
    Inj_Source (EType src cs) src
 => MapTyVar src cs
 -> Type src cs (a::ki)
 -> Type src cs (b::ki)
 -> Either (Error_MGU src cs) (MapTyVar src cs)
mguTy = go
	where
	go
	 :: MapTyVar src cs
	 -> Type src cs (x::k)
	 -> Type src cs (y::k)
	 -> Either (Error_MGU src cs) (MapTyVar src cs)
	go vs x y =
		let ki = kindTyOS x in
		let (xh, xp) = spineTy x in
		let (yh, yp) = spineTy y in
		case ( unTerm `mapEType` xh
		     , unTerm `mapEType` yh ) of
		 ( EType (TyVar _ xkv xv), _ ) | Just Refl <- eqKi xkv ki -> goVar vs xkv xv y
		 ( _, EType (TyVar _ ykv yv) ) | Just Refl <- eqKi ki ykv -> goVar vs ykv yv x
		 ( EType (TyConst _xs xc), EType (TyConst _ys yc) )
		  | Just KRefl <- eqKiTyConst xc yc -> goList vs xp yp
		  | otherwise -> Left $ inj_Error $ Error_MGU_TyConst_mismatch xh yh
		 _ ->
			case (unTerm x, unTerm y) of
			 ( TyApp _ xa xb, TyApp _ ya yb ) ->
				goList vs
				 [EType xa, EType xb]
				 [EType ya, EType yb]
			 _ -> Left $ inj_Error $ Error_MGU_cannot_unify (EType x) (EType y)
	goVar
	 :: MapTyVar src cs
	 -> Kind src k
	 -> Var
	 -> Type src cs (y::k)
	 -> Either (Error_MGU src cs) (MapTyVar src cs)
	goVar vs xkv xv y =
		case Map.lookup xv vs of
		 Just (EType x) -> if_EqKind xkv (kindTyOS x) $ \Refl -> go vs y x
		 Nothing ->
			case substTyVar vs y of
			 KT (TyVar _src _ykv yv) | xv == yv -> Right vs
			 KT y'
			  | xv `elem` freeTyVars y' -> Left $ Error_MGU_TyVar_Loop xv (EType y')
			  | otherwise -> Right $ Map.singleton xv (EType y') `unionSubst` vs
	goList
	 :: MapTyVar src cs
	 -> [EType src cs]
	 -> [EType src cs]
	 -> Either (Error_MGU src cs) (MapTyVar src cs)
	goList vs [] [] = Right vs
	goList vs (EType x:xs) (EType y:ys) =
		if_EqKind (kindTy x) (kindTy y) $ \Refl ->
		go vs x y >>= \vs' -> goList vs' xs ys
	goList _vs _x _y = error "BUG: mguTy: kinds mismatch"

-- | Return the left spine of a 'Type':
-- the root 'Type' and its 'Type' parameters,
-- from the left to the right.
spineTy
 :: forall src cs k h. Inj_Source (EType src cs) src
 => Type src cs (h::k)
 -> (EType src cs, [EType src cs])
spineTy typ = go [] typ
	where
	go :: forall x. Inj_Source (EType src cs) src
	   => [EType src cs]
	   -> Type src cs x
	   -> (EType src cs, [EType src cs])
	go ts (TyApp  _src f a) = go (EType (a `source` EType typ) : ts) f
	go ts (Term _te x) = go ts x
	go ts t = (EType (t `source` EType typ), ts)
{-
spineTy
 :: Type src cs (h::k)
 -> (forall kx (x::kx) xs. Type src cs x -> Types src cs xs -> ret)
 -> ret
spineTy = go TypesZ
	where
	go :: Types src cs hs
	   -> Type src cs (h::k)
	   -> (forall kx (x::kx) xs. Type src cs x -> Types src cs xs -> ret)
	   -> ret
	go ts (TyApp  _src f a) k = go (a `TypesS` ts) f k
	go ts (Term _te x)    k = go ts x k
	go ts (TyAny x)         k = go ts x k
	go ts t k = k t ts
-}

-- | Return, when the given 'Type' is a function,
-- the argument of that function.
unTyFun :: forall src cs h.
 Inj_Source (EType src cs) src
 => Proj_TyConst cs (->)
 => Type src cs h -> Maybe (KTerm src cs, KTerm src cs)
unTyFun t = go t
	where
	go :: forall x. Type src cs x -> Maybe (KTerm src cs, KTerm src cs)
	go (TyApp _ (TyApp _ (TyConst _ f) a) b)
	 | Just KRefl <- projKiTyConst @_ @(->) f
	 = Just (KT (a `source` EType t), KT (b `source` EType t))
	go TyApp{}   = Nothing
	go TyConst{} = Nothing
	go TyQuant{} = Nothing
	go TyVar{}   = Nothing
	go (Term _te x) = go x

-- | Return given 'Type', polymorphized by 'TyQuant'ifying
-- over the given 'Var' of given 'Kind' and 'NameHint'.
polyTyVar :: forall src cs kv h.
 Source src
 => Show_TyConst cs
 => Proj_TyConst cs (->)
 => Proj_TyConst cs (#>)
 => NameHint -> Kind src kv -> Var
 -> Type src cs (h::k)
 -> KT src cs k
polyTyVar nv kv var typ =
	KT $ TyQuant (source_of typ) nv kv $
	go var typ
	where
	go :: forall km (m::km).  Var -> Type src cs m -> (forall (v::kv). Type src cs v -> KT src cs km)
	go _v m@TyConst{} = \_a -> KT m
	go v (TyApp src x y) = \a -> bind2KT (\f' a' -> KT $ TyApp src f' a') (go v x a) (go v y a)
	go v (TyQuant src n kv' f) = \a ->
		KT $ TyQuant src n kv' $ \a' ->
			(\p -> go v p a) `ofKT` f a'
	go v (TyVar _src kv' v') | v == v' =
		case eqKi kv kv' of
		 Nothing -> error "BUG: polyTyVar: polymorphized TyVar has not expected kind"
		 Just Refl -> KT
	go _v m@TyVar{} = \_a -> KT m
	go v (Term te m) = \a ->
		case go v m a of
		 KT p
		  | TyVar _ _ka va <- a
		  , m' <- changeTyVar v va m
		  , Just Refl <- eqTy m' p ->
			KT $ Term te p
			-- NOTE: if @a@ is also a 'TyVar',
			-- the indexed type can be preserved
			-- and thus the 'Term' can be preserved too.
			-- However, the old @v@ used to polymorphize
			-- must be changed to the new @va@
			-- to make 'eqTy' succeed.
			-- Hence, the 'Term' can survive cycles
			-- of 'monoTy' then 'polyTy'.
			-- Because @freeTyVars p == freeTyVars m \\ v@,
			-- @va@ should never be within @freeTyVars m@,
			-- unless it again is @v@.
		 KT p -> KT p

-- | Polymorphize given 'Type':
-- applying successively 'polyTyVar'
-- for all 'freeTyVars' of given 'Type'.
polyTy
 :: Source src
 => Show_TyConst cs
 => Proj_TyConst cs (->)
 => Proj_TyConst cs (#>)
 => Type src cs (h::k)
 -> KT src cs k
polyTy t =
	let fvs = freeKiTyVars t in
	Map.foldrWithKey
	 (\v (EKind kv) (KT x) ->
		polyTyVar "" kv v x)
	 (KT t) fvs