module Quox.Syntax.Subst

import public Quox.Syntax.Shift
import Quox.Var
import Quox.Name

import Data.Nat
import Data.List
import Data.SnocVect
import Data.Singleton
import Derive.Prelude

%default total
%language ElabReflection


public export
data Subst : (Nat -> Type) -> Nat -> Nat -> Type where
  Shift : Shift from to -> Subst env from to
  (:::) : (t : Lazy (env to)) -> Subst env from to -> Subst env (S from) to
%name Subst th, ph, ps

export infixr 7 !:::
||| in case the automatic laziness insertion gets confused
public export
(!:::) : env to -> Subst env from to -> Subst env (S from) to
t !::: ts = t ::: ts


private
Repr : (Nat -> Type) -> Nat -> Type
Repr f to = (List (f to), Nat)

private
repr : Subst f from to -> Repr f to
repr (Shift by) = ([], by.nat)
repr (t ::: th) = let (ts, i) = repr th in (t::ts, i)


export Eq   (f to) => Eq   (Subst f from to) where (==) = (==) `on` repr
export Ord  (f to) => Ord  (Subst f from to) where compare = compare `on` repr
export Show (f to) => Show (Subst f from to) where show = show . repr


export infixl 8 //?, //


public export
interface FromVarR term => CanSubstSelfR term where
  (//?) : {from, to : Nat} ->
          term from -> Lazy (Subst term from to) -> term to

public export
interface (FromVar term, CanSubstSelfR term) => CanSubstSelf term where
  (//) : term from -> Lazy (Subst term from to) -> term to
  0 substSame : (t : term from) -> (th : Subst term from to) ->
                t //? th === t // th


public export
getR : {to : Nat} -> FromVarR term =>
       Subst term from to -> Var from -> Loc -> term to
getR (Shift by) i      loc = fromVarR (shift by i) loc
getR (t ::: th) VZ     _   = t
getR (t ::: th) (VS i) loc = getR th i loc

public export
get : FromVar term => Subst term from to -> Var from -> Loc -> term to
get (Shift by) i      loc = fromVar (shift by i) loc
get (t ::: th) VZ     _   = t
get (t ::: th) (VS i) loc = get th i loc


public export
substVar : Var from -> Lazy (Subst Var from to) -> Var to
substVar i      (Shift by) = shift by i
substVar VZ     (t ::: th) = t
substVar (VS i) (t ::: th) = substVar i th

public export
CanSubstSelfR Var where (//?) = substVar

public export
CanSubstSelf Var where (//) = substVar; substSame _ _ = Refl


public export %inline
shift : (by : Nat) -> Subst env from (by + from)
shift by = Shift $ fromNat by

public export %inline
shift0 : (by : Nat) -> Subst env 0 by
shift0 by = rewrite sym $ plusZeroRightNeutral by in Shift $ fromNat by


export infixr 9 .?

public export
(.?) : CanSubstSelfR f => {from, mid, to : Nat} ->
       Subst f from mid -> Subst f mid to -> Subst f from to
Shift by      .? Shift bz   = Shift $ by . bz
Shift SZ      .? ph         = ph
Shift (SS by) .? (t ::: th) = Shift by .? th
(t ::: th)    .? ph         = (t //? ph) ::: (th .? ph)

public export
(.) : CanSubstSelf f => Subst f from mid -> Subst f mid to -> Subst f from to
Shift by      . Shift bz   = Shift $ by . bz
Shift SZ      . ph         = ph
Shift (SS by) . (t ::: th) = Shift by . th
(t ::: th)    . ph         = (t // ph) ::: (th . ph)


public export %inline
id : Subst f n n
id = shift 0

public export
traverse : Applicative m =>
           (f to -> m (g to)) -> Subst f from to -> m (Subst g from to)
traverse f (Shift by) = pure $ Shift by
traverse f (t ::: th) = [|f t !::: traverse f th|]

-- not in terms of traverse because this map can maintain laziness better
public export
map : (f to -> g to) -> Subst f from to -> Subst g from to
map f (Shift by) = Shift by
map f (t ::: th) = f t ::: map f th


public export
pushNR : {from, to : Nat} -> CanSubstSelfR f => (s : Nat) -> Loc ->
         Subst f from to -> Subst f (s + from) (s + to)
pushNR 0     _   th = th
pushNR (S s) loc th =
  rewrite plusSuccRightSucc s from in
  rewrite plusSuccRightSucc s to in
  pushNR s loc $ fromVarR VZ loc ::: (th .? shift 1)

public export %inline
pushR : {from, to : Nat} -> CanSubstSelfR f =>
        Loc -> Subst f from to -> Subst f (S from) (S to)
pushR = pushNR 1

public export
pushN : CanSubstSelf f => (s : Nat) -> Loc ->
        Subst f from to -> Subst f (s + from) (s + to)
pushN 0     _   th = th
pushN (S s) loc th =
  rewrite plusSuccRightSucc s from in
  rewrite plusSuccRightSucc s to in
  pushN s loc $ fromVar VZ loc ::: (th . shift 1)

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

public export
drop1 : Subst f (S from) to -> Subst f from to
drop1 (Shift by) = Shift $ ssDown by
drop1 (t ::: th) = th


public export
fromSnocVect : SnocVect s (f n) -> Subst f (s + n) n
fromSnocVect [<]       = id
fromSnocVect (xs :< x) = x ::: fromSnocVect xs

public export %inline
one : f n -> Subst f (S n) n
one x = fromSnocVect [< x]


export
getFrom : {to : Nat} -> Subst _ from to -> Singleton from
getFrom (Shift by) = getFrom by
getFrom (t ::: th) = [|S $ getFrom th|]


||| whether two substitutions with the same codomain have the same shape
||| (the same number of terms and the same shift at the end). if so, they
||| also have the same domain
export
cmpShape : Subst env from1 to -> Subst env from2 to ->
           Either Ordering (from1 = from2)
cmpShape (Shift by) (Shift bz) = cmpLen by bz
cmpShape (Shift _)  (_ ::: _)  = Left LT
cmpShape (_ ::: _)  (Shift _)  = Left GT
cmpShape (_ ::: th) (_ ::: ph) = map (\x => cong S x) $ cmpShape th ph


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

export
(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 && s1 == s2

export
(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, s1) (t2, s2)

export %hint
ShowWithSubst : (Show (env n), forall n. Show (tm n)) =>
                Show (WithSubst tm env n)
ShowWithSubst = deriveShow

public export
record WithSubstR tm env n where
  constructor SubR
  {from : Nat}
  term  : tm from
  subst : Lazy (Subst env from n)

export
(Eq (env n), forall n. Eq (tm n)) => Eq (WithSubstR tm env n) where
  SubR {from = m1} t1 s1 == SubR {from = m2} t2 s2 =
    case decEq m1 m2 of
      Yes Refl => t1 == t2 && s1 == s2
      No  _    => False

export
(Ord (env n), forall n. Ord (tm n)) => Ord (WithSubstR tm env n) where
  SubR {from = m1} t1 s1 `compare` SubR {from = m2} t2 s2 =
    case cmp m1 m2 of
      CmpLT _ => LT
      CmpEQ   => compare (t1, s1) (t2, s2)
      CmpGT _ => GT

export %hint
ShowWithSubstR : (Show (env n), forall n. Show (tm n)) =>
                 Show (WithSubstR tm env n)
ShowWithSubstR = deriveShow