module Core.Environment

import Core.Context
import Core.Term
import Core.Thinning
import Core.Var

import Syntax.PreorderReasoning

%prefix_record_projections off

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

public export
data Env : Context -> Type where
  Lin : Env [<]
  Add : Env sx -> Thinned sx -> Env (sx :< n)

%name Env env

-- Constructors ----------------------------------------------------------------

public export
(:<) : Env sx -> Term sx -> Env (sx :< n)
env :< t = Add env (t `Over` id sx)

-- Equality --------------------------------------------------------------------

namespace Eq
  public export
  data EnvEq : Env sx -> Env sx -> Type where
    Base : EnvEq env env
    (:<) : EnvEq env1 env2 -> expand t = expand u -> EnvEq (Add env1 t) (Add env2 u)

  %name EnvEq prf

-- Extension -------------------------------------------------------------------

public export
data IsExtension : sx `Thins` sy -> Env sy -> Env sx -> Type where
  IdWf : IsExtension (id sx) env env
  DropWf : IsExtension thin env2 env1 -> IsExtension (drop thin n) (Add env2 t) env1
  KeepWf :
    IsExtension thin env2 env1 ->
    IsExtension (keep thin n) (Add env2 (shift t thin)) (Add env1 t)

%name IsExtension ext

-- Composition

CompWf :
  IsExtension thin2 env3 env2 ->
  IsExtension thin1 env2 env1 ->
  IsExtension (thin2 . thin1) env3 env1
CompWf IdWf thinWf1 = rewrite compLeftIdentity thin1 in thinWf1
CompWf (DropWf {thin = thin2, n} thinWf2) thinWf1 =
  rewrite compLeftDrop thin2 thin1 n in
  DropWf (CompWf thinWf2 thinWf1)
CompWf (KeepWf {thin = thin2, n} thinWf2) IdWf =
  rewrite compRightIdentity $ keep thin2 n in
  KeepWf thinWf2
CompWf (KeepWf {thin = thin2, n} thinWf2) (DropWf {thin = thin1} thinWf1) =
  rewrite compRightDrop thin2 thin1 n in
  DropWf (CompWf thinWf2 thinWf1)
CompWf (KeepWf {thin = thin2, n} thinWf2) (KeepWf {thin = thin1, t} thinWf1) =
  rewrite compKeep thin2 thin1 n in
  rewrite shiftHomo t thin1 thin2 in
  KeepWf (CompWf thinWf2 thinWf1)

-- Indexing --------------------------------------------------------------------

doIndex : Env sx -> {0 i : Var sx} -> View i -> Thinned sx
doIndex (Add env t) Here = shift t (wkn1 _ _)
doIndex (Add env t) (There i) = shift (doIndex env $ view i) (wkn1 _ _)

export
index : Env sx -> Var sx -> Term sx
index env i = expand $ doIndex env $ view i

-- Respects Equality

doIndexEqIsEq :
  EnvEq env1 env2 ->
  {0 i : Var sx} ->
  (view : View i) ->
  expand (doIndex env1 view) = expand (doIndex env2 view)
doIndexEqIsEq Base view = Refl
doIndexEqIsEq {sx = sx :< n} ((:<) {t, u} prf eq) Here = Calc $
  |~ expand (shift t $ wkn1 sx n)
  ~~ wkn (expand t) (wkn1 sx n)   ...(expandShift t (wkn1 sx n))
  ~~ wkn (expand u) (wkn1 sx n)   ...(cong (\t => wkn t (wkn1 sx n)) $ eq)
  ~~ expand (shift u $ wkn1 sx n) ...(sym $ expandShift u (wkn1 sx n))
doIndexEqIsEq {sx = sx :< n} ((:<) {env1, env2} prf eq) (There i) = Calc $
  |~ expand (shift (doIndex env1 $ view i) (wkn1 sx n))
  ~~ wkn (expand (doIndex env1 $ view i)) (wkn1 sx n)   ...(expandShift (doIndex env1 $ view i) (wkn1 sx n))
  ~~ wkn (expand (doIndex env2 $ view i)) (wkn1 sx n)   ...(cong (\t => wkn t (wkn1 sx n)) $ doIndexEqIsEq prf (view i))
  ~~ expand (shift (doIndex env2 $ view i) (wkn1 sx n)) ...(sym $ expandShift (doIndex env2 $ view i) (wkn1 sx n))

export
indexEqIsEq : EnvEq env1 env2 -> (i : Var sx) -> index env1 i = index env2 i
indexEqIsEq prf i = doIndexEqIsEq prf $ view i