module Total.Term

import Control.Relation
import public Data.SnocList.Elem
import public Subst
import public Thinning

import Syntax.PreorderReasoning
import Syntax.PreorderReasoning.Generic

%prefix_record_projections off

infixr 10 ~>

-- Types -----------------------------------------------------------------------

public export
data Ty : Type where
  N : Ty
  (~>) : Ty -> Ty -> Ty

%name Ty ty

-- Definition ------------------------------------------------------------------

public export
data Term : Ty -> SnocList Ty -> Type where
  Var : (i : Elem ty ctx) -> Term ty ctx
  Abs : Term ty' (ctx :< ty) -> Term (ty ~> ty') ctx
  App : {ty : Ty} -> Term (ty ~> ty') ctx -> Term ty ctx -> Term ty' ctx
  Zero : Term N ctx
  Suc : Term N ctx -> Term N ctx
  Rec : Term N ctx -> Term ty ctx -> Term (ty ~> ty) ctx -> Term ty ctx

%name Term t, u, v

-- Raw Interfaces --------------------------------------------------------------

export
Point Ty Term where
  point = Var

export
Weaken Ty Term where
  wkn (Var i) thin = Var (index thin i)
  wkn (Abs t) thin = Abs (wkn t (keep thin _))
  wkn (App t u) thin = App (wkn t thin) (wkn u thin)
  wkn Zero thin = Zero
  wkn (Suc t) thin = Suc (wkn t thin)
  wkn (Rec t u v) thin = Rec (wkn t thin) (wkn u thin) (wkn v thin)

export
PointWeaken Ty Term where

-- Substitution ----------------------------------------------------------------

public export
subst : Term ty ctx -> Subst Term ctx ctx' -> Term ty ctx'
subst (Var i) sub = index sub i
subst (Abs t) sub = Abs (subst t $ lift sub)
subst (App t u) sub = App (subst t sub) (subst u sub)
subst Zero sub = Zero
subst (Suc t) sub = Suc (subst t sub)
subst (Rec t u v) sub = Rec (subst t sub) (subst u sub) (subst v sub)

-- Properties ------------------------------------------------------------------

-- Utilities

export
cong3 : (0 f : a -> b -> c -> d) -> x1 = x2 -> y1 = y2 -> z1 = z2 -> f x1 y1 z1 = f x2 y2 z2
cong3 f Refl Refl Refl = Refl

-- Weakening

export
FullWeaken Ty Term where
  shiftIsWkn t = Refl

  wknId (Var i) = cong Var $ indexId i
  wknId (Abs t) = cong Abs $ trans (cong (wkn t) (keepId _ _)) (wknId t)
  wknId (App t u) = cong2 App (wknId t) (wknId u)
  wknId Zero = Refl
  wknId (Suc t) = cong Suc (wknId t)
  wknId (Rec t u v) = cong3 Rec (wknId t) (wknId u) (wknId v)

  wknHomo (Var i) thin1 thin2 = cong Var $ indexHomo thin1 thin2 i
  wknHomo (Abs t) thin1 thin2 =
    cong Abs $
      trans
        (wknHomo t (keep thin1 _) (keep thin2 _))
        (cong (wkn t) $ keepHomo thin2 thin1 _)
  wknHomo (App t u) thin1 thin2 = cong2 App (wknHomo t thin1 thin2) (wknHomo u thin1 thin2)
  wknHomo Zero thin1 thin2 = Refl
  wknHomo (Suc t) thin1 thin2 = cong Suc (wknHomo t thin1 thin2)
  wknHomo (Rec t u v) thin1 thin2 =
    cong3 Rec
      (wknHomo t thin1 thin2)
      (wknHomo u thin1 thin2)
      (wknHomo v thin1 thin2)

export
FullPointWeaken Ty Term where
  wknPoint i thin = Refl

-- Substitution & Weakening

export
substCong :
  (t : Term ty ctx) ->
  {0 sub1, sub2 : Subst Term ctx ctx'} ->
  Equiv sub1 sub2 ->
  subst t sub1 = subst t sub2
substCong (Var i) e = e.equiv i
substCong (Abs t) e = cong Abs $ substCong t (liftCong e)
substCong (App t u) e = cong2 App (substCong t e) (substCong u e)
substCong Zero e = Refl
substCong (Suc t) e = cong Suc (substCong t e)
substCong (Rec t u v) e = cong3 Rec (substCong t e) (substCong u e) (substCong v e)

export
wknSubst :
  (t : Term ty ctx) ->
  (sub : Subst Term ctx ctx') ->
  (thin : ctx' `Thins` ctx'') ->
  wkn (subst t sub) thin = subst t (wkn sub thin)
wknSubst (Var i) sub thin =
  indexWkn sub thin i
wknSubst (Abs {ty} t) sub thin =
  cong Abs $ Calc $
    |~ wkn (subst t (lift sub)) (keep thin ty)
    ~~ subst t (wkn (lift sub) (keep thin ty)) ...(wknSubst t (lift sub) (keep thin _))
    ~~ subst t (lift (wkn sub thin))           ...(substCong t (wknLift sub thin))
wknSubst (App t u) sub thin =
  cong2 App (wknSubst t sub thin) (wknSubst u sub thin)
wknSubst Zero sub thin =
  Refl
wknSubst (Suc t) sub thin =
  cong Suc (wknSubst t sub thin)
wknSubst (Rec t u v) sub thin =
  cong3 Rec (wknSubst t sub thin) (wknSubst u sub thin) (wknSubst v sub thin)

export
substWkn :
  (t : Term ty ctx) ->
  (sub : Subst Term ctx' ctx'') ->
  (thin : ctx `Thins` ctx') ->
  subst (wkn t thin) sub = subst t (restrict sub thin)
substWkn (Var i) sub thin =
  indexRestrict sub thin i
substWkn (Abs t) sub thin =
  cong Abs $ Calc $
    |~ subst (wkn t (keep thin _)) (lift sub)
    ~~ subst t (restrict (lift sub) (keep thin _)) ...(substWkn t (lift sub) (keep thin _))
    ~~ subst t (lift (restrict sub thin))          ...(substCong t (restrictLift sub thin))
substWkn (App t u) sub thin =
  cong2 App (substWkn t sub thin) (substWkn u sub thin)
substWkn Zero sub thin =
  Refl
substWkn (Suc t) sub thin =
  cong Suc (substWkn t sub thin)
substWkn (Rec t u v) sub thin =
  cong3 Rec (substWkn t sub thin) (substWkn u sub thin) (substWkn v sub thin)

-- -- substBase : (t : Term ctx ty) -> (thin : ctx `Thins` ctx') -> subst t (Base thin) = wkn t thin
-- -- substBase (Var i) thin = Refl
-- -- substBase (Abs t) thin =
-- --   rewrite identityLeft thin in
-- --   cong Abs $ Calc $
-- --     |~ subst t (Base (Drop thin) :< Var Here)
-- --     ~~ subst t (Base $ Keep thin)             ...(sym $ substCong t Base)
-- --     ~~ wkn t (Keep thin)                      ...(substBase t (Keep thin))
-- -- substBase (App t u) thin = cong2 App (substBase t thin) (substBase u thin)
-- -- substBase Zero thin = Refl
-- -- substBase (Suc t) thin = cong Suc (substBase t thin)
-- -- substBase (Rec t u v) thin = cong3 Rec (substBase t thin) (substBase u thin) (substBase v thin)

-- -- -- Substitution Composition

-- -- indexComp :
-- --   (sub1 : Terms ctx' ctx) ->
-- --   (sub2 : Terms ctx'' ctx') ->
-- --   (i : Elem ty ctx) ->
-- --   index (sub2 . sub1) i = subst (index sub1 i) sub2
-- -- indexComp (Base thin) sub2 i = indexRestrict thin sub2 i
-- -- indexComp (sub1 :< t) sub2 Here = Refl
-- -- indexComp (sub1 :< t) sub2 (There i) = indexComp sub1 sub2 i

-- -- wknAllComp :
-- --   (sub1 : Terms ctx' ctx) ->
-- --   (sub2 : Terms ctx'' ctx') ->
-- --   (thin : ctx'' `Thins` ctx''') ->
-- --   wknAll sub2 thin . sub1 = wknAll (sub2 . sub1) thin
-- -- wknAllComp (Base thin') sub2 thin = wknAllRestrict thin' sub2 thin
-- -- wknAllComp (sub1 :< t) sub2 thin =
-- --   cong2 (:<)
-- --     (wknAllComp sub1 sub2 thin)
-- --     (sym $ wknSubst t sub2 thin)

-- -- compDrop :
-- --   (sub1 : Terms ctx' ctx) ->
-- --   (thin : ctx' `Thins` ctx'') ->
-- --   (sub2 : Terms ctx''' ctx'') ->
-- --   sub2 . wknAll sub1 thin = restrict sub2 thin . sub1
-- -- compDrop (Base thin') thin sub2 = restrictHomo sub2 thin thin'
-- -- compDrop (sub1 :< t) thin sub2 = cong2 (:<) (compDrop sub1 thin sub2) (substWkn t thin sub2)

-- -- export
-- -- compWknAll :
-- --   (sub1 : Terms ctx' ctx) ->
-- --   (sub2 : Terms ctx''' ctx'') ->
-- --   (thin : ctx' `Thins` ctx'') ->
-- --   sub2 . wknAll sub1 thin = restrict sub2 thin . sub1
-- -- compWknAll (Base thin') sub2 thin = restrictHomo sub2 thin thin'
-- -- compWknAll (sub1 :< t) sub2 thin = cong2 (:<) (compWknAll sub1 sub2 thin) (substWkn t thin sub2)

-- -- export
-- -- baseComp :
-- --   (thin : ctx' `Thins` ctx'') ->
-- --   (sub : Terms ctx' ctx) ->
-- --   Base thin . sub = wknAll sub thin
-- -- baseComp thin (Base thin') = Refl
-- -- baseComp thin (sub :< t) = cong2 (:<) (baseComp thin sub) (substBase t thin)

-- -- -- Substitution

-- -- export
-- -- substId : (len : Len ctx) => (t : Term ctx ty) -> subst t (Base $ id @{len}) = t
-- -- substId (Var i) = cong Var (indexId i)
-- -- substId (Abs t) =
-- --   rewrite identitySquared len in
-- --   cong Abs $ trans (sym $ substCong t Base) (substId t)
-- -- substId (App t u) = cong2 App (substId t) (substId u)
-- -- substId Zero = Refl
-- -- substId (Suc t) = cong Suc (substId t)
-- -- substId (Rec t u v) = cong3 Rec (substId t) (substId u) (substId v)

-- -- export
-- -- substHomo :
-- --   (t : Term ctx ty) ->
-- --   (sub1 : Terms ctx' ctx) ->
-- --   (sub2 : Terms ctx'' ctx') ->
-- --   subst (subst t sub1) sub2 = subst t (sub2 . sub1)
-- -- substHomo (Var i) sub1 sub2 =
-- --   sym $ indexComp sub1 sub2 i
-- -- substHomo (Abs t) sub1 sub2 =
-- --   cong Abs $ Calc $
-- --     |~ subst (subst t (wknAll sub1 (Drop id) :< Var Here)) (wknAll sub2 (Drop id) :< Var Here)
-- --     ~~ subst t ((wknAll sub2 (Drop id) :< Var Here) . (wknAll sub1 (Drop id) :< Var Here))
-- --       ...(substHomo t (wknAll sub1 (Drop id) :< Var Here) (wknAll sub2 (Drop id) :< Var Here))
-- --     ~~ subst t ((wknAll sub2 (Drop id) :< Var Here) . wknAll sub1 (Drop id) :< Var Here)
-- --       ...(Refl)
-- --     ~~ subst t (restrict (wknAll sub2 (Drop id)) id . sub1 :< Var Here)
-- --       ...(cong (subst t . (:< Var Here)) $ compDrop sub1 (Drop id) (wknAll sub2 (Drop id) :< Var Here))
-- --     ~~ subst t (wknAll sub2 (Drop id) . sub1 :< Var Here)
-- --       ...(cong (subst t . (:< Var Here) . (. sub1)) $ restrictId (wknAll sub2 (Drop id)))
-- --     ~~ subst t (wknAll (sub2 . sub1) (Drop id) :< Var Here)
-- --       ...(cong (subst t . (:< Var Here)) $ wknAllComp sub1 sub2 (Drop id))
-- -- substHomo (App t u) sub1 sub2 =
-- --   cong2 App (substHomo t sub1 sub2) (substHomo u sub1 sub2)
-- -- substHomo Zero sub1 sub2 =
-- --   Refl
-- -- substHomo (Suc t) sub1 sub2 =
-- --   cong Suc (substHomo t sub1 sub2)
-- -- substHomo (Rec t u v) sub1 sub2 =
-- --   cong3 Rec (substHomo t sub1 sub2) (substHomo u sub1 sub2) (substHomo v sub1 sub2)