{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE UndecidableInstances  #-}

module Language.Rsc.Typecheck.Subst (

  -- * Substitutions
    RSubst
  , RSubstQ (..)
  , Subst
  , toList
  , fromList
  , toSubst

  -- * Free Type Variables
  , Free (..)

  -- * Type-class with operations
  , Substitutable
  , SubstitutableQ (..)

  ) where

import qualified Data.HashMap.Strict          as HM
import qualified Data.HashSet                 as S
import qualified Language.Fixpoint.Types      as F
import           Language.Rsc.Annotations
import           Language.Rsc.AST
import           Language.Rsc.Locations
import           Language.Rsc.Misc            (mapSnd)
import           Language.Rsc.Names
import           Language.Rsc.Typecheck.Types
import           Language.Rsc.Types

-- import           Debug.Trace

---------------------------------------------------------------------------
-- | Substitutions
---------------------------------------------------------------------------

-- | Type alias for Map from @TVar@ to @Type@. Hidden

data RSubstQ q r = Su (HM.HashMap TVar (RTypeQ q r))

type RSubst r = RSubstQ AK r
type Subst    = RSubst ()

toSubst :: RSubst r -> Subst
toSubst (Su m) = Su $ HM.map toType m

toList        :: RSubst r -> [(TVar, RType r)]
toList (Su m) =  HM.toList m

fromList      :: [(TVar, RTypeQ q r)] -> RSubstQ q r
fromList      = Su . HM.fromList


-- | Substitutions form a monoid; not commutative

instance (F.Reftable r, SubstitutableQ q r (RType r)) => Monoid (RSubstQ q r) where
  mempty                    = Su HM.empty
  mappend (Su m) θ'@(Su m') = Su $ (apply θ' <$> m) `HM.union` m'


class Free a where
  free  :: a -> S.HashSet TVar

instance Free (RType r) where
  free (TPrim _ _)          = S.empty
  free (TVar α _)           = S.singleton α
  free (TOr ts _)           = free ts
  free (TAnd ts)            = free $ snd <$> ts
  free (TRef n _)           = free n
  free (TObj _ es _)        = free es
  free (TClass t)           = free t
  free (TMod _)             = S.empty
  free (TAll α t)           = S.delete (btvToTV α) $ free t
  free (TFun xts t _)       = free $ [t] ++ (b_type <$> xts)
  free (TExp _)             = error "free should not be applied to TExp"

instance Free (TGen r) where
  free (Gen _ ts)           = free ts

instance Free (TypeMembers r) where
  free (TM ms sms cl ct s n)
                            = S.unions [free ms, free sms,
                                        free cl, free ct, free s, free n]

instance Free t => Free (F.SEnv t) where
  free                      = free . map snd . F.toListSEnv

instance Free (TypeMember r) where
  free (FI _ _ m t)           = free [m, t]
  free (MI _ _ mts)           = free $ map snd mts

instance Free a => Free [a] where
  free                      = S.unions . map free

instance Free (Fact r) where
  -- free (PhiVarTy (_,t))     = free t
  free (TypInst _ _ ts)     = free ts
  free (EltOverload _ t t') = free [t, t']
  free (VarAnn _ _ _ t )    = free t
  free (MemberAnn f)        = free f
  free (CtorAnn c)          = free c
  free (UserCast t)         = free t
  free (SigAnn _ _ c)       = free c
  free (ClassAnn _ t)       = free t --  foldr S.delete (free $ e ++ i) vs
  free (InterfaceAnn t)     = free t --  foldr S.delete (free $ e ++ i) vs
  -- TODO: TypeCast
  free _                    = S.empty

instance Free (TypeSig r) where
  free (TS _ n h)           = S.unions [free n, free h]

instance Free (TypeDecl r) where
  free (TD s _ m)           = S.unions [free s, free m]

instance Free (BTGen r) where
  free (BGen n ts)          = S.unions [free n, free ts]

instance Free (BTVar r) where
  free (BTV _ _ t)          = free t

instance Free a => Free (RelName, a) where
  free (_, a)               = free a

instance Free a => Free (Maybe a) where
  free Nothing              = S.empty
  free (Just a)             = free a

instance (Free a, Free b) => Free (a,b) where
  free (a,b)                = free a `S.union` free b

instance Free (QN l) where
  free _                    = S.empty

type Substitutable = SubstitutableQ AK

class SubstitutableQ q r a where
  apply                     :: (RSubstQ q r) -> a -> a

instance SubstitutableQ q r a => SubstitutableQ q r [a] where
  apply                     = map . apply

instance (SubstitutableQ q r a, SubstitutableQ q r b) => SubstitutableQ q r (a,b) where
  apply f (x,y)             = (apply f x, apply f y)

instance F.Reftable r => SubstitutableQ q r (RTypeQ q r) where
  apply θ t                 = appTy θ t

instance F.Reftable r => SubstitutableQ q r (BindQ q r) where
  apply θ (B z o t)         = B z o $ appTy θ t

instance SubstitutableQ q r t => SubstitutableQ q r (Located t) where
  apply θ (Loc s v)         = Loc s $ apply θ v

instance F.Reftable r => SubstitutableQ q r (FactQ q r) where
  -- apply θ (PhiVarTy (v,t))     = PhiVarTy . (v,) $ apply θ t
  apply θ (TypInst i ξ ts)     = TypInst i ξ     $ apply θ ts
  apply θ (EltOverload ξ t t') = EltOverload ξ (apply θ t) (apply θ t')
  apply θ (VarAnn x l a t)     = VarAnn x l a    $ apply θ t
  apply θ (MemberAnn f)        = MemberAnn       $ apply θ f
  apply θ (CtorAnn t)          = CtorAnn         $ apply θ t
  apply θ (UserCast t)         = UserCast        $ apply θ t
  apply θ (SigAnn x l t)       = SigAnn x l      $ apply θ t
  apply θ (ClassAnn l t)       = ClassAnn l      $ apply θ t
  apply θ (InterfaceAnn t)     = InterfaceAnn    $ apply θ t
  -- TODO: TypeCast
  apply _ a                    = a

instance F.Reftable r => SubstitutableQ q r (TypeMemberQ q r) where
  apply θ (FI s ms m t)       = FI s ms (apply θ m) (apply θ t)
  apply θ (MI s o mts)        = MI s o (mapSnd (apply θ) <$> mts)

instance SubstitutableQ q r a => SubstitutableQ q r (Maybe a) where
  apply θ (Just a)          = Just $ apply θ a
  apply _ Nothing           = Nothing

instance F.Reftable r => SubstitutableQ q r (TGenQ q r) where
  apply θ (Gen n ts)        = Gen n $ apply θ ts

instance F.Reftable r => SubstitutableQ q r (BTGenQ q r) where
  apply θ (BGen n ts)       = BGen n $ apply θ ts

instance F.Reftable r => SubstitutableQ q r (BTVarQ q r) where
  apply θ (BTV s l c)       = BTV s l $ apply θ c

instance F.Reftable r => SubstitutableQ q r (TypeMembersQ q r) where
  apply θ (TM ms sms cl ct s n)
                            = TM (apply θ ms) (apply θ sms)
                                 (apply θ cl) (apply θ ct)
                                 (apply θ s) (apply θ n)

instance SubstitutableQ q r a => SubstitutableQ q r (F.SEnv a) where
  apply                     = fmap . apply

instance SubstitutableQ q r (Id a) where
  apply _ i                 = i

instance F.Reftable r => SubstitutableQ q r (Annot (FactQ q r) z) where
  apply θ (Ann i z fs)      = Ann i z $ apply θ fs

instance SubstitutableQ q r F.Symbol  where
  apply _ s                 = s

instance SubstitutableQ q r (QN l) where
  apply _ s                 = s

instance SubstitutableQ q r (QP l) where
  apply _ s                 = s

instance F.Reftable r => SubstitutableQ q r (TypeSigQ q r) where
  apply θ (TS k n h)        = TS k n (apply θ h)

instance F.Reftable r => SubstitutableQ q r (TypeDeclQ q r) where
  apply θ (TD s p m)        = TD (apply θ s) p (apply θ m)

instance (F.Reftable r, SubstitutableQ q r a) => SubstitutableQ q r (Statement a) where
  apply θ s                 = fmap (apply θ) s

instance SubstitutableQ q r Assignability where
  apply _ s                 = s

instance SubstitutableQ q r Locality where
  apply _ s                 = s

instance (SubstitutableQ q r a, SubstitutableQ q r b, SubstitutableQ q r c) => SubstitutableQ q r (a,b,c) where
  apply θ (a,b,c)           = (apply θ a, apply θ b, apply θ c)

instance F.Reftable r => SubstitutableQ q r (FAnnQ q r) where
  apply θ (FA i s fs) = FA i s $ apply θ fs


---------------------------------------------------------------------------------
appTy :: F.Reftable r => RSubstQ q r -> RTypeQ q r -> RTypeQ q r
---------------------------------------------------------------------------------
appTy _        (TPrim p r)   = TPrim p r
appTy (Su m) t@(TVar α r)    = (HM.lookupDefault t α m) `strengthen` r
appTy θ        (TOr ts r)    = tOrR (apply θ ts) r
appTy θ        (TAnd ts)     = TAnd (mapSnd (apply θ) <$> ts)
appTy θ        (TRef n r)    = TRef (apply θ n) r
appTy θ        (TObj m es r) = TObj (appTy θ m) (apply θ es) r
appTy θ        (TClass t)    = TClass (apply θ t)
appTy _        (TMod n)      = TMod n
appTy (Su m)   (TAll α t)    = TAll α $ apply (Su $ HM.delete (btvToTV α) m) t
appTy θ        (TFun ts t r) = TFun (apply θ ts) (apply θ t) r
appTy _        (TExp _)      = error "appTy should not be applied to TExp"


---------------------------------------------------------------------------------
-- | Type equality (modulo α-renaming)
---------------------------------------------------------------------------------

instance (F.Reftable r, Eq q) => Eq (TGenQ q r) where
  Gen n ts    == Gen n' ts'    = n == n' && ts == ts'

instance (F.Reftable r, Eq q) => Eq (TypeMembersQ q r) where
  TM m sm c k s n == TM m' sm' c' k' s' n' =
    m == m' && sm == sm' && c == c' && k == k' && s == s' && n == n'

instance (F.Reftable r, Eq q) => Eq (TypeMemberQ q r) where
  FI n o a t == FI n' o' a' t' = n == n' && o == o' && a == a' && t == t'
  MI m k mt  == MI m' k' mt'   = m == m' && k == k' && mt == mt'
  _          == _              = False

instance (F.Reftable r, Eq q) => Eq (BindQ q r) where
  B x o t == B x' o' t' = (x, o, t) == (x', o', t')

instance (F.Reftable r, Eq q) => Eq (BTGenQ q r) where
  BGen n1 b1s == BGen n2 b2s = n1 == n2 && b1s == b2s

instance (F.Reftable r, Eq q) => Eq (BTVarQ q r) where
  BTV _ _ c1 == BTV _ _ c2 = c1 == c2

instance (F.Reftable r, Eq q) => Eq (RTypeQ q r) where
  TPrim p _    == TPrim p' _    = p == p'
  TVar α _     == TVar α' _     = α == α'
  TOr ts _     == TOr ts' _     = ts == ts'
  TAnd ts      == TAnd ts'      = ts == ts'
  TRef g _     == TRef g' _     = g == g'
  TObj m ms _  == TObj m' ms' _ = m == m' && ms == ms'
  TClass g     == TClass g'     = g == g'
  TMod n       == TMod n'       = n == n'
  TAll v@(BTV _ b _) t == TAll v'@(BTV _ b' _) t'
                                = b == b' && t == appTy θ t'
    where
      θ = fromList [(btvToTV v', tVar $ btvToTV v)]
  TFun bs o _ == TFun bs' o' _ = bs == bs' && o == o'
  _           == _             = False

