module RewriteTactic where

open import Level renaming (zero to lzero; suc to lsuc)
open import Data.Nat as N using (ℕ; zero; suc; _+_) renaming (_≟_ to _≟N_)
open import Data.Nat.Properties.Simple
open import Data.Maybe using (Maybe; nothing; just)
open import Data.List
open import Data.Integer as I using (ℤ; ∣_∣)
open import Data.String using (String)
open import Function
open import Reflection
open import Relation.Nullary
open import Relation.Binary.PropositionalEquality

-- Something to return in the error cases
postulate 
  error : ∀ {a} {A : Set a} → String → A
  notImplemented : ∀ {a} {A : Set a} → String → A

unEl : Type → Term
unEl (el _ t) = t

unArg : {A : Set} → Arg A → A
unArg (arg _ x) = x

-- Disabled termination checking out of laziness, could probably get it to work by using sized types.
{-# NO_TERMINATION_CHECK #-}
shift : (n : ℤ) (cutoff : ℕ) → Term → Term
shiftArgs : (n : ℤ) (cutoff : ℕ) → List (Arg Term) → List (Arg Term)
shiftType : (n : ℤ) (cutoff : ℕ) → Type → Type
shiftSort : (n : ℤ) (cutoff : ℕ) → Sort → Sort

shiftArgs n cutoff          = map (λ { (arg i x) → arg i (shift n cutoff x) })
shiftType n cutoff (el s t) = el (shiftSort n cutoff s) (shift n cutoff t)
shiftSort n cutoff (set t)  = set (shift n cutoff t)
shiftSort n cutoff s        = s

shift n cutoff (var i args) =  
  case (cutoff N.≤? i) of λ { 
    (yes _) → case (n I.+ I.+ i) of λ {
      (I.-[1+ _ ]) → error "negative de Bruijn-index" ;
      (I.+ i')     → var i' (shiftArgs n cutoff args)
    } ; 
    (no  _) → var i (shiftArgs n cutoff args) 
  }
shift n cutoff (con c args)       = con c (shiftArgs n cutoff args)
shift n cutoff (def f args)       = def f (shiftArgs n cutoff args)
shift n cutoff (lam v t)          = lam v (shift n (suc cutoff) t)
shift n cutoff (pat-lam cs args)  = notImplemented "shifting pattern-matching lambdas"
shift n cutoff (pi (arg i t₁) t₂) = pi (arg i (shiftType n cutoff t₁)) (shiftType n (suc cutoff) t₂)
shift n cutoff (sort s)           = sort (shiftSort n cutoff s)
shift n cutoff (lit x)            = lit x
shift n cutoff unknown            = unknown

{-# NO_TERMINATION_CHECK #-}
apply : Term → List (Arg Term) → Term
sub : (n : ℕ) (s : Term) → Term → Term
subArgs : (n : ℕ) (s : Term) → List (Arg Term) → List (Arg Term)
subType : (n : ℕ) (s : Term) → Type → Type
subSort : (n : ℕ) (s : Term) → Sort → Sort

apply t [] = t
apply (var i args) xs = var i (args ++ xs)
apply (con c args) xs = con c (args ++ xs)
apply (def f args) xs = def f (args ++ xs)
apply (lam i t) (x ∷ xs) = apply (shift I.-[1+ 0 ] 0 (sub 0 (shift (I.+ 1) 0 (unArg x)) t)) xs
apply (pat-lam cs args) (x ∷ xs) = notImplemented "applying pattern matching lambdas"
apply (pi t₁ t₂) (x ∷ xs) = error "a pi is not a function"
apply (sort x) (x₁ ∷ xs) = error "a sort is not a function"
apply (lit x) (x₁ ∷ xs) = error "a literal is not a function"
apply unknown (x ∷ xs) = unknown

sub n s (var i args) with n ≟N i
sub n s (var .n args) | yes refl = apply s (subArgs n s args)
sub n s (var i args)  | no _     = var i (subArgs n s args)
sub n s (con c args) = con c (subArgs n s args)
sub n s (def f args) = def f (subArgs n s args)
sub n s (lam v t) = lam v (sub (suc n) (shift (I.+ 1) 0 s) t)
sub n s (pat-lam cs args) = notImplemented "substituting in pattern matching lambdas"
sub n s (pi (arg i t₁) t₂) = pi (arg i (subType n s t₁)) (subType (suc n) (shift (I.+ 1) 0 s) t₂)
sub n s (sort x) = sort (subSort n s x)
sub n s (lit x) = lit x
sub n s unknown = unknown

subArgs n s xs = map (λ { (arg i x) → arg i (sub n s x) }) xs
subType n s (el so ty) = el (subSort n s so) (sub n s ty)
subSort n s (set t) = set (sub n s t)
subSort n s (lit l) = lit l
subSort n s unknown = unknown

piApply : Term → List (Arg Term) → Term
piApply ty [] = ty
piApply (pi _ (el _ t)) (x ∷ xs) = piApply (shift I.-[1+ 0 ] 0 (sub 0 (shift (I.+ 1) 0 (unArg x)) t)) xs
piApply _ _ = error "not enough pie"

-- Replace all occurrences of s in t with 'var i'
{-# NO_TERMINATION_CHECK #-}
abstr : ℕ → Term → Term → Term
abstrArgs : ℕ → Term → List (Arg Term) → List (Arg Term)
abstrType : ℕ → Term → Type → Type
abstrSort : ℕ → Term → Sort → Sort

abstrArgs i s          = map (λ { (arg info t) → arg info (abstr i s t) })
abstrType i s (el l t) = el (abstrSort i s l) (abstr i s t)
abstrSort i s (set t)  = set (abstr i s t)
abstrSort i s l        = l

abstr i s t with s ≟ t
abstr i s t | yes _ = var i []
abstr i s (var x args)  | no _ = var x (abstrArgs i s args)
abstr i s (con c args)  | no _ = con c (abstrArgs i s args)
abstr i s (def f args)  | no _ = def f (abstrArgs i s args)
abstr i s (lam v t)     | no _ = lam v (abstr (suc i) (shift (I.+ 1) 0 s) t)
abstr i s (pat-lam _ _) | no _ = notImplemented "abstracting under pattern-matching lambdas"
abstr i s (pi (arg info t₁) t₂) | no _ = pi (arg info (abstrType i s t₁)) (abstrType (suc i) (shift (I.+ 1) 0 s) t₂)
abstr i s (sort x)      | no _ = sort (abstrSort i s x)
abstr i s (lit x)       | no _ = lit x
abstr i s unknown       | no _ = unknown

lookup : ∀ {a} {A : Set a} → List A → ℕ → Maybe A
lookup [] i = nothing
lookup (x ∷ xs) zero = just x
lookup (x ∷ xs) (suc i) = lookup xs i

typeOf : Term → Term
typeOf (var i args) = notImplemented "type inference for variables"
typeOf (con c args) = piApply (unEl (type c)) args
typeOf (def f args) = piApply (unEl (type f)) args
typeOf (lam v t) = notImplemented "type inference for lambdas"
typeOf (pat-lam cs args) = notImplemented "type inference for pattern matching lambdas"
typeOf (pi t₁ t₂) = sort (notImplemented "level inference")
typeOf (sort x) = sort (notImplemented "level inference")
typeOf (lit (nat x)) = def (quote ℕ) []
typeOf (lit (float x)) = def (notImplemented "floats") []
typeOf (lit (char x)) = def (notImplemented "chars") []
typeOf (lit (string x)) = def (quote String) []
typeOf (lit (name x)) = def (quote Name) []
typeOf unknown = unknown

rewrite′ : Term → Term → Term
rewrite′ e g = def (quote subst) (arg (arg-info visible relevant) motive ∷ arg (arg-info visible relevant) sym-e ∷ [])
  where
    sym-e = def (quote sym) (arg (arg-info visible relevant) e ∷ [])
    to-rewrite = case (typeOf e) of λ { (def (quote _≡_) (_ ∷ _ ∷ t ∷ _ ∷ [])) → unArg t ; _ → error "rewrite expects an element of type x ≡ y" }
    motive = lam visible (abstr 0 (shift (I.+ 1) 0 to-rewrite) (shift (I.+ 1) 0 g))

-- using the build-in rewrite
test : (x y : ℕ) → (x + y) + 0 ≡ y + (x + 0)
test x y 
  rewrite +-right-identity x
        | +-right-identity (x + y)
        | +-comm x y
  = refl

--using the rewrite′ tactic
test′ : (x y : ℕ) → (x + y) + 0 ≡ y + (x + 0)
test′ x y = 
   tactic (rewrite′ (quoteTerm (+-right-identity x)))
 | tactic (rewrite′ (quoteTerm (+-right-identity (x + y))))
 | tactic (rewrite′ (quoteTerm (+-comm x y))) 
 | refl