module Quox.Syntax.Subst

import Quox.Thin
import Quox.Loc

import Data.DPair
import Data.List
import Data.SnocVect
import Derive.Prelude

%default total
%language ElabReflection


public export
Subst : (Nat -> Type) -> Nat -> Nat -> Type
Subst env from to = SnocVect from (Lazy (Thinned env to))

public export
Subst2 : (Nat -> Nat -> Type) -> Nat -> Nat -> Nat -> Type
Subst2 env d from to = SnocVect from (Lazy (Thinned2 env d to))


public export
get : Subst env f t -> Fin f -> Thinned env t
get (sx :< x) FZ     = x
get (sx :< x) (FS i) = get sx i


public export
interface FromVar (0 term : Nat -> Type) where
  var : Loc -> term 1

public export
0 FromVar2 : (Nat -> Nat -> Type) -> Type
FromVar2 t = FromVar (t 0)

public export
varT : FromVar term => Fin n -> Loc -> Thinned term n
varT i loc = Th (one' i) (var loc)

public export
varT2 : FromVar2 term => Fin n -> Loc -> Thinned2 term d n
varT2 i loc = Th2 zero (one' i) (var loc)

infixl 8 //
namespace CanSubstSelf
  public export
  interface FromVar term => CanSubstSelf term where
    (//) : {f : Nat} -> Thinned term f -> Subst term f t -> Thinned term t

namespace CanSubstSelf2
  public export
  interface FromVar2 term => CanSubstSelf2 term where
    (//) : {f : Nat} -> Thinned2 term d f ->
           Subst2 term d f t -> Thinned2 term d t

public export
(.) : {mid : Nat} -> CanSubstSelf f =>
      Subst f from mid -> Subst f mid to -> Subst f from to
th . ph = map (\(Delay x) => x // ph) th

infixr 9 .%

public export
(.%) : {mid : Nat} -> CanSubstSelf2 f =>
       Subst2 f d from mid -> Subst2 f d mid to -> Subst2 f d from to
th .% ph = map (\(Delay x) => x // ph) th


public export
tabulate : (n : Nat) -> SnocVect n (Fin n)
tabulate n = go n id where
  go : (n : Nat) -> (Fin n -> Fin n') -> SnocVect n (Fin n')
  go 0     f = [<]
  go (S n) f = go n (f . FS) :< f FZ

public export
id : FromVar term => {n : Nat} -> (under : Nat) -> Loc ->
     Subst term n (n + under)
id under loc =
  map (\i => delay $ varT (weakenN under i) loc) (tabulate n)

public export
id2 : FromVar2 term => {n : Nat} -> Loc -> Subst2 term d n n
id2 loc = map (\i => delay $ varT2 i loc) $ tabulate n

export
select : {n, mask : Nat} -> (0 ope : OPE m n mask) ->
         SnocVect n a -> SnocVect m a
select ope sx with %syntactic (view ope)
  select _ [<]       | StopV       = [<]
  select _ (sx :< x) | DropV _ ope = select ope sx
  select _ (sx :< x) | KeepV _ ope = select ope sx :< x

export
opeToFins : {n, mask : Nat} ->
            (0 ope : OPE m n mask) -> SnocVect m (Fin n)
opeToFins ope = select ope $ tabulate n

export
shift : FromVar term => {from : Nat} ->
        (n : Nat) -> Loc -> Subst term from (n + from)
shift n loc = map (\i => delay $ varT (shift n i) loc) $ tabulate from

public export
pushN : CanSubstSelf term => {to : Nat} -> (by : Nat) ->
        Subst term from to -> Loc -> Subst term (by + from) (by + to)
pushN by th loc =
  rewrite plusCommutative by from in
  (th . shift by loc) ++ id to loc

public export %inline
push : CanSubstSelf f => {to : Nat} ->
       Subst f from to -> Loc -> Subst f (S from) (S to)
push = pushN 1


public export %inline
one : Thinned f n -> Subst f 1 n
one x = [< x]


||| whether two substitutions with the same codomain have the same domain
export
cmpShape : SnocVect m a -> SnocVect n a -> Either Ordering (m = n)
cmpShape [<]       [<]       = Right Refl
cmpShape [<]       (sx :< _) = Left LT
cmpShape (sx :< _) [<]       = Left GT
cmpShape (sx :< _) (sy :< _) = cong S <$> cmpShape sx sy


public export
record WithSubst tm env n where
  constructor Sub
  term  : tm from
  subst : Subst env from n

{-
export
(forall n. Eq (env n), forall n. Eq (tm n)) =>
Eq (WithSubst tm env n) where
  Sub t1 s1 == Sub t2 s2 =
    case cmpShape s1 s2 of
      Left  _    => False
      Right Refl =>
        t1 == t2 && concat @{All} (zipWith ((==) `on` force) s1 s2)

export
(forall n. Ord (env n), forall n. Ord (tm n)) =>
Ord (WithSubst tm env n) where
  Sub t1 s1 `compare` Sub t2 s2 =
    case cmpShape s1 s2 of
      Left  o    => o
      Right Refl =>
        compare t1 t2 <+> concat (zipWith (compare `on` force) s1 s2)

export %hint
ShowWithSubst : {n : Nat} ->
                (forall n. Show (env n), forall n. Show (tm n)) =>
                Show (WithSubst tm env n)
ShowWithSubst = deriveShow where
  Show (Lazy (Thinned env n)) where showPrec d = showPrec d . force
-}


public export
record WithSubst2 tm env d n where
  constructor Sub2
  term  : tm d from
  subst : Subst2 env d from n