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

-- Interaction with Extension

doWknIndex :
  {0 thin : sx `Thins` sy} ->
  IsExtension thin env2 env1 ->
  (v : View i) ->
  wkn (expand $ doIndex env1 v) thin = expand (doIndex env2 $ view $ wkn i thin)
doWknIndex IdWf v =
  rewrite wknId i in
  Calc $
  |~ wkn (expand $ doIndex env2 v) (id sy)
  ~~ expand (doIndex env2 v)               ...(wknId (expand $ doIndex env2 v))
  ~~ expand (doIndex env2 $ view i)        ...(cong (expand . doIndex env2) $ viewUnique v (view i))
doWknIndex {sy = sy :< n} (DropWf {env2, thin, n, t} ext) v =
  rewrite wknDropIsThere i thin n in
  rewrite viewUnique (view $ there {n} $ wkn i thin) (There $ wkn i thin) in
  Calc $
  |~ wkn (expand (doIndex env1 v)) (drop thin n)
  ~~ wkn (expand (doIndex env1 v)) (wkn1 sy n . thin)            ...(sym $ cong (wkn _) $ compLeftWkn1 thin n)
  ~~ wkn (wkn (expand $ doIndex env1 v) thin) (wkn1 sy n)        ...(sym $ wknHomo (expand $ doIndex env1 v) thin (wkn1 sy n))
  ~~ wkn (expand $ doIndex env2 $ view $ wkn i thin) (wkn1 sy n) ...(cong (\t => wkn t (wkn1 sy n)) $ doWknIndex ext v)
  ~~ expand (doIndex (Add env2 t) (There $ wkn i thin))          ...(sym $ expandShift (doIndex env2 $ view $ wkn i thin) (wkn1 sy n))
doWknIndex {sx = sx :< n, sy = sy :< n} (KeepWf {thin, n, t} ext) Here =
  rewrite wknKeepHereIsHere thin n in
  rewrite viewUnique (view $ here {sx = sy, n}) Here in
  Calc $
  |~ wkn (expand $ shift t (wkn1 sx n)) (keep thin n)
  ~~ wkn (wkn (expand t) (wkn1 sx n)) (keep thin n)   ...(cong (\t => wkn t (keep thin n)) $ expandShift t (wkn1 sx n))
  ~~ wkn (expand t) (keep thin n . wkn1 sx n)         ...(wknHomo (expand t) (wkn1 sx n) (keep thin n))
  ~~ wkn (expand t) (drop thin n)                     ...(cong (wkn $ expand t) $ compRightWkn1 thin n)
  ~~ wkn (expand t) (wkn1 sy n . thin)                ...(sym $ cong (wkn $ expand t) $ compLeftWkn1 thin n)
  ~~ wkn (wkn (expand t) thin) (wkn1 sy n)            ...(sym $ wknHomo (expand t) thin (wkn1 sy n))
  ~~ wkn (expand $ shift t thin) (wkn1 sy n)          ...(sym $ cong (\t => wkn t (wkn1 sy n)) $ expandShift t thin)
  ~~ expand (shift (shift t thin) (wkn1 sy n))        ...(sym $ expandShift (shift t thin) (wkn1 sy n))
doWknIndex {sx = sx :< n, sy = sy :< n} (KeepWf {env1, env2, thin, n, t} ext) (There i) =
  rewrite wknKeepThereIsThere i thin n in
  rewrite viewUnique (view $ there {n} $ wkn i thin) (There $ wkn i thin) in
  Calc $
  |~ wkn (expand $ shift (doIndex env1 $ view i) (wkn1 sx n)) (keep thin n)
  ~~ wkn (wkn (expand $ doIndex env1 $ view i) (wkn1 sx n)) (keep thin n)   ...(cong (\t => wkn t (keep thin n)) $ expandShift (doIndex env1 $ view i) (wkn1 sx n))
  ~~ wkn (expand $ doIndex env1 $ view i) (keep thin n . wkn1 sx n)         ...(wknHomo _ (wkn1 sx n) (keep thin n))
  ~~ wkn (expand $ doIndex env1 $ view i) (drop thin n)                     ...(cong (wkn (expand $ doIndex env1 $ view i)) $ compRightWkn1 thin n)
  ~~ wkn (expand $ doIndex env1 $ view i) (wkn1 sy n . thin)                ...(sym $ cong (wkn _) $ compLeftWkn1 thin n)
  ~~ wkn (wkn (expand $ doIndex env1 $ view i) thin) (wkn1 sy n)            ...(sym $ wknHomo _ thin (wkn1 sy n))
  ~~ wkn (expand $ doIndex env2 $ view $ wkn i thin) (wkn1 sy n)            ...(cong (\t => wkn t (wkn1 sy n)) $ doWknIndex ext $ view i)
  ~~ expand (doIndex (Add env2 (shift t thin)) (There $ wkn i thin))        ...(sym $ expandShift (doIndex env2 $ view $ wkn i thin) (wkn1 sy n))

export
wknIndex :
  IsExtension thin env2 env1 ->
  (i : Var sx) ->
  wkn (index env1 i) thin = index env2 (wkn i thin)
wknIndex ext i = doWknIndex ext (view i)