module Quox.Syntax.Term.Reduce

import Quox.Syntax.Term.Base
import Quox.Syntax.Term.Subst

%default total


mutual
  ||| true if a term has a closure or dimension closure at the top level,
  ||| or is `E` applied to such an elimination
  public export %inline
  topCloT : Term d n -> Bool
  topCloT (CloT  _ _) = True
  topCloT (DCloT _ _) = True
  topCloT (E e)       = topCloE e
  topCloT _           = False

  ||| true if an elimination has a closure or dimension closure at the top level
  public export %inline
  topCloE : Elim d n -> Bool
  topCloE (CloE  _ _) = True
  topCloE (DCloE _ _) = True
  topCloE _           = False


public export IsNotCloT : Term d n -> Type
IsNotCloT = So . not . topCloT

||| a term which is not a top level closure
public export NotCloTerm : Nat -> Nat -> Type
NotCloTerm d n = Subset (Term d n) IsNotCloT

public export IsNotCloE : Elim d n -> Type
IsNotCloE = So . not . topCloE

||| an elimination which is not a top level closure
public export NotCloElim : Nat -> Nat -> Type
NotCloElim d n = Subset (Elim d n) IsNotCloE

public export %inline
ncloT : (t : Term d n) -> (0 _ : IsNotCloT t) => NotCloTerm d n
ncloT t @{p} = Element t p

public export %inline
ncloE : (t : Elim d n) -> (0 _ : IsNotCloE t) => NotCloElim d n
ncloE e @{p} = Element e p



mutual
  ||| if the input term has any top-level closures, push them under one layer of
  ||| syntax
  export %inline
  pushSubstsT : Term d n -> NotCloTerm d n
  pushSubstsT s = pushSubstsTWith id id s

  ||| if the input elimination has any top-level closures, push them under one
  ||| layer of syntax
  export %inline
  pushSubstsE : Elim d n -> NotCloElim d n
  pushSubstsE e = pushSubstsEWith id id e

  export
  pushSubstsTWith : DSubst dfrom dto -> TSubst dto from to ->
                    Term dfrom from -> NotCloTerm dto to
  pushSubstsTWith th ph (TYPE l) =
    ncloT $ TYPE l
  pushSubstsTWith th ph (Pi qty x a body) =
    ncloT $ Pi qty x (subs a th ph) (subs body th ph)
  pushSubstsTWith th ph (Lam x body) =
    ncloT $ Lam x $ subs body th ph
  pushSubstsTWith th ph (E e) =
    let Element e _ = pushSubstsEWith th ph e in ncloT $ E e
  pushSubstsTWith th ph (CloT  s ps) =
    pushSubstsTWith th (comp' th ps ph) s
  pushSubstsTWith th ph (DCloT s ps) =
    pushSubstsTWith (ps . th) ph s

  export
  pushSubstsEWith : DSubst dfrom dto -> TSubst dto from to ->
                    Elim dfrom from -> NotCloElim dto to
  pushSubstsEWith th ph (F x) =
    ncloE $ F x
  pushSubstsEWith th ph (B i) =
    let res = ph !! i in
    case choose $ topCloE res of
         Left  _ => assert_total pushSubstsE res
         Right _ => ncloE res
  pushSubstsEWith th ph (f :@ s) =
    ncloE $ subs f th ph :@ subs s th ph
  pushSubstsEWith th ph (s :# a) =
    ncloE $ subs s th ph :# subs a th ph
  pushSubstsEWith th ph (CloE  e ps) =
    pushSubstsEWith th (comp' th ps ph) e
  pushSubstsEWith th ph (DCloE e ps) =
    pushSubstsEWith (ps . th) ph e


parameters (th : DSubst dfrom dto) (ph : TSubst dto from to)
  public export %inline
  pushSubstsTWith' : Term dfrom from -> Term dto to
  pushSubstsTWith' s = (pushSubstsTWith th ph s).fst

  public export %inline
  pushSubstsEWith' : Elim dfrom from -> Elim dto to
  pushSubstsEWith' e = (pushSubstsEWith th ph e).fst


public export %inline
pushSubstsT' : Term d n -> Term d n
pushSubstsT' s = (pushSubstsT s).fst

public export %inline
pushSubstsE' : Elim d n -> Elim d n
pushSubstsE' e = (pushSubstsE e).fst


mutual
  -- tightening a term/elim also causes substitutions to be pushed through.
  -- this is because otherwise a variable in an unused part of the subst
  -- would cause it to incorrectly fail

  export covering
  Tighten (Term d) where
    tighten p (TYPE l) =
      pure $ TYPE l
    tighten p (Pi qty x arg res) =
      Pi qty x <$> tighten p arg
               <*> tighten p res
    tighten p (Lam x body) =
      Lam x <$> tighten p body
    tighten p (E e) =
      E <$> tighten p e
    tighten p (CloT tm th) =
      tighten p $ pushSubstsTWith' id th tm
    tighten p (DCloT tm th) =
      tighten p $ pushSubstsTWith' th id tm

  export covering
  Tighten (Elim d) where
    tighten p (F x) =
      pure $ F x
    tighten p (B i) =
      B <$> tighten p i
    tighten p (fun :@ arg) =
      [|tighten p fun :@ tighten p arg|]
    tighten p (tm :# ty) =
      [|tighten p tm :# tighten p ty|]
    tighten p (CloE el th) =
      tighten p $ pushSubstsEWith' id th el
    tighten p (DCloE el th) =
      tighten p $ pushSubstsEWith' th id el

  export covering
  Tighten (ScopeTerm d) where
    tighten p (TUsed   body) = TUsed   <$> tighten (Keep p) body
    tighten p (TUnused body) = TUnused <$> tighten p body


public export %inline
weakT : Term d n -> Term d (S n)
weakT t = t //. shift 1

public export %inline
weakE : Elim d n -> Elim d (S n)
weakE t = t //. shift 1


mutual
  public export
  data IsRedexT : Term d n -> Type where
    IsUpsilonT : IsRedexT $ E (_ :# _)
    IsCloT     : IsRedexT $ CloT {}
    IsDCloT    : IsRedexT $ DCloT {}
    IsERedex   : IsRedexE e -> IsRedexT $ E e

  public export %inline
  NotRedexT : Term d n -> Type
  NotRedexT = Not . IsRedexT

  public export
  data IsRedexE : Elim d n -> Type where
    IsUpsilonE : IsRedexE $ E _ :# _
    IsBetaLam  : IsRedexE $ (Lam {} :# Pi {}) :@ _
    IsCloE     : IsRedexE $ CloE {}
    IsDCloE    : IsRedexE $ DCloE {}

  public export %inline
  NotRedexE : Elim d n -> Type
  NotRedexE = Not . IsRedexE


mutual
  export %inline
  isRedexT : (t : Term d n) -> Dec (IsRedexT t)
  isRedexT (E (tm :# ty)) = Yes IsUpsilonT
  isRedexT (CloT {})      = Yes IsCloT
  isRedexT (DCloT {})     = Yes IsDCloT
  isRedexT (E (CloE  {})) = Yes $ IsERedex IsCloE
  isRedexT (E (DCloE {})) = Yes $ IsERedex IsDCloE
  isRedexT (E e@(_ :@ _)) with (isRedexE e)
    _ | Yes yes = Yes $ IsERedex yes
    _ | No  no  = No \case IsERedex p => no p
  isRedexT (TYPE {}) = No $ \x => case x of {}
  isRedexT (Pi   {}) = No $ \x => case x of {}
  isRedexT (Lam  {}) = No $ \x => case x of {}
  isRedexT (E (F _)) = No $ \x => case x of IsERedex _ impossible
  isRedexT (E (B _)) = No $ \x => case x of IsERedex _ impossible

  export %inline
  isRedexE : (e : Elim d n) -> Dec (IsRedexE e)
  isRedexE (E _ :# _)                    = Yes IsUpsilonE
  isRedexE ((Lam {} :# Pi {}) :@ _)      = Yes IsBetaLam
  isRedexE (CloE {})                     = Yes IsCloE
  isRedexE (DCloE {})                    = Yes IsDCloE
  isRedexE (F x)                         = No $ \x => case x of {}
  isRedexE (B i)                         = No $ \x => case x of {}
  isRedexE (F _                    :@ _) = No $ \x => case x of {}
  isRedexE (B _                    :@ _) = No $ \x => case x of {}
  isRedexE (_ :@ _                 :@ _) = No $ \x => case x of {}
  isRedexE ((TYPE _   :# _)        :@ _) = No $ \x => case x of {}
  isRedexE ((Pi {}    :# _)        :@ _) = No $ \x => case x of {}
  isRedexE ((Lam {}   :# TYPE _)   :@ _) = No $ \x => case x of {}
  isRedexE ((Lam {}   :# Lam {})   :@ _) = No $ \x => case x of {}
  isRedexE ((Lam {}   :# E _)      :@ _) = No $ \x => case x of {}
  isRedexE ((Lam {}   :# CloT {})  :@ _) = No $ \x => case x of {}
  isRedexE ((Lam {}   :# DCloT {}) :@ _) = No $ \x => case x of {}
  isRedexE ((E _      :# _)        :@ _) = No $ \x => case x of {}
  isRedexE ((CloT {}  :# _)        :@ _) = No $ \x => case x of {}
  isRedexE ((DCloT {} :# _)        :@ _) = No $ \x => case x of {}
  isRedexE (CloE {}                :@ _) = No $ \x => case x of {}
  isRedexE (DCloE {}               :@ _) = No $ \x => case x of {}
  isRedexE (TYPE _   :# _)               = No $ \x => case x of {}
  isRedexE (Pi {}    :# _)               = No $ \x => case x of {}
  isRedexE (Lam {}   :# _)               = No $ \x => case x of {}
  isRedexE (CloT {}  :# _)               = No $ \x => case x of {}
  isRedexE (DCloT {} :# _)               = No $ \x => case x of {}


public export %inline
RedexTerm : Nat -> Nat -> Type
RedexTerm d n = Subset (Term d n) IsRedexT

public export %inline
NonRedexTerm : Nat -> Nat -> Type
NonRedexTerm d n = Subset (Term d n) NotRedexT


public export %inline
RedexElim : Nat -> Nat -> Type
RedexElim d n = Subset (Elim d n) IsRedexE

public export %inline
NonRedexElim : Nat -> Nat -> Type
NonRedexElim d n = Subset (Elim d n) NotRedexE


||| substitute a term with annotation for the bound variable of a `ScopeTerm`
export %inline
substScope : (arg, argTy : Term d n) -> (body : ScopeTerm d n) -> Term d n
substScope arg argTy (TUsed   body) = body // one (arg :# argTy)
substScope arg argTy (TUnused body) = body

mutual
  export %inline
  stepT' : (s : Term d n) -> IsRedexT s -> Term d n
  stepT' (E (s :# _)) IsUpsilonT   = s
  stepT' (CloT  s th) IsCloT       = pushSubstsTWith' id th s
  stepT' (DCloT s th) IsDCloT      = pushSubstsTWith' th id s
  stepT' (E e)        (IsERedex p) = E $ stepE' e p

  export %inline
  stepE' : (e : Elim d n) -> IsRedexE e -> Elim d n
  stepE' (E e :# _) IsUpsilonE = e
  stepE' ((Lam {body, _} :# Pi {arg, res, _}) :@ s) IsBetaLam =
    substScope s arg body :# substScope s arg res
  stepE' (CloE  e th) IsCloE  = pushSubstsEWith' id th e
  stepE' (DCloE e th) IsDCloE = pushSubstsEWith' th id e

export %inline
stepT : (s : Term d n) -> Either (NotRedexT s) (Term d n)
stepT s = case isRedexT s of Yes y => Right $ stepT' s y; No n => Left n

export %inline
stepE : (e : Elim d n) -> Either (NotRedexE e) (Elim d n)
stepE e = case isRedexE e of Yes y => Right $ stepE' e y; No n => Left n

export covering
whnfT : Term d n -> NonRedexTerm d n
whnfT s = case stepT s of
               Right s'   => whnfT s'
               Left  done => Element s done

export covering
whnfE : Elim d n -> NonRedexElim d n
whnfE e = case stepE e of
               Right e'   => whnfE e'
               Left  done => Element e done