{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fno-warn-tabs #-}
module Language.LOL.Calculus.Type where

import Control.Arrow
import Control.Monad
import Data.Either (Either(..))
import Data.Eq (Eq(..))
import Data.Foldable (Foldable(..))
import Data.Function (($), (.), const)
import Data.Functor ((<$>))
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (Maybe(..))
import Data.Monoid (Monoid(..), (<>))
import Data.Ord (Ord(..))
import Data.Text.Buildable (Buildable(..))
import Data.Traversable (Traversable(..))
import Data.Typeable as Typeable
import Text.Show (Show(..))

import Language.LOL.Calculus.Abstraction
import Language.LOL.Calculus.Term
import Language.LOL.Calculus.Form

-- * Type 'Type'

-- | Construct the 'Type' of the given 'Term',
-- effectively checking for the /well-formedness/
-- and /well-typedness/ of a 'Term' (or 'Type').
--
-- Note that a 'Type' is always to be considered
-- according to a given 'Context':
-- 'type_of' applied to the same 'Term'
-- but on different 'Context's
-- may return a different 'Type' or 'Type_Error'.
type_of
 :: (Eq var, Ord var, Variable var, Typeable var)
 => Context var
 -> Term var
 -> Either Type_Error (Type var)
type_of ctx term =
	case term of
	 Type_Sort s -> (err +++ Type_Sort) $ sort_of_sort s
	 TeTy_Var v ->
		case form_normal
		 .   context_item_type
		 <$> context_lookup ctx v of
		 Just ty -> return ty
		 Nothing -> Left $ err $ Type_Error_Msg_Unbound_variable v
	 TeTy_Axiom ax ->
		return $ axiom_type_of ctx ax
	 TeTy_App f x -> do
		f_ty <- whnf ctx <$> type_of ctx f
		(f_in, f_out) <-
			case f_ty of
			 Type_Abst _ i o -> return (i, o)
			 _ -> Left $ err $ Type_Error_Msg_Not_a_function f f_ty
		x_ty <- type_of ctx x
		if equiv ctx x_ty f_in
			then return $ const x `unabstract` f_out
			else Left $ err $ Type_Error_Msg_Function_argument_mismatch f_in x_ty
	 Term_Abst (Suggest x) f_in f -> do
		_ <- type_of ctx f_in
		let new_ctx =
			if x == ""
			then context_push_nothing ctx
			else context_push_type ctx (Suggest x) f_in
		f_out <- type_of new_ctx (abstract_normalize f)
		let abst_ty = Type_Abst (Suggest x) f_in (abstract_generalize f_out)
		_ <- type_of ctx abst_ty
		return abst_ty
	 Type_Abst (Suggest x) f_in f -> do
		f_in_ty <- type_of ctx f_in
		f_in_so <- case whnf ctx f_in_ty of
		 Type_Sort s -> return s
		 f_in_ty_whnf -> Left $ err $ Type_Error_Msg_Invalid_input_type f_in_ty_whnf
		let new_ctx =
			if x == ""
			then context_push_nothing ctx
			else context_push_type ctx (Suggest x) f_in
		f_out_ty <- type_of new_ctx (abstract_normalize f)
		f_out_so <- case whnf new_ctx f_out_ty of
		 Type_Sort s -> return s
		 f_out_ty_whnf -> Left $ Type_Error ctx term $
			Type_Error_Msg_Invalid_output_type f_out_ty_whnf (x, f_out_ty_whnf)
		(err +++ Type_Sort) $
			sort_of_type_abst f_in_so f_out_so
	where
	err = Type_Error ctx term

-- | Check that the given 'Term' has the given 'Type'.
check
 :: (Eq var, Ord var, Variable var, Typeable var)
 => Context var
 -> Type var
 -> Term var
 -> Either Type_Error ()
check ctx expect_ty te =
	type_of ctx te >>= \actual_ty ->
		if equiv ctx expect_ty actual_ty
		then Right ()
		else Left Type_Error
			 { type_error_ctx  = ctx
			 , type_error_term = te
			 , type_error_msg  = Type_Error_Msg_Type_mismatch expect_ty actual_ty (normalize ctx expect_ty) (normalize ctx actual_ty)
			 }

-- | Check that a 'Term' is closed, i.e. has no /unbound variables/.
close :: Term Var_Name -> Either Type_Error (Term ())
close te =
	traverse go te
	where
		go var = Left $
			Type_Error context te $
			Type_Error_Msg_Unbound_variable var

-- | Return the /unbound variables/ of given 'Term'.
unbound_vars :: Ord var => Term var -> Map var ()
unbound_vars = foldr (`Map.insert` ()) mempty

-- | /Dependent product/ rules: @s ↝ t : u@, i.e.
--   "abstracting something of type @s@ out of something of type @t@ gives something of type @u@".
--
-- Given two 'Sort': @s@ and @t@, return a 'Sort': @u@,
-- when ('Type_Abst' @s@ @t@) is ruled legal
-- and has 'Type': 'Type_Sort' ('Sort' @u@).
--
-- The usual /PTS/ rules for /λω/
-- (or equivalently /Type Assignment Systems/ (TAS) rules for /System Fω/)
-- are used here:
--
-- * RULE: @⊦ * ↝ * : *@, aka. /simple types/:
--   "abstracting a term out of a term is valid and gives a term",
--   as in /PTS λ→/ or /TAS F1/.
-- * RULE: @⊦ □ ↝ * : *@, aka. /parametric polymorphism/:
--   "abstracting a type out of a term is valid and gives a term",
--   as in /PTS λ2/ or /TAS F2/ (aka. /System F/).
-- * RULE: @⊦ □ ↝ □ : □@, aka. /constructors of types/:
--   "abstracting a type out of a type is valid and gives a type",
--   as in /PTS λω/ or /TAS Fω/.
--
-- Note that the fourth usual rule is not ruled valid here:
--
-- * RULE: @⊦ * ↝ □ : □@, aka. /dependent types/:
--   "abstracting a term out of a type is valid and gives a type",
--   as in /PTS λPω/ or /TAS DFω/ (aka. /Calculus of constructions/).
--
-- However, to contain /impredicativity/ (see 'Axiom_MonoPoly')
-- the above /sort constants/ are split in two,
-- and the above rules adapted
-- to segregate between /monomorphic types/ (aka. /monotypes/)
-- and /polymorphic types/ (aka. /polytypes/):
--
-- * RULE: @⊦ *m ↝ *m : *m@, i.e. /simple types/, without /parametric polymorphism/.
-- * RULE: @⊦ *m ↝ *p : *p@, i.e. /simple types/, preserving /parametric polymorphism/ capture.
--
-- * RULE: @⊦ *p ↝ *m : *p@, i.e. /higher-rank polymorphism/, preserving /parametric polymorphism/ capture.
-- * RULE: @⊦ *p ↝ *p : *p@, i.e. /higher-rank polymorphism/, preserving /parametric polymorphism/ capture.
--
-- * RULE: @⊦ □m ↝ *m : *p@, i.e. /parametric polymorphism/, captured within @*p@ ('sort_star_poly').
-- * RULE: @⊦ □m ↝ *p : *p@, i.e. /parametric polymorphism/, preserving capture.
--
-- * RULE: @⊦ □m ↝ □m : □m@, i.e. /constructors of types/, without /parametric polymorphism/.
-- * RULE: @⊦ □m ↝ □p : □p@, i.e. /constructors of types/, preserving /parametric polymorphism/ capture.
--
-- Note that what is important here is
-- that there is no rule of the form: @⊦ □p ↝ _ : _@,
-- which forbids abstracting a /polymorphic type/ out of anything,
-- in particular the type @*p -> *m@ is forbidden,
-- though 'Axiom_MonoPoly'
-- is given to make it possible within it.
--
-- __Ressources:__
--
-- * /Henk: a typed intermediate language/,
--   Simon Peyton Jones, Erik Meijer, 20 May 1997,
--   https://research.microsoft.com/en-us/um/people/simonpj/papers/henk.ps.gz
sort_of_type_abst
 :: Sort
 -> Sort
 -> Either Type_Error_Msg Sort

-- Simple types
sort_of_type_abst
          (Type_Level_0, Type_Morphism_Mono)
          (Type_Level_0, m)
 = return (Type_Level_0, m)
 -- RULE: *m ↝ *m : *m
 -- RULE: *m ↝ *p : *p
 -- abstracting: a MONOMORPHIC term
 -- out of     : a MONOMORPHIC (resp. POLYMORPHIC) term
 -- forms      : a MONOMORPHIC (resp. POLYMORPHIC) term

-- Higher-rank
sort_of_type_abst
          (Type_Level_0, Type_Morphism_Poly)
          (Type_Level_0, _)
 = return (Type_Level_0, Type_Morphism_Poly)
 -- RULE: *p ↝ *m : *p
 -- RULE: *p ↝ *p : *p
 -- abstracting: a POLYMORPHIC term
 -- out of     : a             term
 -- forms      : a POLYMORPHIC term

-- Polymorphism
sort_of_type_abst
          (Type_Level_1, Type_Morphism_Mono)
          (Type_Level_0, _)
 = return (Type_Level_0, Type_Morphism_Poly)
 -- RULE: □m ↝ *m : *p
 -- RULE: □m ↝ *p : *p
 -- abstracting: a MONOMORPHIC type of a term
 -- out of     : a             term
 -- forms      : a POLYMORPHIC term

-- Type constructors
sort_of_type_abst
          (Type_Level_1, Type_Morphism_Mono)
          (Type_Level_1, m)
 = return (Type_Level_1, m)
 -- RULE: □m ↝ □m : □m
 -- RULE: □m ↝ □p : □p
 -- abstracting: a MONOMORPHIC type of a term
 -- out of     : a MONOMORPHIC (resp. POLYMORPHIC) type of a term
 -- forms      : a MONOMORPHIC (resp. POLYMORPHIC) type of a term
 --
 -- NOTE: □m ↝ □p : □p is useful for instance to build List_:
 -- let List_ : (A:*m) -> *p = \(A:*m) -> (R:*m) -> (A -> R -> R) -> R -> R
 -- let List  : (A:*m) -> *m = \(A:*m) -> Monotype (List_ A)

-- Dependent types

{-
sort_of_type_abst
          (Type_Level_0, Type_Morphism_Mono)
          (Type_Level_1, m)
 = return (Type_Level_1, m)
 -- RULE: *m ↝ □m : □m
 -- RULE: *m ↝ □p : □p
 -- abstracting: a MONOMORPHIC term
 -- out of     : a MONOMORPHIC (resp. POLYMORPHIC) type of a term
 -- forms      : a MONOMORPHIC (resp. POLYMORPHIC) type of a term

sort_of_type_abst
 (Type_Level_0, Type_Morphism_Poly)
 (Type_Level_1, _)
 = return (Type_Level_1, Type_Morphism_Poly)
 -- RULE: *p ↝ □m : □p
 -- RULE: *p ↝ □p : □p
 -- abstracting: a POLYMORPHIC term
 -- out of     : a             type of a term
 -- forms      : a POLYMORPHIC type of a term
-}

sort_of_type_abst
 s@(Type_Level_0, _)
 t@(Type_Level_1, _)
 = Left $ Type_Error_Msg_Illegal_type_abstraction s t
 -- RULE: * ↝ □ : Illegal
 -- abstracting: a term
 -- out of     : a type of a term
 -- is illegal

-- No impredicativity (only allowed through 'Axiom_MonoPoly')
sort_of_type_abst
 s@(Type_Level_1, Type_Morphism_Poly)
 t@(_, _)
 = Left $ Type_Error_Msg_Illegal_type_abstraction s t
 -- RULE: □p ↝ _ : Illegal
 -- abstracting: a POLYMORPHIC type of a term
 -- out of     : anything
 -- is illegal

-- ** Type 'Type_Error'
data Type_Error
 = forall var. (Ord var, Show var, Buildable var)
 =>  Type_Error
 {   type_error_ctx  :: Context var
 ,   type_error_term :: Term var
 ,   type_error_msg  :: Type_Error_Msg
 }
deriving instance Show Type_Error
instance Buildable Type_Error where
	build (Type_Error ctx te msg) =
		"Error: Type_Error"
		 <> "\n " <> build msg
		 <> "\n Term:" <> " " <> build te
		 <> (
			let vars =
				Map.keys $
				Map.intersection
				 (unbound_vars te)
				 (Map.fromList $ (, ()) <$> context_vars ctx) in
			case vars of
			 [] -> "\n"
			 _ -> "\n Context:\n" <> build ctx{context_vars=vars}
		 )

-- ** Type 'Type_Error_Msg'
data Type_Error_Msg
 =   Type_Error_Msg_No_sort_for_sort Sort
 |   Type_Error_Msg_Illegal_type_abstraction Sort Sort
 |   forall var. Variable var => Type_Error_Msg_Invalid_input_type  (Type var)
 |   forall var. Variable var => Type_Error_Msg_Invalid_output_type (Type var) (Var_Name, Type var)
 |   forall var. Variable var => Type_Error_Msg_Not_a_function (Term var) (Type var)
 |   forall var. Variable var => Type_Error_Msg_Function_argument_mismatch (Type var) (Type var)
 |   forall var. Variable var => Type_Error_Msg_Unbound_variable var
 |   forall var. Variable var => Type_Error_Msg_Unbound_axiom var
 |   forall var. Variable var => Type_Error_Msg_Type_mismatch (Type var) (Type var) (Type var) (Type var)
deriving instance Show Type_Error_Msg
instance Buildable Type_Error_Msg where
	build msg =
		case msg of
		 Type_Error_Msg_No_sort_for_sort x ->
			"No_sort_for_sort: "
			 <> build x
		 Type_Error_Msg_Illegal_type_abstraction x y ->
			"Illegal_type_abstraction: "
			 <> build x <> " -> " <> build y
		 Type_Error_Msg_Invalid_input_type ty ->
			"Invalid_input_type: "
			 <> build ty
		 Type_Error_Msg_Invalid_output_type f_out (x, f_in) ->
			"Invalid_output_type: "
			 <> build f_out <> "\n"
			 <> " Input binding: "
			 <> "(" <> build x <> " : " <> build f_in <> ")"
		 Type_Error_Msg_Not_a_function f f_ty ->
			"Not_a_function: "
			 <> build f
			 <> " : "
			 <> build f_ty
		 Type_Error_Msg_Function_argument_mismatch f_in x_ty ->
			"Function_argument_mismatch: \n"
			 <> " Function domain: " <> build f_in <> "\n"
			 <> " Argument type:   " <> build x_ty
		 Type_Error_Msg_Unbound_variable var ->
			"Unbound_variable: "
			 <> build var
		 Type_Error_Msg_Unbound_axiom var ->
			"Unbound_axiom: "
			 <> build var
		 Type_Error_Msg_Type_mismatch x y nx ny ->
			"Type_mismatch: \n"
			 <> " Expected type: " <> build x <> " == " <> build nx <> "\n"
			 <> " Actual   type: " <> build y <> " == " <> build ny

-- ** Type 'Sort'

-- *** Type 'Type_Level'

-- | /PTS/ axioms for 'Sort':
--
-- * AXIOM: @⊦ *m : □m@
-- * AXIOM: @⊦ *p : □p@
sort_of_sort :: Sort -> Either Type_Error_Msg Sort
sort_of_sort (Type_Level_0, Type_Morphism_Mono)
 = return (Type_Level_1, Type_Morphism_Mono)
 -- AXIOM: @*m : □m@
 -- The type of MONOMORPHIC types of terms,
 -- is of type: the type of types of MONOMORPHIC types of terms
sort_of_sort (Type_Level_0, Type_Morphism_Poly)
 = return (Type_Level_1, Type_Morphism_Poly)
 -- AXIOM: @*p : □p@
 -- The type of POLYMORPHIC types of terms,
 -- is of type: the type of types of POLYMORPHIC types of terms
sort_of_sort s = Left $ Type_Error_Msg_No_sort_for_sort s