{-# 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.Helper.Data.Type.Equality
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 ctx ss cs (a::ka) -> Type src ctx ss cs b -> Maybe (KT src ctx ss 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 ctx ss cs (h::k)
 -> SS.StateT Var m (KT src ctx ss 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 t@(TyFam _src _f _hs) = return $ KT t
monoTy (Term x te) =
	monoTy x >>= \case
	 KT t | Just Refl <- eqTy x t -> return $ KT $ Term t te
	 t -> return t

-- * Type 'MapTyVar'
-- | /type variable substitution/
type MapTyVar src ctx ss cs = Map Var (EType src ctx ss 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 ctx ss cs k h.
 Var -> MapTyVar src ctx ss cs
 -> Type src ctx ss cs (h::k)
 -> KT   src ctx ss cs k
substTyQuant var vs = go var
	where
	go :: forall kx x. Var -> Type src ctx ss cs (x::kx) -> KT src ctx ss 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) -> bindKT (go (v + 1)) $ 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 (TyFam src f hs) = goArgs v hs $ \hs' -> KT $ TyFam src f hs'
		where
		goArgs :: forall hs ret. Var
		 -> Types src ctx ss cs hs
		 -> (forall hs'. Types src ctx ss cs hs' -> ret)
		 -> ret
		goArgs _v TypesZ k = k TypesZ
		goArgs v' (TypesS h as) k =
			case go v' h of
			 KT h' -> goArgs v' as $ \as' -> k $ TypesS h' as'
	go v (Term x te) =
		case go v x of
		 KT t | Just Refl <- eqTy x t -> KT $ Term t te -- FIXME: likely useless
		 t -> t

-- | Substitute the 'TyVar's of the given 'Type',
-- according to the given 'MapTyVar'.
substTyVar
 :: forall src ctx ss cs k (h::k).
    MapTyVar src ctx ss cs
 -> Type src ctx ss cs h
 -> KT   src ctx ss cs k
substTyVar vs = go
	where
	go :: forall kx x. Type src ctx ss cs (x::kx) -> KT src ctx ss 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 (TyFam src f hs) = goArgs hs $ \hs' -> KT $ TyFam src f hs'
		where
		goArgs :: forall hs ret.
		    Types src ctx ss cs hs
		 -> (forall hs'. Types src ctx ss cs hs' -> ret)
		 -> ret
		goArgs TypesZ k = k TypesZ
		goArgs (TypesS h as) k =
			case go h of
			 KT h' -> goArgs as $ \as' -> k $ TypesS h' as'
	go (Term x te) =
		case go x of
		 KT t | Just Refl <- eqTy x t -> KT $ Term t te -- 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 ctx ss cs (h::k) -> Type src ctx ss cs (h::k)
changeTyVar ov nv = go
	where
	go :: Type src ctx ss cs (h::k) -> Type src ctx ss 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 (TyFam src f hs) = TyFam src f $ go `mapTys` hs
	go (Term x te) = Term (go x) te

-- | 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 ctx ss cs -> MapTyVar src ctx ss cs -> MapTyVar src ctx ss cs
unionSubst x y = x `Map.union` ((\(EType t) -> EType `ofKT` substTyVar x t) <$> y)

-- | Return the smallest 'Var' not used in given 'Type'.
freeTyVar :: Type src ctx ss cs (h::k) -> Var
freeTyVar = (+ 1) . maxTyVar (-1)
	where
	maxTyVar :: Var -> Type src ctx ss cs h -> Var
	maxTyVar m TyConst{}             = m
	maxTyVar m (TyApp   _src f a)    = maxTyVar m f `max` maxTyVar m a
	maxTyVar m (TyQuant src _n kv f) = maxTyVar m `ofKT` f (TyVar src kv (-1))
	maxTyVar m (TyVar   _src _kv v)  = max m $ if v < 0 then -1 else v + 1
	maxTyVar m (TyFam _src _f as)    = foldrTys (flip maxTyVar) as m
	maxTyVar m (Term x _te)          = maxTyVar m x

-- | Return the 'Var' of the 'TyVar's in given 'Type'.
freeTyVars :: Type src ctx ss cs (h::k) -> Set Var
freeTyVars = go
	where
	go :: Type src ctx ss 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 (TyFam _src _f as) = foldrTys (\h acc -> Set.union acc $ go h) as Set.empty
	go (Term x _te) = go x

-- | Like 'freeTyVars', but also return the 'Kind' of each 'Var'.
freeKiTyVars :: Type src ctx ss cs (h::k) -> Map Var (EKind src)
freeKiTyVars = go
	where
	go :: Type src ctx ss 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 (TyFam _src _f as) = foldrTys (\h acc -> uni acc $ go h) as Map.empty
	go (Term x _te) = 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 ctx ss cs (h::k) -> Type src ctx ss cs (h::k)
cleanTyQuants t = go (freeTyVar t) t
	where
	go :: Var -> Type src ctx ss cs (h::k) -> Type src ctx ss 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 x te)    = Term (go v x) te
-}

-- ** Type 'Error_MGU'
-- | Reasons why two 'Type's cannot be unified.
data Error_MGU src ss cs
 =   Error_MGU_TyVar_Loop Var (EType src '[] ss cs)
     -- ^ /occurence check/: a 'TyVar' is unified with a 'Type'
     --   which contains this same 'TyVar'.
 |   Error_MGU_TyConst_mismatch (EType src '[] ss cs) (EType src '[] ss 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 '[] ss cs) (EType src '[] ss cs)
     -- ^ Cannot unify those two 'Type's.
deriving instance
 ( Eq src
 , Eq_Token ss
 , Inj_Source (EType src '[] ss cs) src
 ) => Eq (Error_MGU src ss cs)
deriving instance
 ( Show src
 , Show_TyConst cs
 , Show (EType src '[] ss cs)
 , Fixity_TyConst cs
 ) => Show (Error_MGU src ss cs)

instance Inj_Error (Error_MGU src ss cs) (Error_MGU src ss cs) where
	inj_Error = id
instance Inj_Error (Con_Kind src) (Error_MGU src ss 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 ctx ss cs ki a b.
    Inj_Source (EType src '[] ss cs) src
 => MapTyVar src ctx ss cs
 -> Type src ctx ss cs (a::ki)
 -> Type src ctx ss cs (b::ki)
 -> Either (Error_MGU src ss cs) (MapTyVar src ctx ss cs)
mguTy = go
	where
	go
	 :: MapTyVar src ctx ss cs
	 -> Type src ctx ss cs (x::k)
	 -> Type src ctx ss cs (y::k)
	 -> Either (Error_MGU src ss cs) (MapTyVar src ctx ss cs)
	go vs x y =
		let ki = kindTyOS x in
		case (spineTy x, spineTy y) of
		 ((EType (unTerms -> hx), px), (EType (unTerms -> hy), py)) ->
			case (hx, hy) of
			 (TyVar _ kvx vx, _) | Just Refl <- eqKi kvx ki -> goVar vs kvx vx y
			 (_, TyVar _ kvy yv) | Just Refl <- eqKi ki kvy -> goVar vs kvy yv x
			 (TyConst _sx cx, TyConst _sy cy)
			  | Just KRefl <- eqKiTyConst cx cy -> goList vs px py
			  | otherwise -> Left $ inj_Error $
				Error_MGU_TyConst_mismatch (EType hx) (EType hy)
			 _ ->
				case (unTerm x, unTerm y) of
				 (TyApp _ ax bx, TyApp _ ay by) ->
					goList vs
					 [EType ax, EType bx]
					 [EType ay, EType by]
				 _ -> Left $ inj_Error $
					Error_MGU_cannot_unify
					 (EType $ unCtxTe x)
					 (EType $ unCtxTe y)
	goVar
	 :: MapTyVar src ctx ss cs
	 -> Kind src k
	 -> Var
	 -> Type src ctx ss cs (y::k)
	 -> Either (Error_MGU src ss cs) (MapTyVar src ctx ss 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 $ unCtxTe y')
			       | otherwise -> Right $ Map.singleton xv (EType y') `unionSubst` vs
	goList
	 :: MapTyVar src ctx ss cs
	 -> [EType src ctx ss cs]
	 -> [EType src ctx ss cs]
	 -> Either (Error_MGU src ss cs) (MapTyVar src ctx ss 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 ctx ss cs k h.
    Inj_Source (EType src '[] ss cs) src
 => Type src ctx ss cs (h::k)
 -> (EType src ctx ss cs, [EType src ctx ss cs])
spineTy typ = go [] typ
	where
	go :: forall x. Inj_Source (EType src '[] ss cs) src
	   => [EType src ctx ss cs]
	   -> Type src ctx ss cs x
	   -> (EType src ctx ss cs, [EType src ctx ss cs])
	go ts (TyApp  _src f a) = go (EType (a `source` EType (unCtxTe typ)) : ts) f
	go ts (Term x _te) = go ts x
	go ts t = (EType (t `source` EType (unCtxTe typ)), ts)
{-
spineTy
 :: Type src ctx ss cs (h::k)
 -> (forall kx (x::kx) xs. Type src ctx ss cs x -> Types src ctx ss cs xs -> ret)
 -> ret
spineTy = go TypesZ
	where
	go :: Types src ctx ss cs hs
	   -> Type src ctx ss cs (h::k)
	   -> (forall kx (x::kx) xs. Type src ctx ss cs x -> Types src ctx ss cs xs -> ret)
	   -> ret
	go ts (TyApp  _src f a) k = go (a `TypesS` ts) f k
	go ts (Term x _te)    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 ctx ss cs h.
 Inj_Source (EType src '[] ss cs) src
 => Proj_TyConst cs (->)
 => Type src ctx ss cs h -> Maybe (KTerm src ctx ss cs, KTerm src ctx ss cs)
unTyFun t = go t
	where
	go :: forall x. Type src ctx ss cs x -> Maybe (KTerm src ctx ss cs, KTerm src ctx ss cs)
	go (TyApp _ (TyApp _ (TyConst _ f) a) b)
	 | Just KRefl <- proj_TyConstKi @_ @(->) f
	 = Just (KT (a `source` EType (unCtxTe t)), KT (b `source` EType (unCtxTe t)))
	go TyApp{}      = Nothing
	go TyConst{}    = Nothing
	go TyQuant{}    = Nothing
	go TyVar{}      = Nothing
	go TyFam{}      = Nothing
	go (Term x _te) = go x

-- | Return given 'Type', polymorphized by 'TyQuant'ifying
-- over the given 'Var' of given 'Kind' and 'NameHint'.
polyTyVar :: forall src ctx ss cs kv h.
 Source src
 => Show_TyConst cs
 => Proj_TyConst cs (->)
 => Proj_TyConst cs (#>)
 => NameHint -> Kind src kv -> Var
 -> Type src ctx ss cs (h::k)
 -> KT src ctx ss 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 ctx ss cs m
	 -> (forall (v::kv). Type src ctx ss cs v -> KT src ctx ss 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') a | v == v' =
		case eqKi kv kv' of
		 Nothing -> error "[BUG] polyTyVar: polymorphized TyVar has not expected kind"
		 Just Refl -> KT a
	go _v m@TyVar{} _a = KT m
	go v (TyFam src f hs) a =
		goArgs hs $ \hs' -> KT $ TyFam src f hs'
		where
		goArgs :: forall hs ret.
		    Types src ctx ss cs hs
		 -> (forall hs'. Types src ctx ss cs hs' -> ret)
		 -> ret
		goArgs TypesZ k = k TypesZ
		goArgs (TypesS h as) k =
			case go v h a of
			 KT h' -> goArgs as $ \as' -> k $ TypesS h' as'
	go v (Term m te) 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 p te
			-- 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 ctx ss cs (h::k)
 -> KT src ctx ss cs k
polyTy t =
	let fvs = freeKiTyVars t in
	Map.foldrWithKey
	 (\v (EKind kv) (KT x) ->
		polyTyVar "" kv v x)
	 (KT t) fvs