{-# LANGUAGE CPP #-}
module Agda.TypeChecking.Substitute where

import Control.Monad.Identity
import Control.Monad.Reader
import Data.Generics
import Data.List hiding (sort)
import Data.Map (Map)
import qualified Data.Map as Map

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

import Agda.TypeChecking.Monad.Base
import Agda.TypeChecking.Free

import Agda.Utils.Monad
import Agda.Utils.Size
import Agda.Utils.Permutation

#include "../undefined.h"
import Agda.Utils.Impossible

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

instance Apply Term where
    apply m [] = m
    apply m args@(Arg _ v:args0) =
	case m of
	    Var i args'   -> Var i (args' ++ args)
	    Def c args'   -> Def c (args' ++ args)
	    Con c args'   -> Con c (args' ++ args)
	    Lam _ u	  -> absApp u v `apply` args0
	    MetaV x args' -> MetaV x (args' ++ args)
	    Lit l	  -> __IMPOSSIBLE__
	    Pi _ _	  -> __IMPOSSIBLE__
	    Fun _ _	  -> __IMPOSSIBLE__
	    Sort _	  -> __IMPOSSIBLE__

instance Apply Type where
  apply = piApply

instance Apply Sort where
  apply s [] = s
  apply s _  = __IMPOSSIBLE__

instance Apply Telescope where
  apply tel		  []	   = tel
  apply EmptyTel	  _	   = __IMPOSSIBLE__
  apply (ExtendTel _ tel) (t : ts) = absApp tel (unArg t) `apply` ts

instance Apply Definition where
    apply (Defn x t df m d) args = Defn x (piApply t args) df m (apply d args)

instance Apply Defn where
  apply d args = case d of
    Axiom{} -> d
    Function{ funClauses = cs, funInv = inv } ->
      d { funClauses = apply cs args, funInv = apply inv args }
    Datatype{ dataPars = np, dataClause = cl } ->
      d { dataPars = np - size args, dataClause = apply cl args }
    Record{ recPars = np, recClause = cl, recTel = tel } ->
      d { recPars = np - size args, recClause = apply cl args, recTel = apply tel args }
    Constructor{ conPars = np } ->
      d { conPars = np - size args }
    Primitive{ primClauses = cs } ->
      d { primClauses = apply cs args }

instance Apply PrimFun where
    apply (PrimFun x ar def) args   = PrimFun x (ar - size args) $ \vs -> def (args ++ vs)

instance Apply Clause where
    apply (Clause r tel perm ps b) args =
      Clause r (apply tel args) (apply perm args)
             (drop (size args) ps) (apply b args)

instance Apply FunctionInverse where
  apply NotInjective  args = NotInjective
  apply (Inverse inv) args = Inverse $ apply inv args

instance Apply ClauseBody where
    apply  b		   []		  = b
    apply (Bind (Abs _ b)) (Arg _ v:args) = subst v b `apply` args
    apply (NoBind b)	   (_:args)	  = b `apply` args
    apply (Body _)	   (_:_)	  = __IMPOSSIBLE__
    apply  NoBody	    _		  = NoBody

instance Apply DisplayTerm where
  apply (DTerm v)	   args = DTerm $ apply v args
  apply (DWithApp v args') args = DWithApp v $ args' ++ args

instance Apply t => Apply [t] where
    apply ts args = map (`apply` args) ts

instance Apply t => Apply (Blocked t) where
    apply b args = fmap (`apply` args) b

instance Apply t => Apply (Maybe t) where
  apply x args = fmap (`apply` args) x

instance Apply v => Apply (Map k v) where
  apply x args = fmap (`apply` args) x

instance (Apply a, Apply b) => Apply (a,b) where
    apply (x,y) args = (apply x args, apply y args)

instance (Apply a, Apply b, Apply c) => Apply (a,b,c) where
    apply (x,y,z) args = (apply x args, apply y args, apply z args)

instance Apply Permutation where
  -- The permutation must start with [0..m - 1]
  apply (Perm n xs) args = Perm (n - m) $ map (flip (-) m) $ genericDrop m xs
    where
      m = size args

instance Abstract Permutation where
  abstract tel (Perm n xs) = Perm (n + m) $ [0..m - 1] ++ map (+ m) xs
    where
      m = size tel

-- | The type must contain the right number of pis without have to perform any
-- reduction.
piApply :: Type -> Args -> Type
piApply t []				= t
piApply (El _ (Pi  _ b)) (Arg _ v:args) = absApp b v `piApply` args
piApply (El _ (Fun _ b)) (_:args)	= b `piApply` args
piApply _ _				= __IMPOSSIBLE__

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

instance Abstract Term where
    abstract = teleLam

instance Abstract Type where
    abstract = telePi_

instance Abstract Sort where
    abstract EmptyTel s = s
    abstract _	      s = __IMPOSSIBLE__

instance Abstract Telescope where
  abstract  EmptyTel	        tel = tel
  abstract (ExtendTel arg tel') tel = ExtendTel arg $ fmap (`abstract` tel) tel'

instance Abstract Definition where
    abstract tel (Defn x t df m d) = Defn x (abstract tel t) df m (abstract tel d)

instance Abstract Defn where
  abstract tel d = case d of
    Axiom{} -> d
    Function{ funClauses = cs, funInv = inv } ->
      d { funClauses = abstract tel cs, funInv = abstract tel inv }
    Datatype{ dataPars = np, dataClause = cl } ->
      d { dataPars = np + size tel, dataClause = abstract tel cl }
    Record{ recPars = np, recClause = cl, recTel = tel' } ->
      d { recPars = np + size tel, recClause = abstract tel cl, recTel = abstract tel tel' }
    Constructor{ conPars = np } ->
      d { conPars = np + size tel }
    Primitive{ primClauses = cs } ->
      d { primClauses = abstract tel cs }

instance Abstract PrimFun where
    abstract tel (PrimFun x ar def) = PrimFun x (ar + n) $ \ts -> def $ genericDrop n ts
	where n = size tel

instance Abstract Clause where
  abstract tel (Clause r tel' perm ps b) =
    Clause r (abstract tel tel') (abstract tel perm)
           (telVars tel ++ ps) (abstract tel b)

telVars EmptyTel		    = []
telVars (ExtendTel arg (Abs x tel)) = fmap (const $ VarP x) arg : telVars tel

instance Abstract FunctionInverse where
  abstract tel NotInjective  = NotInjective
  abstract tel (Inverse inv) = Inverse $ abstract tel inv

instance Abstract ClauseBody where
    abstract EmptyTel		 b = b
    abstract (ExtendTel _ tel)	 b = Bind $ fmap (`abstract` b) tel

instance Abstract t => Abstract [t] where
    abstract tel = map (abstract tel)

instance Abstract t => Abstract (Maybe t) where
  abstract tel x = fmap (abstract tel) x

instance Abstract v => Abstract (Map k v) where
  abstract tel m = fmap (abstract tel) m

abstractArgs :: Abstract a => Args -> a -> a
abstractArgs args x = abstract tel x
    where
	tel   = foldr (\(Arg h x) -> ExtendTel (Arg h $ sort Prop) . Abs x) EmptyTel
	      $ zipWith (fmap . const) names args
	names = cycle $ map (:[]) ['a'..'z']

-- | Substitute a term for the nth free variable.
--
class Subst t where
    substs     :: [Term] -> t -> t
    substUnder :: Nat -> Term -> t -> t

idSub :: Telescope -> [Term]
idSub tel = [ Var i [] | i <- [0..size tel - 1] ]

subst :: Subst t => Term -> t -> t
subst u t = substUnder 0 u t

instance Subst Term where
    substs us t =
	case t of
	    Var i vs   -> (us !!! i) `apply` substs us vs
	    Lam h m    -> Lam h $ substs us m
	    Def c vs   -> Def c $ substs us vs
	    Con c vs   -> Con c $ substs us vs
	    MetaV x vs -> MetaV x $ substs us vs
	    Lit l      -> Lit l
	    Pi a b     -> uncurry Pi $ substs us (a,b)
	    Fun a b    -> uncurry Fun $ substs us (a,b)
	    Sort s     -> Sort s
        where
            []     !!! n = error "unbound variable"
            (x:xs) !!! 0 = x
            (_:xs) !!! n = xs !!! (n - 1)
    substUnder n u t =
	case t of
	    Var i vs
	      | i == n	  -> raise n u `apply` substUnder n u vs
	      | i < n	  -> Var i $ substUnder n u vs
	      | otherwise -> Var (i - 1) $ substUnder n u vs
	    Lam h m    -> Lam h $ substUnder n u m
	    Def c vs   -> Def c $ substUnder n u vs
	    Con c vs   -> Con c $ substUnder n u vs
	    MetaV x vs -> MetaV x $ substUnder n u vs
	    Lit l      -> Lit l
	    Pi a b     -> uncurry Pi $ substUnder n u (a,b)
	    Fun a b    -> uncurry Fun $ substUnder n u (a,b)
	    Sort s     -> Sort s

instance Subst Type where
    substs us (El s t) = El s $ substs us t
    substUnder n u (El s t) = El s $ substUnder n u t

instance Subst Pattern where
  substs us p = case p of
    VarP s    -> VarP s
    LitP l    -> LitP l
    ConP c ps -> ConP c $ substs us ps
    DotP t    -> DotP $ substs us t
  substUnder n u p = case p of
    VarP s    -> VarP s
    LitP l    -> LitP l
    ConP c ps -> ConP c $ substUnder n u ps
    DotP t    -> DotP $ substUnder n u t

instance Subst t => Subst (Blocked t) where
    substs us b	     = fmap (substs us) b
    substUnder n u b = fmap (substUnder n u) b

instance Subst DisplayTerm where
  substs us	 (DTerm v)	  = DTerm $ substs us v
  substs us	 (DWithApp vs ws) = uncurry DWithApp $ substs us (vs, ws)
  substUnder n u (DTerm v)	  = DTerm $ substUnder n u v
  substUnder n u (DWithApp vs ws) = uncurry DWithApp $ substUnder n u (vs, ws)

instance Subst Telescope where
  substs us  EmptyTel		   = EmptyTel
  substs us (ExtendTel t tel)	   = uncurry ExtendTel $ substs us (t, tel)
  substUnder n u  EmptyTel	   = EmptyTel
  substUnder n u (ExtendTel t tel) = uncurry ExtendTel $ substUnder n u (t, tel)

instance (Data a, Subst a) => Subst (Abs a) where
    substs us	   (Abs x t) = Abs x $ substs (Var 0 [] : raise 1 us) t
    substUnder n u (Abs x t) = Abs x $ substUnder (n + 1) u t

instance Subst a => Subst (Arg a) where
    substs us	   = fmap (substs us)
    substUnder n u = fmap (substUnder n u)

instance Subst a => Subst (Maybe a) where
  substs us	 = fmap (substs us)
  substUnder n u = fmap (substUnder n u)

instance Subst a => Subst [a] where
    substs us	   = map (substs us)
    substUnder n u = map (substUnder n u)

instance (Subst a, Subst b) => Subst (a,b) where
    substs us (x,y)	 = (substs us x, substs us y)
    substUnder n u (x,y) = (substUnder n u x, substUnder n u y)

instance Subst ClauseBody where
    substs us (Body t)	      = Body $ substs us t
    substs us (Bind b)	      = Bind $ substs us b
    substs us (NoBind b)      = NoBind $ substs us b
    substs _   NoBody	      = NoBody
    substUnder n u (Body t)   = Body $ substUnder n u t
    substUnder n u (Bind b)   = Bind $ substUnder n u b
    substUnder n u (NoBind b) = NoBind $ substUnder n u b
    substUnder _ _   NoBody   = NoBody

-- | Instantiate an abstraction
absApp :: Subst t => Abs t -> Term -> t
absApp (Abs _ v) u = subst u v

-- | Add @k@ to index of each open variable in @x@.
class Raise t where
    raiseFrom :: Nat -> Nat -> t -> t

instance Raise Term where
    raiseFrom m k v =
	case v of
	    Var i vs
		| i < m	    -> Var i $ rf vs
		| otherwise -> Var (i + k) $ rf vs
	    Lam h m	    -> Lam h $ rf m
	    Def c vs	    -> Def c $ rf vs
	    Con c vs	    -> Con c $ rf vs
	    MetaV x vs	    -> MetaV x $ rf vs
	    Lit l	    -> Lit l
	    Pi a b	    -> uncurry Pi $ rf (a,b)
	    Fun a b	    -> uncurry Fun $ rf (a,b)
	    Sort s	    -> Sort s
	where
	    rf x = raiseFrom m k x

instance Raise Type where
    raiseFrom m k (El s t) = El s $ raiseFrom m k t

instance Raise Telescope where
    raiseFrom m k EmptyTel	    = EmptyTel
    raiseFrom m k (ExtendTel a tel) = uncurry ExtendTel $ raiseFrom m k (a, tel)

instance Raise DisplayForm where
  raiseFrom m k (Display n ps v) = Display n (raiseFrom (m + 1) k ps)
					     (raiseFrom (m + n) k v)

instance Raise DisplayTerm where
  raiseFrom m k (DWithApp xs ys) = uncurry DWithApp $ raiseFrom m k (xs, ys)
  raiseFrom m k (DTerm v)	 = DTerm $ raiseFrom m k v

instance Raise t => Raise (Abs t) where
    raiseFrom m k = fmap (raiseFrom (m + 1) k)

instance Raise t => Raise (Arg t) where
    raiseFrom m k = fmap (raiseFrom m k)

instance Raise t => Raise (Blocked t) where
    raiseFrom m k = fmap (raiseFrom m k)

instance Raise t => Raise [t] where
    raiseFrom m k = fmap (raiseFrom m k)

instance Raise t => Raise (Maybe t) where
    raiseFrom m k = fmap (raiseFrom m k)

instance Raise v => Raise (Map k v) where
    raiseFrom m k = fmap (raiseFrom m k)

instance (Raise a, Raise b) => Raise (a,b) where
    raiseFrom m k (x,y) = (raiseFrom m k x, raiseFrom m k y)

raise :: Raise t => Nat -> t -> t
raise = raiseFrom 0

data TelView = TelV Telescope Type

telView :: Type -> TelView
telView t = case unEl t of
  Pi a (Abs x b)  -> absV a x $ telView b
  Fun a b	  -> absV a "_" $ telView (raise 1 b)
  _		  -> TelV EmptyTel t
  where
    absV a x (TelV tel t) = TelV (ExtendTel a (Abs x tel)) t

telePi :: Telescope -> Type -> Type
telePi  EmptyTel	 t = t
telePi (ExtendTel u tel) t = el $ fn u b
  where
    el = El (sLub s1 s2)
    b = fmap (flip telePi t) tel
    s1 = getSort $ unArg u
    s2 = getSort $ absBody b

    fn a b
      | 0 `freeIn` absBody b = Pi a b
      | otherwise	     = Fun a $ absApp b __IMPOSSIBLE__

-- | Everything will be a pi.
telePi_ :: Telescope -> Type -> Type
telePi_  EmptyTel	 t = t
telePi_ (ExtendTel u tel) t = el $ Pi u b
  where
    el = El (sLub s1 s2)
    b = fmap (flip telePi_ t) tel
    s1 = getSort $ unArg u
    s2 = getSort $ absBody b