module Core.Environment.Extension

import Core.Environment
import Core.Term
import Core.Term.Thinned
import Core.Thinning

import Syntax.PreorderReasoning

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

public export
data Extends : m `Thins` n -> Env n -> Env m -> Type where
  Id : Extends (id n) env env
  Drop : Extends thin env2 env1 -> Extends (drop thin) (env2 :< a) env1
  Keep : Extends thin env2 env1 -> Extends (keep thin) (env2 :< wkn a thin) (env1 :< a)

%name Extends ext

-- Smart Constructor -----------------------------------------------------------

export
Keep' : Extends thin env1 env2 -> Extends (keep thin) (env1 :< (a `Over` thin)) (env2 :< pure a)
Keep' ext = rewrite sym $ cong (a `Over`) $ compRightId thin in Keep ext

-- Composition -----------------------------------------------------------------

export
(.) : Extends thin2 env3 env2 -> Extends thin1 env2 env1 -> Extends (thin2 . thin1) env3 env1
Id . ext1 =
  rewrite compLeftId thin1 in
  ext1
Drop {thin = thin2} ext2 . ext1 =
  rewrite compLeftDrop thin2 thin1 in
  Drop (ext2 . ext1)
Keep {thin = thin2} ext2 . Id =
  rewrite compRightId (keep thin2) in
  Keep ext2
Keep {thin = thin2} ext2 . Drop {thin = thin1} ext1 =
  rewrite compRightDrop thin2 thin1 in
  Drop (ext2 . ext1)
Keep {thin = thin2} ext2 . Keep {thin = thin1, a} ext1 =
  rewrite compKeep thin2 thin1 in
  rewrite wknHomo a thin1 thin2 in
  Keep (ext2 . ext1)

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

export
indexHomo :
  Extends thin env2 env1 ->
  (i : Fin m) ->
  index env2 (wkn i thin) = wkn (index env1 i) thin
indexHomo Id i = Calc $
  |~ index env1 (wkn i $ id m)
  ~~ index env1 i              ...(cong (index env1) $ wknId i)
  ~~ wkn (index env1 i) (id m) ...(sym $ wknId (index env1 i))
indexHomo (Drop {thin, env2, env1, a} ext) i = Calc $
  |~ index (env2 :< a) (wkn i $ drop thin)
  ~~ wkn (index env2 (wkn i thin)) (wkn1 _) ...(cong (index $ env2 :< a) $ wknDrop i thin)
  ~~ wkn (wkn (index env1 i) thin) (wkn1 _) ...(cong (flip wkn (wkn1 _)) $ indexHomo ext i)
  ~~ wkn (index env1 i) (wkn1 _ . thin)     ...(wknHomo (index env1 i) thin (wkn1 _))
  ~~ wkn (index env1 i) (drop thin)         ...(cong (wkn (index env1 i)) $ compLeftWkn1 thin)
indexHomo (Keep {thin, env2, env1, a} ext) FZ = Calc $
  |~ index (env2 :< wkn a thin) (wkn FZ $ keep thin)
  ~~ wkn (wkn a thin) (wkn1 _)                       ...(cong (index $ env2 :< wkn a thin) $ wknKeepFZ thin)
  ~~ wkn (wkn a (wkn1 _)) (keep thin)                ...(wkn1Comm a thin)
indexHomo (Keep {thin, env2, env1, a} ext) (FS i) = Calc $
  |~ index (env2 :< wkn a thin) (wkn (FS i) $ keep thin)
  ~~ wkn (index env2 $ wkn i thin) (wkn1 _)              ...(cong (index $ env2 :< wkn a thin) $ wknKeepFS i thin)
  ~~ wkn (wkn (index env1 i) thin) (wkn1 _)              ...(cong (flip wkn (wkn1 _)) $ indexHomo ext i)
  ~~ wkn (wkn (index env1 i) (wkn1 _)) (keep thin)       ...(wkn1Comm (index env1 i) thin)