{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# OPTIONS_GHC -fno-warn-tabs #-}
module Language.LOL.Typing.Type.Substitution where

import Data.Bool
import Data.Either (Either(..), either)
import Data.Eq (Eq(..))
import Data.Function (($), (.))
import Data.Functor ((<$>))
import qualified Data.List as List
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (Maybe(..), fromMaybe, maybe)
import Data.Monoid (Monoid(..), (<>))
import Data.Sequence (Seq)
import qualified Data.Set as Set
import Data.Text.Buildable (Buildable(..))
import Prelude (Num(..), error)
import Text.Show (Show(..))

import Language.LOL.Typing.Type.Monotype
import Language.LOL.Typing.Lib.Data.Empty (Empty(..))

-- * Class 'Substitution'

-- | A /monomorphic type substitution/.
--
-- TODO: maybe rename to Monosub?
class Substitution s where
	substitution_codomain        :: s -> [Monotype]
	substitution_domain          :: s -> [Monovar]
	substitution_domain_remove   :: [Monovar] -> s -> s
	substitution_domain_restrict :: [Monovar] -> s -> s
	substitution_lookup          :: Monovar -> s -> Maybe Monotype

-- ** Type 'Substitution_Finite'

-- | A 'Substitution' represented by a finite map.
type Substitution_Finite = Map Monovar Monotype

instance Empty Substitution_Finite where
	empty = mempty
instance Buildable Substitution_Finite where
	build =
		mconcat . List.intersperse "\n" .
		Map.foldrWithKey (\var ty ->
			((build (Monotype_Var var) <> " == " <> build ty) :)
		 ) []
instance Substitution Substitution_Finite where
	substitution_codomain = Map.elems
	substitution_domain   = Map.keys
	substitution_domain_remove vars s =
		List.foldr Map.delete s vars
	substitution_domain_restrict vars =
		Map.filterWithKey (\v _ -> Set.member v vars_kept)
		where vars_kept = Set.fromList vars
	substitution_lookup = Map.lookup

substitution_finite_singleton :: Monovar -> Monotype -> Substitution_Finite
substitution_finite_singleton = Map.singleton

substitution_finite :: [(Monovar, Monotype)] -> Substitution_Finite
substitution_finite = Map.fromList

-- | Safely unify two 'Substitution_Finite's.
--
-- NOTE: the union is left-biased:
-- it prefers the first 'Substitution_Finite'
-- when duplicate keys are encountered.
substitution_finite_union :: Substitution_Finite -> Substitution_Finite -> Substitution_Finite
substitution_finite_union x y = x `Map.union` ((x `substitute`) <$> y)

-- | Unsafely unify two 'Substitution_Finite's.
substitution_finite_union_unsafe :: Substitution_Finite -> Substitution_Finite -> Substitution_Finite
substitution_finite_union_unsafe = Map.union

-- ** Type 'Substitution_Fixpoint'

-- | A fixpoint is computed when looking up
-- the target of a 'Monovar' in this 'Substitution'.
-- 'mappend'ing two 'Substitution's is cheap,
-- whereas a 'substitution_lookup' is more expensive
-- than with 'Substitution_Finite'.
newtype Substitution_Fixpoint
 =      Substitution_Fixpoint (Map Monovar Monotype)
 deriving (Show)

instance Empty Substitution_Fixpoint where
	empty = Substitution_Fixpoint Map.empty
instance Buildable Substitution_Fixpoint where
	build (Substitution_Fixpoint sub) = build sub
-- | WARNING: only disjoint 'Substitution_Fixpoint's can be 'mappend'ed.
instance Monoid Substitution_Fixpoint where
	mempty  = empty
	mappend = substitution_fixpoint_union
instance Substitution Substitution_Fixpoint where
	substitution_codomain (Substitution_Fixpoint s) = Map.elems s
	substitution_domain   (Substitution_Fixpoint s) = Map.keys  s
	substitution_domain_remove vars (Substitution_Fixpoint s) =
		Substitution_Fixpoint $
		Map.filterWithKey (\v _ -> v `List.notElem` vars) s
	substitution_domain_restrict vars (Substitution_Fixpoint s) =
		Substitution_Fixpoint $
		Map.filterWithKey (\v _ -> Set.member v vars_kept) s
		where vars_kept = Set.fromList vars
		-- Map.filterWithKey (\v _ -> v `List.notElem` vars_removed) s
		-- where vars_removed = Map.keys s List.\\ vars
	substitution_lookup v sub@(Substitution_Fixpoint s) =
		case Map.lookup v s of
		 Nothing -> Nothing
		 Just ty | ty == Monotype_Var v -> Nothing
		         | otherwise -> Just $ sub `substitute` ty

substitution_fixpoint :: [(Monovar, Monotype)] -> Substitution_Fixpoint
substitution_fixpoint = Substitution_Fixpoint . Map.fromList

-- | WARNING: only valid with disjoint 'Substitution_Fixpoint's
substitution_fixpoint_union
 :: Substitution_Fixpoint
 -> Substitution_Fixpoint
 -> Substitution_Fixpoint
substitution_fixpoint_union
 (Substitution_Fixpoint x)
 (Substitution_Fixpoint y) =
	Substitution_Fixpoint $
	Map.unionWith err x y
	where err = error "substitution_fixpoint_union: the two substitutions are not disjoint"

-- ** Type 'Substitution_Instance'

data Substitution_Instance
 = forall s. Substitution s
 => Substitution_Instance
	 { substitution_instance_substitution :: s
	 , substitution_instance_codomain     :: s -> [Monotype]
	 , substitution_instance_domain       :: s -> [Monovar]
	 , substitution_instance_remove       :: [Monovar] -> s -> s
	 , substitution_instance_restrict     :: [Monovar] -> s -> s
	 , substitution_instance_lookup       :: Monovar -> s -> Maybe Monotype
	 }

substitution_instance :: Substitution s => s -> Substitution_Instance
substitution_instance s =
	Substitution_Instance s
	 substitution_codomain
	 substitution_domain
	 substitution_domain_remove
	 substitution_domain_restrict
	 substitution_lookup

instance Substitution Substitution_Instance where
	substitution_codomain           (Substitution_Instance s f _ _ _ _) = f s
	substitution_domain             (Substitution_Instance s _ f _ _ _) = f s
	substitution_domain_remove   vs (Substitution_Instance s _ _ f _ _) = substitution_instance (f vs s)
	substitution_domain_restrict vs (Substitution_Instance s _ _ _ f _) = substitution_instance (f vs s)
	substitution_lookup          v  (Substitution_Instance s _ _ _ _ f) = f v s

-- * Class 'Substitutable'

class Substitutable a where
	-- | /substitutable type variables/ (aka. /free type variables/).
	subvars :: a -> [Monovar]
	-- | /idempotent substitution application/.
	substitute :: Substitution s => s -> a -> a

instance Substitutable Monotype where
	subvars ty =
		case ty of
		 Monotype_Var v -> [v]
		 Monotype_Const _ -> []
		 Monotype_App t1 t2 -> subvars t1 `List.union` subvars t2
	sub `substitute` ty =
		case ty of
		 Monotype_Var v -> fromMaybe (Monotype_Var v) $ substitution_lookup v sub
		 Monotype_Const _ -> ty
		 Monotype_App t1 t2 -> Monotype_App (sub `substitute` t1) (sub `substitute` t2)
instance Substitutable a => Substitutable [a] where
	subvars = List.foldr (List.union . subvars) []
	substitute sub = ((sub `substitute`) <$>)
instance Substitutable a => Substitutable (Seq a) where
	subvars        = List.foldr (List.union . subvars) mempty
	substitute sub = ((sub `substitute`) <$>)
instance Substitutable a => Substitutable (Maybe a) where
	subvars = maybe [] subvars
	sub `substitute` m  = (sub `substitute`) <$> m
instance (Substitutable a, Substitutable b) => Substitutable (Either a b) where
	subvars = either subvars subvars
	sub `substitute` e = either (Left . (sub `substitute`)) (Right . (sub `substitute`)) e

-- | Return the next 'Monovar'
-- that is not in the 'subvars'
-- of the given 'Substitutable' value.
--
-- NOTE: return @0@ when 'subvars' is empty
subvar_next :: Substitutable a => a -> Monovar
subvar_next a =
	case subvars a of
	 [] -> 0
	 vs -> 1 + List.maximum vs