module FunctionBuilder2

import Data.Nat
import Data.Fin
import Data.Vect
import Decidable.Equality

%default total

namespace Untyped

  data Term : (n : Nat) -> Type where -- n max de brujin index
    Var : (i : Fin n) -> Term n
    App : (l : Term n) -> (r : Term n) -> Term n
    Lam : (b : Term (S n)) -> Term n

  shift : Term n -> Term (S n)
  shift (Var i) = Var (FS i)
  shift (App l r) = App (shift l) (shift r)
  shift (Lam b) = Lam (shift b)

  predFin : Fin (S n) -> Maybe (Fin n)
  predFin FZ = Nothing
  predFin (FS k) = Just k

  unshift : Fin (S n) -> Term (S n) -> Maybe (Term n)
  unshift i (Var j) = if i == j then Nothing else do
    k <- predFin j
    Just (Var k)
  unshift i (App l r) = do
    l <- unshift i l
    r <- unshift i r
    Just (App l r)
  unshift i (Lam b) = do
    b <- unshift (FS i) b
    Just (Lam b)

  replace : Term n -> Term (S n) -> Term n
  replace t (Var FZ) = t
  replace t (Var (FS i)) = Var i
  replace t (App l r) = App (replace t l) (replace t r)
  replace t (Lam b) = Lam (replace (shift t) b)

  covering
  eval : Term n -> Term n
  eval (Var i) = Var i
  eval (App l r) = case eval l of
    Lam b => eval (replace (eval r) b)
    l => App l (eval r)
  eval (Lam b) = case eval b of
    App l (Var FZ) => case unshift FZ l of
      Just l => eval l
      Nothing => Lam (App l (Var FZ))
    b => Lam b

  EvalTest1 : eval (Lam (Var 0)) = Lam (Var 0)
  EvalTest1 = Refl
  EvalTest2 : eval (App {n=100} (Lam (Var 0)) (Var 99)) = (Var 99)
  EvalTest2 = Refl
  EvalTest3 : eval (App (Lam (Var 0)) (Lam (Var 0))) = Lam (Var 0)
  EvalTest3 = Refl
  EvalTest4 : eval (Lam (App (Lam (Var 0)) (Lam (Var 0)))) = Lam (Lam (Var 0))
  EvalTest4 = Refl

VectDel :
  (i : Fin (S n)) -> (h : g) -> (t : Vect (S n) g) ->
  deleteAt (FS i) (h :: t) = h :: (deleteAt i t)
VectDel FZ h (x :: xs) = Refl
VectDel (FS k) h (x :: xs) = Refl

namespace SimplyTyped

  data TermType : Type where
    Ground : TermType
    Pi : TermType -> TermType -> TermType

  data Term : (tc : Vect n TermType) -> (tt : TermType) -> Type where
    Var : (i : Fin n) -> Term tc (index i tc)
    App : (l : Term tc (Pi itt ott)) -> (r : Term tc itt) -> Term tc ott
    Lam : (itt : TermType) -> (b : Term (itt :: tc) ott) -> Term tc (Pi itt ott)

  replaceVar :
    (n : Nat) ->
    (i : Fin (S n)) -> -- var index to replace
    (j : Fin (S n)) -> -- current var index
    (tc : Vect (S n) TermType) ->
    Either (i = j) (k : Fin n ** index k (deleteAt i tc) = (index j tc)) -- this is needed to allow recursion
  replaceVar Z FZ FZ tc = Left Refl
  replaceVar (S m) FZ FZ tc = Left Refl
  replaceVar (S m) FZ (FS y) (l :: ls) = Right (y ** Refl) 
  replaceVar (S m) (FS x) FZ (l :: ls@(_::_)) = Right (FZ ** Refl)
  replaceVar (S m) (FS x) (FS y) (l :: ls@(_::_)) = case replaceVar m x y ls of
    Left prf => Left (cong FS prf)
    Right (t ** p) => Right (FS t ** p)

  shiftVar : 
    (n : Nat) ->
    (i : Fin (S n)) -> -- var index to shift
    (j : Fin n) -> -- current var index
    (tc : Vect n TermType) ->
    (k : Fin (S n) ** index k (insertAt i u tc) = index j tc) -- this is needed to allow recursion
  shiftVar (S m) FZ FZ (l :: ls) = (FS FZ ** Refl)
  shiftVar (S m) FZ (FS x) (l :: ls) = (FS (FS x) ** Refl)
  shiftVar (S m) (FS x) FZ (l :: ls) = (FZ ** Refl)
  shiftVar (S m) (FS x) (FS y) (l :: ls) = case shiftVar m x y ls of
    (k ** prf) => (FS k ** prf)
  
  shiftAt :
    (n : Nat) -> (i : Fin (S n)) -> (ntt : TermType) -> (tc : Vect n TermType) ->
    Term tc tt -> Term (insertAt i ntt tc) tt
  shiftAt n i ntt tc (Var j) = case shiftVar n i j tc of
    (k ** prf) => rewrite sym prf in Var k
  shiftAt n i ntt tc (App l r) = App (shiftAt n i ntt tc l) (shiftAt n i ntt tc r)
  shiftAt n i ntt tc (Lam itt b) = Lam itt (shiftAt (S n) (FS i) ntt (itt :: tc) b)

  shift : {ntt : TermType} -> {n : Nat} -> {tc : Vect n TermType} -> 
    Term tc tt -> Term (ntt :: tc) tt
  shift t = shiftAt n FZ ntt tc t

  replaceAt :
    {n : Nat} ->
    (i : Fin (S n)) ->
    (tc : Vect (S n) TermType) ->
    Term (deleteAt i tc) (index i tc) ->
    Term tc tt ->
    Term (deleteAt i tc) tt
  replaceAt i tc t (Var j) = case replaceVar n i j tc of
    Left prf => rewrite sym prf in t
    Right (k ** prf) => rewrite sym prf in Var k
  replaceAt i tc t (App l r) =
    App (replaceAt i tc t l) (replaceAt i tc t r)
  replaceAt i tc t (Lam itt b) =
    Lam itt (rewrite sym (VectDel i itt tc) in replaceAt (FS i) (itt :: tc) (rewrite (VectDel i itt tc) in (shift t)) b)

  covering
  eval :
    {n : Nat} -> {tc : Vect n TermType} ->
    Term tc tt -> Term tc tt
  eval (Var i) = Var i
  eval (App l r) = case eval l of
    Lam itt b => eval (replaceAt FZ (itt :: tc) (eval r) b)
    l => App l (eval r) 
  eval (Lam itt b) = Lam itt (eval b)