{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Language.Symantic.Typing.Type
 ( module Language.Symantic.Typing.Type
 -- , Any
 , Proxy(..)
 , (:~:)(..)
 , (:~~:)(..)
 ) where

import Data.Maybe (isJust)
import Data.Proxy
import Data.Text (Text)
import Data.Type.Equality
import GHC.Exts (Constraint)
import GHC.Prim (Any)
import Data.String (IsString(..))
-- import Data.Typeable (Typeable, eqT)
import qualified Data.Kind as K

import Language.Symantic.Helper.Data.Type.Equality
import Language.Symantic.Grammar
import Language.Symantic.Typing.Kind
import Language.Symantic.Typing.Constant
import Language.Symantic.Parsing
import Language.Symantic.Transforming

-- * Type 'T', for both 'Type' and 'Term'

-- | /Type of terms/, also embedding… /terms/ to be able
-- to have /polymorphic terms/ (ie. handle /type variables/).
--
-- Parameters:
--
-- * @src@: /source/.
-- * @cs@: /type constants/, fixed at compile time.
-- * @h@: indexed Haskell type.
-- * @k@: indexed Haskell kind of @h@.
--
-- Constructors:
--
-- * 'TyConst': /type constant/, selected amongst @cs@.
--
-- * 'TyApp': /type parameter application/.
--
-- * 'TyQuant': /universal type quantification/ (aka. /type abstraction/):
--   used to introduce (and bind) a /type variable/.
--    
--   Note that the type indexing is lost (though not the kind @k@),
--   because GHC cannot represent a quantified type in a type
--   aka. /impredicative polymorphism/.
--
-- * 'TyVar': /type variable/:,
--   used to represent a /bounded type variable/,
--   when passing through a 'TyQuant'.
--    
--   Note that only the kind @k@ is flexible,
--   @h@ being set to 'Any' to loose the type indexing
--   and enabling the return of a type equality when the 'Var's are the same.
--    
--   Note also that there is no DeBruijn encoding done to link 'TyQuant's and 'TyVar's.
--
-- * 'Term': /term/, useful to embed a term
--   requiring quantifications and qualifications.
--
-- * 'TyFam': /type family/.
--
-- See also: https://ghc.haskell.org/trac/ghc/wiki/Typeable
-- which currently concludes:
-- "Why kind equalities, then? Given the fact that Richard's branch
-- doesn't solve this problem, what is it good for?
-- It still works wonders in the mono-kinded case,
-- such as for decomposing ->.
-- It's just that poly-kinded constructors are still a pain."
data T (src::K.Type)
       (ctx::[K.Type])
       -- (ss::([K.Type],[K.Type],[K.Type]))
       (ss::[K.Type])
       (cs::[K.Type])
       (h::k) where
	TyConst :: src -> TyConst src cs c
	               -> Type src ctx ss cs c
	TyApp   :: src -> Type src ctx ss cs f
	               -> Type src ctx ss cs a
	               -> Type src ctx ss cs (f a)
	               -- NOTE: @f a@ decomposition made possible
	               -- at the cost of requiring @TypeInType@.
	TyQuant :: src -> NameHint
	               -> Kind src kv
	               -> (forall (v::kv). Type src ctx ss cs v -> KType src ctx ss cs k)
	               -> Type src ctx ss cs (h::k)
	TyVar   :: src -> Kind src kv
	               -> Var
	               -> Type src ctx ss cs (Any::kv)
	TyFam   :: src -> TyConst src cs fam
	               -> Types src ctx ss cs hs
	               -> Type src ctx ss cs (TyFam fam hs)
	Term    :: Type src ctx ss cs h
	        -> TeSym ctx ss h
	        -> Term src ctx ss cs h

-- ** Type 'Type'
type Type = T

-- ** Type 'Name'
type Name     = Text
type NameHint = Name

-- ** Type 'Term'
type Term src ctx ss cs h = T src ctx ss cs (h::K.Type)

-- | Remove all top 'Term's.
--
-- Useful before pattern-matching.
unTerm :: Type src ctx ss cs (h::k) -> Type src ctx ss cs (h::k)
unTerm (Term x _te) = unTerm x
unTerm t = t

-- | Remove all 'Term's.
--
-- Useful before pattern-matching, or to reset the context.
unTerms :: Type src ctx ss cs (h::k) -> Type src ctx' ss cs (h::k)
unTerms (TyConst src c)      = TyConst src c
unTerms (TyApp src f a)      = TyApp src (unTerms f) (unTerms a)
unTerms (TyQuant src n kv f) = TyQuant src n kv $ \v -> case f (unTerms v) of KT t -> KT $ unTerms t
unTerms (TyVar src kv v)     = TyVar src kv v
unTerms (TyFam src f hs)     = TyFam src f (unTerms `mapTys` hs)
unTerms (Term x _te)         = unTerms x

unCtxTe :: T src ctx ss cs h -> T src '[] ss cs h
unCtxTe = unTerms

-- | Like 'TyConst', but 'sourceLess'.
-- 
-- NOTE: to be used with @TypeApplications@ and @NoMonomorphismRestriction@.
ty :: forall c ss cs src ctx.
      Source src
   => Inj_TyConst cs c
   => Type src ctx ss cs c
ty = TyConst sourceLess inj_TyConst
 -- NOTE: The forall brings @c@ first in the type variables,
 -- which is convenient to use @TypeApplications@.

-- | Like 'ty', but for a known 'Source'.
tySourced
 :: forall c ss cs src ctx.
    Source src
 => Inj_TyConst cs c
 => src
 -> Type src ctx ss cs c
tySourced src = TyConst src inj_TyConst

-- | Like 'TyApp', but 'sourceLess'.
tyApp :: Source src
      => Type src ctx ss cs f
      -> Type src ctx ss cs a
      -> Type src ctx ss cs (f a)
tyApp = TyApp sourceLess
infixl 9 `tyApp`

-- | Like 'TyQuant', but 'sourceLess'.
tyQuant
 :: forall src ctx ss cs kv k.
    Source src
 => Inj_Kind kv
 => NameHint
 -> (forall (v::kv). Type src ctx ss cs v -> KT src ctx ss cs k)
 -> KT src ctx ss cs k
tyQuant n f = KT $ TyQuant sourceLess n inj_Kind f

-- | Return the 'Kind' of the given 'Type'.
kindTy
 :: Inj_Source (EType src '[] ss cs) src
 => Type src ctx ss cs (h::k)
 -> Kind src k
kindTy t = kindTyOS t `source` EType (unCtxTe t)

-- | Like 'kindTy' but keep the old 'Source'.
kindTyOS :: Type src ctx ss cs (h::k) -> Kind src k
kindTyOS = go
	where
	go :: Type src ctx ss cs (h::k) -> Kind src k
	go = \case
	 TyConst _src c -> kindTyConst c
	 TyApp _src f _x ->
		case kindTyOS f of
		 KiArrow _src_kf _kx kf -> kf
	 TyQuant src _n k f -> kindTyOS `unKT` f (TyVar src k 0)
	 TyVar  _src k _v   -> k
	 TyFam _src f _hs   -> kindTyConst f
	 Term x _te         -> kindTyOS x

instance -- TestEquality
 Inj_Source (EType src ctx ss cs) src
 => TestEquality (Type src ctx ss cs) where
	testEquality = eqTy
instance Eq (Type src ctx ss cs h) where
	_x == _y = True

-- | Return a proof when two 'Type's are equals.
--
-- NOTE: 'TyQuant's are not compared:
-- they first have to be monomorphized.
eqTy :: Type src ctx ss cs (x::k) -> Type src ctx ss cs (y::k) -> Maybe (x:~:y)
eqTy = go
	where
	go :: Type src ctx ss cs (x::k) -> Type src ctx ss cs (y::k) -> Maybe (x:~:y)
	go (TyConst _sx x) (TyConst _sy y)
	 | Just Refl <- testEquality x y
	 = Just Refl
	go (TyApp _sx fx xx) (TyApp _sy fy xy)
	 | Just Refl <- eqKi (kindTyOS fx) (kindTyOS fy)
	 , Just Refl <- go fx fy
	 , Just Refl <- go xx xy
	 = Just Refl
	go TyQuant{} TyQuant{} = Nothing
	{-
	go v (TyQuant sx _nx kx x::Type src ctx ss cs vx)
	     (TyQuant sy _ny ky y::Type src ctx ss cs vy)
	 | Just Refl <- eqKi kx ky
	 , Just Refl <- go (v - 1) (x $ TyVar sx kx v) (y $ TyVar sy ky v)
	 = Just Refl
	-}
	go (TyVar _sx kx x) (TyVar _sy ky y)
	 | Just Refl <- eqKi kx ky
	 , True <- x == y
	 = Just Refl
	-- NOTE: 'Term' are transparent to 'eqTy'.
	go (Term x _tx) (Term y _ty) | Just Refl <- go x y = Just Refl
	go (Term x _tx) y            | Just Refl <- go x y = Just Refl
	go x            (Term y _ty) | Just Refl <- go x y = Just Refl
	go _x _y = Nothing

-- | Return two proofs when two 'Type's are equals:
-- one for the type and one for the kind.
eqKiTy :: Type src ctx ss cs (x::kx) -> Type src ctx ss cs (y::ky) -> Maybe (x:~~:y)
eqKiTy = go
	where
	go :: Type src ctx ss cs (x::kx) -> Type src ctx ss cs (y::ky) -> Maybe (x:~~:y)
	go (TyConst _sx x) (TyConst _sy y) = eqKiTyConst x y
	go (TyApp _sx fx xx) (TyApp _sy fy xy)
	 | Just KRefl <- go fx fy
	 , Just KRefl <- go xx xy
	 = Just KRefl
	go TyQuant{} TyQuant{} = Nothing
	go (TyVar _sx kx _x) (TyVar _sy ky _y)
	 | Just Refl <- eqKi kx ky
	 = Just KRefl
	go (Term x _tx) (Term y _ty) = go x y
	go _x _y = Nothing

-- * Type @(#>)@
-- | 'TyConst' to qualify a 'Type'.
data (#>) (q::Constraint) (a::K.Type)
instance Fixity_TyConstC (#>) where
	fixity_TyConstC _ = FixityInfix $ infixR 0

(#>) :: Source src
     => Inj_TyConst cs (#>)
     => Type src ctx ss cs a
     -> Type src ctx ss cs b
     -> Type src ctx ss cs (a #> b)
(#>) a b = (ty @(#>) `tyApp` a) `tyApp` b
infixr 0 #>
 -- NOTE: should actually be (-1)
 -- to compose well with @infixr 0 (->)@
 -- but this is not allowed by GHC.

-- * Type family 'TyFam'
-- | Apply the /type family/ selected by @fam@
-- to a list of types (within a 'Proxy').
type family TyFam (fam::k) (hs::[K.Type]) :: k
type family TyFamKi fam :: K.Type

-- ** Type 'TyFamName'
type TyFamName = Text

-- * Type 'TeSym'
data TeSym ctx ss (h::K.Type)
 =   TeSym
 ( forall term. Syms ss term =>
                Sym_Lambda term =>
                CtxTe term ctx -> term h )

-- * Class 'Sym_Lambda'
class Sym_Lambda term where
	-- | /Function application/.
	apply :: term ((a -> b) -> a -> b)
	
	-- | /Lambda application/.
	app :: term (a -> b) -> (term a -> term b); infixr 0 `app`
	default app :: Trans t term => t term (arg -> res) -> t term arg -> t term res
	app f x = trans_lift (app (trans_apply f) (trans_apply x))
	
	-- | /Lambda abstraction/.
	lam :: (term a -> term b) -> term (a -> b)
	default lam :: Trans t term => (t term arg -> t term res) -> t term (arg -> res)
	lam f = trans_lift (lam (trans_apply . f . trans_lift))
	
	-- | Convenient 'lam' and 'app' wrapper.
	let_ :: term var -> (term var -> term res) -> term res
	let_ x y = lam y `app` x

-- * Type 'TeName'
newtype TeName = TeName Text
 deriving (Eq, Ord, Show)
instance IsString TeName where
	fromString = TeName . fromString

-- * Type 'CtxTe'
-- | GADT for an /interpreting context/:
-- accumulating at each /lambda abstraction/
-- the @term@ of the introduced variable.
data CtxTe (term::K.Type -> K.Type) (hs::[K.Type]) where
	CtxTeZ :: CtxTe term '[]
	CtxTeS :: term h
	       -> CtxTe term hs
	       -> CtxTe term (h ': hs)
infixr 5 `CtxTeS`

-- ** Type family 'Sym'
type family Sym (s::K.Type) :: {-term-}(K.Type -> K.Type) -> Constraint

-- ** Type family 'Syms'
type family Syms (ss::[K.Type]) (term:: K.Type -> K.Type) :: Constraint where
	Syms '[] term = ()
	Syms (s ': ss) term = (Sym s term, Syms ss term)


instance Source src => Sourced (Type src ctx ss cs h) where
	type Source_of (Type src ctx ss cs h) = src
	
	source_of (TyConst src _c)        = src
	source_of (TyApp   src _f _x)     = src
	source_of (TyQuant src _n _kv _f) = src
	source_of (TyVar   src _kv _v)    = src
	source_of (TyFam   src _f _hs)    = src
	source_of (Term x _te)            = source_of x
	
	source_is (TyConst _src c)      src = TyConst src c
	source_is (TyApp   _src f x)    src = TyApp   src f x
	source_is (TyQuant _src n kv f) src = TyQuant src n kv f
	source_is (TyVar   _src kv v)   src = TyVar   src kv v
	source_is (TyFam   _src f hs)   src = TyFam   src f hs
	source_is (Term x te)           src = Term (source_is x src) te

-- ** Type 'Var'
type Var = Int

-- * Type 'Con_Type'
data Con_Type src ctx ss cs
 =   Con_EqType  (EType src ctx ss cs) (EType src ctx ss cs)
 |   Con_TyApp   (EType src ctx ss cs)
 |   Con_TyCon   (KT src ctx ss cs Constraint)
 |   Con_TyFam   (At src TyFamName) [EType src ctx ss cs]
deriving instance
 ( Eq src
 , Inj_Source (EType src ctx ss cs) src
 ) => Eq (Con_Type src ctx ss cs)
deriving instance
 ( Show src
 , Show_TyConst cs
 , Show (EType src ctx ss cs)
 , Show (KT    src ctx ss cs Constraint)
 , Fixity_TyConst cs
 ) => Show (Con_Type src ctx ss cs)

if_TyApp
 :: ( Inj_Error (Con_Type src '[] ss cs) err
    , Inj_Source (EType src '[] ss cs) src )
 => Type src ctx ss cs (fa::kfa)
 -> (forall ka f a. (fa :~: f a)
  -> Type src ctx ss cs (f::ka -> kfa)
  -> Type src ctx ss cs (a::ka)
  -> Either err ret)
 -> Either err ret
if_TyApp typ k =
	case typ of
	 TyApp _src a b -> k Refl (a `source` EType (unCtxTe typ)) (b `source` EType (unCtxTe typ))
	 _ -> Left $ inj_Error $ Con_TyApp (EType (unCtxTe typ))

if_EqType
 :: Inj_Error (Con_Type src '[] ss cs) err
 => Type src ctx ss cs x
 -> Type src ctx ss cs y
 -> ((x :~: y) -> Either err ret) -> Either err ret
if_EqType x y k =
	case x `eqTy` y of
	 Just Refl -> k Refl
	 Nothing -> Left $ inj_Error $
		Con_EqType (EType (unCtxTe x)) (EType (unCtxTe y))

if_EqType1
 :: ( Inj_Error (Con_Type src '[] ss cs) err
    , Inj_Error (Con_Kind src) err
    , Inj_Source (EType src '[] ss cs) src
    )
 => Type src ctx ss cs (f:: K.Type -> K.Type)
 -> Type src ctx ss cs fa
 -> (forall a. (fa :~: f a)
  -> Type src ctx ss cs a
  -> Either err ret)
 -> Either err ret
if_EqType1 typ ty_fa k =
	if_TyApp ty_fa $ \Refl ty_f ty_a ->
	if_EqKind (inj_Kind @(K.Type -> K.Type))
	          (kindTy ty_f) $ \Refl ->
	if_EqType typ ty_f $ \Refl ->
	k Refl ty_a

if_EqType2
 :: ( Inj_Error (Con_Type src '[] ss cs) err
    , Inj_Error (Con_Kind src) err
    , Inj_Source (EType src '[] ss cs) src )
 => Type src ctx ss cs (f:: K.Type -> K.Type -> K.Type)
 -> Type src ctx ss cs fab
 -> (forall a b. (fab :~: f a b)
  -> Type src ctx ss cs a
  -> Type src ctx ss cs b
  -> Either err ret)
 -> Either err ret
if_EqType2 typ ty_fab k =
	if_TyApp ty_fab $ \Refl ty_fa ty_b ->
	if_TyApp ty_fa  $ \Refl ty_f  ty_a ->
	if_EqKind (inj_Kind @(K.Type -> K.Type -> K.Type))
	          (kindTy ty_f) $ \Refl ->
	if_EqType typ ty_f $ \Refl ->
	k Refl ty_a ty_b

if_TyFun ::
 ( Inj_TyConst cs (->)
 , Inj_Source (EType src '[] ss cs) src
 , Inj_Error (Con_Kind src) err
 , Inj_Error (Con_Type src '[] ss cs) err
 ) => Type src ctx ss cs fab
 -> (forall a b. (:~:) fab (a -> b)
    -> Type src ctx ss cs a -> Type src ctx ss cs b -> Either err ret)
 -> Either err ret
if_TyFun = if_EqType2 (ty @(->))

-- * Type 'EType'
-- | Existential for 'Type'.
data EType src ctx ss cs
 = forall h. EType (Type src ctx ss cs h)

instance
 Inj_Source (EType src ctx ss cs) src =>
 Eq (EType src ctx ss cs) where
	EType x == EType y = isJust $ eqKiTy x y

mapEType
 :: (forall k (h::k). Type src ctx ss cs h -> Type src ctx ss cs h)
 -> EType src ctx ss cs
 -> EType src ctx ss cs
mapEType f (EType t) = EType (f t)

bindEType
 :: (forall k (h::k). Type src ctx ss cs h -> EType src ctx ss cs)
 -> EType src ctx ss cs
 -> EType src ctx ss cs
bindEType f (EType t) = f t

-- * Type 'KT'
-- | Existential for 'T' of a known 'Kind'.
data KT src ctx ss cs k = forall (h::k). KT (Type src ctx ss cs h)
type KType = KT
type KTerm src ctx ss cs = KT src ctx ss cs K.Type

instance Eq (KT src ctx ss cs ki) where
	KT x == KT y = isJust $ eqTy x y
instance
 Inj_Source (EType src ctx ss cs) src =>
 Inj_Source (KT src ctx ss cs k) src where
	inj_Source (KT t) = inj_Source $ EType t

unKT :: (forall (h::k). Type src ctx ss cs h -> a k) -> KT src ctx ss cs k -> a k
unKT f (KT t) = f t

ofKT :: (forall (h::k). Type src ctx ss cs h -> a) -> KT src ctx ss cs k -> a
ofKT f (KT t) = f t

mapKT
 :: (forall (h::k). Type src ctx ss cs h -> Type src ctx ss cs h)
 -> KT src ctx ss cs k
 -> KT src ctx ss cs k
mapKT f (KT t) = KT (f t)

bindKT
 :: (forall (h::k). Type src ctx ss cs h -> KT src ctx ss cs k)
 -> KT src ctx ss cs k
 -> KT src ctx ss cs k
bindKT f (KT t) = f t

bind2KT
 :: (forall (x::ka) (y::kb). Type src ctx ss cs x -> Type src ctx ss cs y -> KT src ctx ss cs kxy)
 -> KT src ctx ss cs ka
 -> KT src ctx ss cs kb
 -> KT src ctx ss cs kxy
bind2KT f (KT x) (KT y) = f x y

-- * Type 'Types'
data Types src ctx ss cs (hs::[K.Type]) where
	TypesZ :: Types src ctx ss cs '[]
	TypesS :: Type  src ctx ss cs h
	       -> Types src ctx ss cs hs
	       -> Types src ctx ss cs (Proxy h ': hs)
infixr 5 `TypesS`

eTypes :: Types src ctx ss cs hs -> [EType src ctx ss cs]
eTypes TypesZ = []
eTypes (TypesS t ts) = EType t : eTypes ts

mapTys
 :: (forall k (h::k). Type src ctx ss cs h -> Type src ctx' ss cs h)
 -> Types src ctx ss cs hs
 -> Types src ctx' ss cs hs
mapTys _f TypesZ = TypesZ
mapTys f (TypesS h hs) = TypesS (f h) (mapTys f hs)

foldlTys
 :: (forall k (h::k). Type src ctx ss cs h -> acc -> acc)
 -> Types src ctx ss cs hs
 -> acc -> acc
foldlTys _f TypesZ acc = acc
foldlTys f (TypesS h hs) acc = foldlTys f hs (f h acc)

foldrTys
 :: (forall k (h::k). Type src ctx ss cs h -> acc -> acc)
 -> Types src ctx ss cs hs
 -> acc -> acc
foldrTys _f TypesZ acc = acc
foldrTys f (TypesS h hs) acc = f h (foldrTys f hs acc)

{-
traverseTys :: (Traversable t, Applicative f) => (a -> f b) -> t a -> f (t b)
traverseTys f TypesZ        = TypesZ
traverseTys f (TypesS h hs) = TypesS <$> f h <*> traverseTys f hs
-}

-- * Type 'TypesKi'
data TypesKi src ctx ss cs hs where
	TypesKiZ :: TypesKi src ctx ss cs '[]
	TypesKiS :: Type src ctx ss cs (h::k)
	         -> TypesKi src ctx ss cs hs
	         -> TypesKi src ctx ss cs ((Proxy h, k) ': hs)
infixr 5 `TypesKiS`

eTypesKi :: TypesKi src ctx ss cs hs -> [EType src ctx ss cs]
eTypesKi TypesKiZ = []
eTypesKi (TypesKiS t ts) = EType t : eTypesKi ts

-- * Type 'UnProxy'
type family UnProxy (x::K.Type) :: k where
	UnProxy (Proxy x) = x