{-# OPTIONS_GHC -Wunused-imports #-}

module Agda.TypeChecking.Substitute.Class where

import Control.Arrow ((***), second)
import Control.DeepSeq
import GHC.Generics

import Agda.Syntax.Common
import Agda.Syntax.Internal

import Agda.TypeChecking.Free
import Agda.TypeChecking.Substitute.DeBruijn

import Agda.Utils.Impossible

---------------------------------------------------------------------------
-- * Application
---------------------------------------------------------------------------

-- | Apply something to a bunch of arguments.
--   Preserves blocking tags (application can never resolve blocking).
class Apply t where
  apply  :: t -> Args -> t
  applyE :: t -> Elims -> t

  apply t args = applyE t $ map Apply args
  -- Andreas, 2018-06-18, issue #3136
  -- This default instance should be removed to get more precise
  -- crash locations (raise the IMPOSSIBLE in a more specific place).
  -- applyE t es  = apply  t $ fromMaybe __IMPOSSIBLE__ $ allApplyElims es
    -- precondition: all @es@ are @Apply@s

-- | Apply to some default arguments.
applys :: Apply t => t -> [Term] -> t
applys t vs = apply t $ map defaultArg vs

-- | Apply to a single default argument.
apply1 :: Apply t => t -> Term -> t
apply1 t u = applys t [ u ]

-- | Apply to two default arguments.
apply2 :: Apply t => t -> Term -> Term -> t
apply2 t u v = applys t [u, v]

---------------------------------------------------------------------------
-- * Abstraction
---------------------------------------------------------------------------

-- | @(abstract args v) `apply` args --> v[args]@.
class Abstract t where
  abstract :: Telescope -> t -> t

---------------------------------------------------------------------------
-- * Substitution and shifting\/weakening\/strengthening
---------------------------------------------------------------------------

-- | Apply a substitution.

-- For terms:
--
--  Γ ⊢ ρ : Δ
--  Δ ⊢ t : A
-- -----------
-- Γ ⊢ tρ : Aρ

class DeBruijn (SubstArg a) => Subst a where
  type SubstArg a
  applySubst :: Substitution' (SubstArg a) -> a -> a

  default applySubst :: (a ~ f b, Functor f, Subst b, SubstArg a ~ SubstArg b) => Substitution' (SubstArg a) -> a -> a
  applySubst rho = fmap (applySubst rho)

-- | Simple constraint alias for a `Subst` instance `a` with arg type `t`.
type SubstWith t a = (Subst a, SubstArg a ~ t)

-- | `Subst` instance whose agument type is itself
type EndoSubst a = SubstWith a a

-- | `Subst` instance whose argument type is `Term`
type TermSubst a = SubstWith Term a

-- | Raise de Bruijn index, i.e. weakening
raise :: Subst a => Nat -> a -> a
raise = raiseFrom 0

raiseFrom :: Subst a => Nat -> Nat -> a -> a
raiseFrom n k = applySubst (raiseFromS n k)

-- | Replace de Bruijn index i by a 'Term' in something.
subst :: Subst a => Int -> SubstArg a -> a -> a
subst i u = applySubst $ singletonS i u

strengthen :: Subst a => Impossible -> a -> a
strengthen err = applySubst (strengthenS err 1)

-- | Replace what is now de Bruijn index 0, but go under n binders.
--   @substUnder n u == subst n (raise n u)@.
substUnder :: Subst a => Nat -> SubstArg a -> a -> a
substUnder n u = applySubst (liftS n (singletonS 0 u))

-- | Checks whether the variable bound by the abstraction is actually
-- used, and, if /not/, returns the term within, 'strengthen'ed to live in
-- the context /outside/ the abstraction.
-- See also 'Agda.TypeChecking.Free.isBinderUsed'.
isNoAbs :: (Free a, Subst a) => Abs a -> Maybe a
isNoAbs (NoAbs _ b) = Just b
isNoAbs (Abs _ b)
  | not (0 `freeIn` b) = Just (strengthen __IMPOSSIBLE__ b)
  | otherwise          = Nothing

-- ** Identity instances

instance Subst QName where
  type SubstArg QName = Term
  applySubst _ q = q

-- | Wrapper for types that do not contain variables (so applying a substitution is the identity).
--   Useful if you have a structure of types that support substitution mixed with types that don't
--   and need to apply a substitution to the full structure.
newtype NoSubst t a = NoSubst { unNoSubst :: a }
  deriving (Generic, NFData, Functor)

instance DeBruijn t => Subst (NoSubst t a) where
  type SubstArg (NoSubst t a) = t
  applySubst _ x = x

---------------------------------------------------------------------------
-- * Explicit substitutions
---------------------------------------------------------------------------

-- See Syntax.Internal for the definition.

idS :: Substitution' a
idS = IdS

wkS :: Int -> Substitution' a -> Substitution' a
wkS 0 rho        = rho
wkS n (Wk m rho) = Wk (n + m) rho
wkS n (EmptyS err) = EmptyS err
wkS n rho        = Wk n rho

raiseS :: Int -> Substitution' a
raiseS n = wkS n idS

{-# INLINABLE consS #-}
consS :: DeBruijn a => a -> Substitution' a -> Substitution' a
consS t (Wk m rho)
  | Just n <- deBruijnView t,
    n + 1 == m = wkS (m - 1) (liftS 1 rho)
consS u rho = seq u (u :# rho)
{-# SPECIALIZE consS :: Term  -> Substitution' Term  -> Substitution' Term #-}
{-# SPECIALIZE consS :: Level -> Substitution' Level -> Substitution' Level #-}

{-# INLINABLE singletonS #-}
-- | To replace index @n@ by term @u@, do @applySubst (singletonS n u)@.
--   @
--               Γ, Δ ⊢ u : A
--    ---------------------------------
--    Γ, Δ ⊢ singletonS |Δ| u : Γ, A, Δ
--   @
singletonS :: DeBruijn a => Int -> a -> Substitution' a
singletonS n u = map deBruijnVar [0..n-1] ++# consS u (raiseS n)
  -- ALT: foldl (\ s i -> deBruijnVar i `consS` s) (consS u $ raiseS n) $ downFrom n
{-# SPECIALIZE singletonS :: Int -> Term -> Substitution' Term #-}
{-# SPECIALIZE singletonS :: Int -> Level -> Substitution' Level #-}

-- | Single substitution without disturbing any deBruijn indices.
--   @
--             Γ, A, Δ ⊢ u : A
--    ---------------------------------
--    Γ, A, Δ ⊢ inplace |Δ| u : Γ, A, Δ
--   @
inplaceS :: EndoSubst a => Int -> a -> Substitution' a
inplaceS k u = singletonS k u `composeS` liftS (k + 1) (raiseS 1)

-- | Lift a substitution under k binders.
liftS :: Int -> Substitution' a -> Substitution' a
liftS 0 rho          = rho
liftS k IdS          = IdS
liftS k (Lift n rho) = Lift (n + k) rho
liftS k rho          = Lift k rho

-- | @
--         Γ ⊢ ρ : Δ, Ψ
--      -------------------
--      Γ ⊢ dropS |Ψ| ρ : Δ
--   @
dropS :: Int -> Substitution' a -> Substitution' a
dropS 0 rho                = rho
dropS n IdS                = raiseS n
dropS n (Wk m rho)         = wkS m (dropS n rho)
dropS n (u :# rho)         = dropS (n - 1) rho
dropS n (EmptyS err)       = throwImpossible err
dropS n (Lift m rho)
  -- dropS n (Lift m rho) =
  --   wkS 1 $ dropS (n - 1) $ liftS (m - 1) rho
  | n > m     = wkS m $ dropS (n - m) rho
  | otherwise = wkS n $ liftS (m - n) rho
dropS n (Strengthen err m rho)
  | n < m     = Strengthen err (m - n) rho
  | otherwise = dropS (n - m) rho

{-# INLINABLE composeS #-}
-- | @applySubst (ρ `composeS` σ) v == applySubst ρ (applySubst σ v)@
composeS :: EndoSubst a => Substitution' a -> Substitution' a -> Substitution' a
composeS rho IdS = rho
composeS IdS sgm = sgm
composeS rho (EmptyS err) = EmptyS err
composeS rho (Wk n sgm) = composeS (dropS n rho) sgm
composeS rho (u :# sgm) = applySubst rho u :# composeS rho sgm
composeS rho (Strengthen err n sgm) = Strengthen err n (composeS rho sgm)
composeS rho (Lift 0 sgm) = __IMPOSSIBLE__
composeS (u :# rho) (Lift n sgm) = u :# composeS rho (liftS (n - 1) sgm)
composeS rho (Lift n sgm) = lookupS rho 0 :# composeS rho (wkS 1 (liftS (n - 1) sgm))

-- If Γ ⊢ ρ : Δ, Θ then splitS |Θ| ρ = (σ, δ), with
--   Γ ⊢ σ : Δ
--   Γ ⊢ δ : Θσ
splitS :: Int -> Substitution' a -> (Substitution' a, Substitution' a)
splitS 0 rho                    = (rho, EmptyS impossible)
splitS n (u :# rho)             = second (u :#) $ splitS (n - 1) rho
splitS n (Lift 0 _)             = __IMPOSSIBLE__
splitS n (Wk m rho)             = wkS m *** wkS m $ splitS n rho
splitS n IdS                    = ( raiseS n
                                  , liftS n $ EmptyS impossible
                                  )
splitS n (Lift m rho)           = wkS 1 *** liftS 1 $
                                  splitS (n - 1) (liftS (m - 1) rho)
splitS n (EmptyS err)           = __IMPOSSIBLE__
splitS n (Strengthen err m rho)
  -- splitS n (Strengthen err 1 rho) =
  --   second (Strengthen err) $ splitS (n - 1) rho
  | n > m     = second (Strengthen err m) $
                splitS (n - m) rho
  | otherwise = second (Strengthen err n) $
                splitS 0 (Strengthen err (m - n) rho)

infixr 4 ++#

(++#) :: DeBruijn a => [a] -> Substitution' a -> Substitution' a
us ++# rho = foldr consS rho us

-- | @
--      Γ ⊢ ρ : Δ  Γ ⊢ reverse vs : Θ
--      ----------------------------- (treating Nothing as having any type)
--        Γ ⊢ prependS vs ρ : Δ, Θ
--   @
prependS :: DeBruijn a => Impossible -> [Maybe a] -> Substitution' a -> Substitution' a
prependS err us rho = go 0 us
  where
  -- The function strengthenS' is not used here, to avoid replacing
  -- the "error message" of a potential outermost Strengthen
  -- constructor in rho.
  str 0 = id
  str n = Strengthen err n

  go !n (Just u  : us) = str n (consS u (go 0 us))
  go  n (Nothing : us) = go (1 + n) us
  go  n []             = str n rho

-- | @
--        Γ ⊢ reverse vs : Δ
--      -----------------------------
--        Γ ⊢ parallelS vs ρ : Γ, Δ
--   @
--
--   Note the @Γ@ in @Γ, Δ@.
parallelS :: DeBruijn a => [a] -> Substitution' a
parallelS us = us ++# idS

-- | Γ ⊢ (strengthenS ⊥ |Δ|) : Γ,Δ
strengthenS :: Impossible -> Int -> Substitution' a
strengthenS err n = case compare n 0 of
  LT -> __IMPOSSIBLE__
  EQ -> idS
  GT -> Strengthen err n idS

-- | A \"smart\" variant of 'Strengthen'. If 'strengthenS' is applied
-- to a substitution with an outermost 'Strengthen' constructor, then
-- the \"error message\" of that constructor is discarded in favour of
-- the 'Impossible' argument of this function.
strengthenS' :: Impossible -> Int -> Substitution' a -> Substitution' a
strengthenS' err m rho = case compare m 0 of
  LT -> __IMPOSSIBLE__
  EQ -> rho
  GT -> case rho of
    Strengthen _ n rho -> Strengthen err (m + n) rho
    _                  -> Strengthen err m       rho

{-# INLINABLE lookupS #-}
lookupS :: EndoSubst a => Substitution' a -> Nat -> a
lookupS rho i = case rho of
  IdS                    -> deBruijnVar i
  Wk n IdS               -> let j = i + n in
                            if  j < 0 then __IMPOSSIBLE__ else deBruijnVar j
  Wk n rho               -> applySubst (raiseS n) (lookupS rho i)
  u :# rho   | i == 0    -> u
             | i < 0     -> __IMPOSSIBLE__
             | otherwise -> lookupS rho (i - 1)
  Strengthen err n rho
             | i < 0     -> __IMPOSSIBLE__
             | i < n     -> throwImpossible err
             | otherwise -> lookupS rho (i - n)
  Lift n rho | i < n     -> deBruijnVar i
             | otherwise -> raise n $ lookupS rho (i - n)
  EmptyS err             -> throwImpossible err


-- | lookupS (listS [(x0,t0)..(xn,tn)]) xi = ti, assuming x0 < .. < xn.


listS :: EndoSubst a => [(Int,a)] -> Substitution' a
listS ((i,t):ts) = singletonS i t `composeS` listS ts
listS []         = IdS

-- | @Γ, Ξ, Δ ⊢ raiseFromS |Δ| |Ξ| : Γ, Δ@
raiseFromS :: Nat -> Nat -> Substitution' a
raiseFromS n k = liftS n $ raiseS k


---------------------------------------------------------------------------
-- * Functions on abstractions
--   and things we couldn't do before we could define 'absBody'
---------------------------------------------------------------------------

-- | Instantiate an abstraction. Strict in the term.
absApp :: Subst a => Abs a -> SubstArg a -> a
absApp (Abs   _ v) u = subst 0 u v
absApp (NoAbs _ v) _ = v

-- | Instantiate an abstraction. Lazy in the term, which allow it to be
--   __IMPOSSIBLE__ in the case where the variable shouldn't be used but we
--   cannot use 'noabsApp'. Used in Apply.
lazyAbsApp :: Subst a => Abs a -> SubstArg a -> a
lazyAbsApp (Abs   _ v) u = applySubst (u :# IdS) v  -- Note: do not use consS here!
lazyAbsApp (NoAbs _ v) _ = v

-- | Instantiate an abstraction that doesn't use its argument.
noabsApp :: Subst a => Impossible -> Abs a -> a
noabsApp err (Abs   _ v) = strengthen err v
noabsApp _   (NoAbs _ v) = v

absBody :: Subst a => Abs a -> a
absBody (Abs   _ v) = v
absBody (NoAbs _ v) = raise 1 v

mkAbs :: (Subst a, Free a) => ArgName -> a -> Abs a
mkAbs x v | 0 `freeIn` v = Abs x v
          | otherwise    = NoAbs x (raise (-1) v)

reAbs :: (Subst a, Free a) => Abs a -> Abs a
reAbs (NoAbs x v) = NoAbs x v
reAbs (Abs x v)   = mkAbs x v
