module Quox.Syntax.Term.Reduce

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

%default total


mutual
  public export
  data NotCloT : Term {} -> Type where
    NCTYPE : NotCloT $ TYPE _
    NCPi   : NotCloT $ Pi {}
    NCLam  : NotCloT $ Lam {}
    NCE    : NotCloE e -> NotCloT $ E e

  public export
  data NotCloE : Elim {} -> Type where
    NCF   : NotCloE $ F _
    NCB   : NotCloE $ B _
    NCApp : NotCloE $ _ :@ _
    NCAnn : NotCloE $ _ :# _

mutual
  export
  notCloT : (t : Term {}) -> Dec (NotCloT t)
  notCloT (TYPE _) = Yes NCTYPE
  notCloT (Pi {})  = Yes NCPi
  notCloT (Lam {}) = Yes NCLam
  notCloT (E e) with (notCloE e)
    notCloT (E e) | Yes nc = Yes $ NCE nc
    notCloT (E e) | No  c  = No  $ \case NCE nc => c nc
  notCloT (CloT {})  = No $ \case _ impossible
  notCloT (DCloT {}) = No $ \case _ impossible

  export
  notCloE : (e : Elim {}) -> Dec (NotCloE e)
  notCloE (F _)      = Yes NCF
  notCloE (B _)      = Yes NCB
  notCloE (_ :@ _)   = Yes NCApp
  notCloE (_ :# _)   = Yes NCAnn
  notCloE (CloE  {}) = No $ \case _ impossible
  notCloE (DCloE {}) = No $ \case _ impossible

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

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

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

public export %inline
ncloE : (e : Elim q d n) -> (0 _ : NotCloE e) => NotCloElim q 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 q d n -> NotCloTerm q 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 q d n -> NotCloElim q d n
  pushSubstsE e = pushSubstsEWith id id e

  export
  pushSubstsTWith : DSubst dfrom dto -> TSubst q dto from to ->
                    Term q dfrom from -> NotCloTerm q 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 q dto from to ->
                    Elim q dfrom from -> NotCloElim q dto to
  pushSubstsEWith th ph (F x) =
    ncloE $ F x
  pushSubstsEWith th ph (B i) =
    let res = ph !! i in
    case notCloE res of
         Yes _ => ncloE res
         No  _ => assert_total pushSubstsE 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 q dto from to)
  public export %inline
  pushSubstsTWith' : Term q dfrom from -> Term q dto to
  pushSubstsTWith' s = (pushSubstsTWith th ph s).fst

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


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

public export %inline
pushSubstsE' : Elim q d n -> Elim q 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 q d n -> Term q d (S n)
weakT t = t //. shift 1

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


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

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

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

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


mutual
  export %inline
  isRedexT : (t : Term {}) -> 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 $ \case _ impossible
  isRedexT (Pi   {}) = No $ \case _ impossible
  isRedexT (Lam  {}) = No $ \case _ impossible
  isRedexT (E (F _)) = No $ \case IsERedex _ impossible
  isRedexT (E (B _)) = No $ \case IsERedex _ impossible

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


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

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

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

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


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

mutual
  export %inline
  stepT' : (s : Term q d n) -> IsRedexT s -> Term q 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 q d n) -> IsRedexE e -> Elim q 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 q d n) -> Either (NotRedexT s) (Term q d n)
stepT s = case isRedexT s of Yes y => Right $ stepT' s y; No n => Left n

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

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

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