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

import Control.Arrow
import Control.Monad
import Data.Bool
import Data.Either (Either(..))
import Data.Eq (Eq(..))
import Data.Foldable (Foldable(..))
import Data.Function (($), (.), id, const, flip, on)
import Data.Functor ((<$>))
import Data.List ((++))
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (Maybe(..), fromMaybe)
import Data.Monoid ()
import Data.Ord (Ord(..))
import Data.Text.Buildable (Buildable(..))
import Data.Traversable (Traversable(..))
import Data.Typeable as Typeable
import Prelude (error)
import Text.Show (Show(..))

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

-- * Type 'Form'

-- | Return a 'Form' from a given 'Term'.
form :: (Eq var, Ord var, Variable var, Typeable var) => Context var -> Term var -> Form var
form ctx te = Form te (normalize ctx te)

-- | Reduce given 'Term' (or 'Type') to /normal form/ (NF)
-- by recursively performing /β-reduction/ and /η-reduction/.
-- 
-- Note that any well-typed 'Term' (i.e. for which 'type_of' returns a 'Type')
-- is /strongly normalizing/ (i.e. 'normalize' always returns,
-- and its returned 'Term' is unique up to /α-equivalence/).
normalize :: (Eq var, Ord var, Variable var, Typeable var) => Context var -> Term var -> Term var
normalize = go []
	where
	{-
	go_debug :: (Eq var, Ord var, Variable var, Typeable var)
	 => [Term var]
	 -> Context var
	 -> Term var
	 -> Term var
	go_debug args ctx te =
		go args ctx $
		trace ("normalize: "
		 ++ "\n te = " ++ show te
		 ++ "\n args = " ++ show args
		 ) $ te
	-}
	
	go :: (Eq var, Ord var, Variable var, Typeable var)
	 => [Term var]
	 -> Context var
	 -> Term var
	 -> Term var
	go args ctx (TeTy_Var ((form_normal <$>) . context_item_term <=< context_lookup ctx -> Just te))
	 -- NOTE: Replace variable mapped by the context
	 = case args of
		 [] -> te -- NOTE: no need to normalize again
		 _  -> go args ctx te
	go args ctx (TeTy_App f x)
	 -- NOTE: Normalize and collect applied arguments,
	 -- this to normalize an argument only once for all patterns.
	 = go (go [] ctx x : args) ctx f
	go [] ctx (Term_Abst x f_in f)
	 -- NOTE: η-reduce: Term_Abst _ (TeTy_App f (TeTy_Var (Var_Bound _)) ==> f
	 = (Term_Abst x (go [] ctx f_in) ||| id)
		 (abst_eta_reduce ctx f)
	go (x:args) ctx (Term_Abst _ _ f)
	 -- NOTE: β-reduce: TeTy_App (Term_Abst _ _ f) x ==> f x
	 = go args ctx (const x `unabstract` f)
	go args ctx (Type_Abst x f_in f_out)
	 -- NOTE: Recurse in Type_Abst fields
	 = term_apps (Type_Abst x (go [] ctx f_in) (go_scope ctx f_out)) args
	go args ctx (TeTy_Axiom (flip (axiom_normalize ctx) args -> Just (r_te, r_args)))
	 -- NOTE: Normalize axiom
	 = go r_args ctx r_te
	go args _ctx te
	 -- NOTE: Reapply remaining arguments, normalized
	 = term_apps te args
	
	abst_eta_reduce
	 :: (Eq var, Eq bound, Ord var, Ord bound, Variable var, Variable bound, Typeable var, Typeable bound)
	 => Context var
	 -> Abstraction bound Term var
	 -> Either (Abstraction bound Term var) -- could not η-reduce
	           (Term var)                     -- could     η-reduce
	abst_eta_reduce ctx =
		(\abst_body ->
			let new_ctx = context_push_nothing ctx in
			case go [] new_ctx abst_body of
			 te@(TeTy_App t (TeTy_Var (Var_Bound _{-Term_Abst's variable-}))) ->
				traverse (\var ->
					case var of
					 Var_Free v -> Right v
						 -- NOTE: decrement the DeBruijn indexing by one
						 -- to reflect the removal of the Term_Abst.
					 Var_Bound _ -> Left $ abstract_generalize te
						 -- NOTE: cannot η-reduce because Term_Abst's variable
						 -- is also used within t.
				 ) t
			 te -> Left $ abstract_generalize te
		 ) .
		abstract_normalize
	
	go_scope
	 :: (Eq var, Eq bound, Ord var, Ord bound, Variable var, Variable bound, Typeable var, Typeable bound)
	 => Context var
	 -> Abstraction bound Term var
	 -> Abstraction bound Term var
	go_scope ctx =
		abstract_generalize .
		go [] (context_push_nothing ctx) .
		abstract_normalize

-- | Reduce given 'Term' to /Weak Head Normal Form/ (WHNF).
--
-- __Ressources:__
--
-- * https://wiki.haskell.org/Weak_head_normal_form
whnf :: (Ord var, Variable var, Typeable var) => Context var -> Term var -> Term var
whnf = go []
	where
	go :: (Ord var, Variable var, Typeable var) => [Term var] -> Context var -> Term var -> Term var
	go args ctx (TeTy_Var ((form_given <$>) .
	 context_item_term <=< context_lookup ctx
	 -> Just te))
	 -- NOTE: Replace any variable mapped by the context
	 = go args ctx te
	go args ctx (TeTy_App f a)
	 -- NOTE: Collect applied arguments
	 = go (a:args) ctx f
	go (x:args) ctx (Term_Abst _ _ f)
	 -- NOTE: β-reduce: TeTy_App (Term_Abst _ _ f) x ==> f x
	 = go args ctx (const x `unabstract` f)
	go args ctx (TeTy_Axiom (flip (axiom_normalize ctx) args -> Just (r_te, r_args)))
	 -- NOTE: Normalize axiom
	 -- TODO: maybe introduce an 'axiom_whnf' instead of 'axiom_normalize'
	 = go r_args ctx r_te
	go args _ctx te
	 -- NOTE: Reapply remaining arguments, normalized
	 = term_apps te args

-- | Apply multiple 'Term's to a 'Term',
-- useful to reconstruct a 'Term' while normalizing (see 'normalize' or 'whnf').
term_apps :: Term var -> [Term var] -> Term var
term_apps = foldl TeTy_App

-- | /α-equivalence relation/, synonym of @(==)@.
--
-- Return 'True' iif. both given 'Term's are the same,
-- up to renaming the /bound variables/ it contains (see instance 'Eq' of 'Suggest').
alpha_equiv :: (Eq var, Show var) => Term var -> Term var -> Bool
alpha_equiv = (==)

-- | /αβη-equivalence relation/.
equiv :: (Eq var, Ord var, Variable var, Typeable var) => Context var -> Term var -> Term var -> Bool
equiv ctx = (==) `on` normalize ctx

-- * Type 'Context'

context :: Context Var_Name
context = Context [] (const Nothing) id id

context_apply :: Context var -> Term var -> Term var
context_apply ctx te =
	te >>= \var ->
		fromMaybe (TeTy_Var var) $ form_normal <$>
		(context_item_term =<< context_lookup ctx var)

context_push_type
 :: (Eq var, Ord var, Variable var, Typeable var)
 => Context var
 -> bound
 -> Type var
 -> Context (Var bound var)
context_push_type
 ctx@(Context keys lookup var_lift var_unlift)
 bound ty =
	Context
	 { context_vars =
		Var_Bound bound :
		Var_Free `fmap` keys
	 , context_lookup = \var ->
		(Var_Free `fmap`) `fmap`
			case var of
			 Var_Bound _ ->
				Just Context_Item
				 { context_item_term = Nothing
				 , context_item_type = form ctx ty
				 }
			 Var_Free v -> lookup v
	 , context_lift = Var_Free . var_lift
	 , context_unlift = \var ->
		case var of
		 Var_Bound _ -> error "context_push_type: context_unlift"
		 Var_Free v  -> var_unlift v
	 }

context_push_nothing
 :: (Show bound, Buildable (Context var), Typeable var, Typeable bound)
 => Context var
 -> Context (Var bound var)
context_push_nothing
 (Context keys lookup var_lift var_unlift) =
	Context
	 { context_vars = Var_Free `fmap` keys
	 , context_lookup = \var ->
		(Var_Free `fmap`) `fmap`
			case var of
			 Var_Bound _ -> Nothing
			 Var_Free v  -> lookup v
	 , context_lift = Var_Free . var_lift
	 , context_unlift = \var ->
		case var of
		 -- DEBUG: Var_Bound (cast -> Just (Suggest b)) -> b
		 -- DEBUG: Var_Bound (cast -> Just b) -> b
		 Var_Bound b -> error $ "context_push_nothing: context_unlift: " ++ show b
		 Var_Free v  -> var_unlift v
	 }

context_from_env :: Env -> Context Var_Name
context_from_env env =
	Context
	 { context_vars   = Map.keys env
	 , context_lookup = \var ->
		(\item -> Context_Item
		 { context_item_term = Just $ env_item_term item
		 , context_item_type =        env_item_type item
		 }) <$> env_lookup env var
	 , context_lift   = id
	 , context_unlift = id
	 }

context_relift :: forall from_var to_var.
 ( Typeable from_var
 , Typeable to_var )
 => Context from_var
 -> Type    from_var
 -> Context to_var
 -> Type    to_var
context_relift from_ctx ty to_ctx =
	( context_lift to_ctx
	. context_unlift from_ctx
	) <$> ty

{-
context_from_list :: Eq var => [(var, Type var)] -> Context var
context_from_list l =
	Context
	 { context_vars   = fst <$> l
	 , context_lookup_type = \var -> List.lookup var l
	 }

-- | Return /free term variables/ of an 'Abstraction'-ed 'Term'.
ftv
 :: Abstraction bound (Term) var
 -> Term (Var bound var)
ftv = abstract_normalize
-}

-- * Type 'Env'

type Env
 = Map Var_Name Env_Item
data Env_Item
 =   Env_Item
 {   env_item_term :: Form Var_Name
 ,   env_item_type :: Form Var_Name
 }
env_item
 :: Context Var_Name
 -> Term Var_Name
 -> Type Var_Name
 -> Env_Item
env_item ctx te ty =
	Env_Item
	 { env_item_term = form ctx te
	 , env_item_type = form ctx ty
	 }

env_lookup :: Env -> Var_Name -> Maybe Env_Item
env_lookup env var = Map.lookup var env

env_insert :: Var_Name -> Term Var_Name -> Type Var_Name -> Env -> Env
env_insert var te ty env =
	let ctx = context_from_env env in
	Map.insert var (env_item ctx te ty) env