module Quox.Syntax.Term.Reduce

import Quox.No
import Quox.Syntax.Term.Base
import Quox.Syntax.Term.Subst
import Data.Maybe

%default total


namespace Elim
  public export %inline
  isClo : Elim {} -> Bool
  isClo (CloE  {}) = True
  isClo (DCloE {}) = True
  isClo _          = False

  public export
  0 NotClo : Pred $ Elim {}
  NotClo = No . isClo

namespace Term
  public export %inline
  isClo : Term {} -> Bool
  isClo (CloT  {}) = True
  isClo (DCloT {}) = True
  isClo (E e)      = isClo e
  isClo _          = False

  public export
  0 NotClo : Pred $ Term {}
  NotClo = No . isClo

public export
0 NonCloElim : TermLike
NonCloElim q d n = Subset (Elim q d n) NotClo

public export
0 NonCloTerm : TermLike
NonCloTerm q d n = Subset (Term q d n) NotClo


public export %inline
ncloT : (t : Term q d n) -> (0 nc : NotClo t) => NonCloTerm q d n
ncloT t = Element t nc

public export %inline
ncloE : (e : Elim q d n) -> (0 nc : NotClo e) => NonCloElim q d n
ncloE e = Element e nc


mutual
  namespace Term
    ||| if the input term has any top-level closures, push them under one layer of
    ||| syntax
    export %inline
    pushSubsts : Term q d n -> NonCloTerm q d n
    pushSubsts s = pushSubstsWith id id s

    export
    pushSubstsWith : DSubst dfrom dto -> TSubst q dto from to ->
                      Term q dfrom from -> NonCloTerm q dto to
    pushSubstsWith th ph (TYPE l) =
      ncloT $ TYPE l
    pushSubstsWith th ph (Pi qty x a body) =
      ncloT $ Pi qty x (subs a th ph) (subs body th ph)
    pushSubstsWith th ph (Lam x body) =
      ncloT $ Lam x $ subs body th ph
    pushSubstsWith th ph (Eq i ty l r) =
      ncloT $ Eq i (subs ty th ph) (subs l th ph) (subs r th ph)
    pushSubstsWith th ph (DLam i body) =
      ncloT $ DLam i $ subs body th ph
    pushSubstsWith th ph (E e) =
      let Element e nc = pushSubstsWith th ph e in ncloT $ E e
    pushSubstsWith th ph (CloT s ps) =
      pushSubstsWith th (comp th ps ph) s
    pushSubstsWith th ph (DCloT s ps) =
      pushSubstsWith (ps . th) ph s

  namespace Elim
    ||| if the input elimination has any top-level closures, push them under one
    ||| layer of syntax
    export %inline
    pushSubsts : Elim q d n -> NonCloElim q d n
    pushSubsts e = pushSubstsWith id id e

    export
    pushSubstsWith : DSubst dfrom dto -> TSubst q dto from to ->
                     Elim q dfrom from -> NonCloElim q dto to
    pushSubstsWith th ph (F x) =
      ncloE $ F x
    pushSubstsWith th ph (B i) =
      let res = ph !! i in
      case nchoose $ isClo res of
        Left  yes => assert_total pushSubsts res
        Right no  => Element res no
    pushSubstsWith th ph (f :@ s) =
      ncloE $ subs f th ph :@ subs s th ph
    pushSubstsWith th ph (f :% d) =
      ncloE $ subs f th ph :% (d // th)
    pushSubstsWith th ph (s :# a) =
      ncloE $ subs s th ph :# subs a th ph
    pushSubstsWith th ph (CloE e ps) =
      pushSubstsWith th (comp th ps ph) e
    pushSubstsWith th ph (DCloE e ps) =
      pushSubstsWith (ps . th) ph e


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

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


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


public export 0
Lookup : TermLike
Lookup q d n = Name -> Maybe $ Elim q d n

public export %inline
isLamHead : Elim {} -> Bool
isLamHead (Lam {} :# Pi {}) = True
isLamHead _                 = False

public export %inline
isDLamHead : Elim {} -> Bool
isDLamHead (DLam {} :# Eq {}) = True
isDLamHead _                  = False

public export %inline
isE : Term {} -> Bool
isE (E _) = True
isE _     = False

public export %inline
isAnn : Elim {} -> Bool
isAnn (_ :# _) = True
isAnn _        = False

parameters (g : Lookup q d n)
  mutual
    namespace Elim
      public export
      isRedex : Elim q d n -> Bool
      isRedex (F x)      = isJust $ g x
      isRedex (B _)      = False
      isRedex (f :@ _)   = isRedex f || isLamHead f
      isRedex (f :% _)   = isRedex f || isDLamHead f
      isRedex (t :# a)   = isE t || isRedex t || isRedex a
      isRedex (CloE  {}) = True
      isRedex (DCloE {}) = True

    namespace Term
      public export
      isRedex : Term q d n -> Bool
      isRedex (CloT  {}) = True
      isRedex (DCloT {}) = True
      isRedex (E e)      = isAnn e || isRedex e
      isRedex _          = False

  namespace Elim
    public export
    0 IsRedex, NotRedex : Pred $ Elim q d n
    IsRedex  = So . isRedex
    NotRedex = No . isRedex

  namespace Term
    public export
    0 IsRedex, NotRedex : Pred $ Term q d n
    IsRedex  = So . isRedex
    NotRedex = No . isRedex

public export
0 NonRedexElim, NonRedexTerm : (q, d, n : _) -> Lookup q d n -> Type
NonRedexElim q d n g = Subset (Elim q d n) (NotRedex g)
NonRedexTerm q d n g = Subset (Term q d n) (NotRedex g)


parameters (g : Lookup q d n)
  mutual
    namespace Elim
      export covering
      whnf : Elim q d n -> NonRedexElim q d n g
      whnf (F x) with (g x) proof eq
        _ | Just y  = whnf y
        _ | Nothing = Element (F x) $ rewrite eq in Ah

      whnf (B i) = Element (B i) Ah

      whnf (f :@ s) =
        let Element f fnf = whnf f in
        case nchoose $ isLamHead f of
          Left _ =>
            let Lam {body, _} :# Pi {arg, res, _} = f
                s = s :# arg
            in
            whnf $ sub1 body s :# sub1 res s
          Right nlh => Element (f :@ s) $ fnf `orNo` nlh

      whnf (f :% p) =
        let Element f fnf = whnf f in
        case nchoose $ isDLamHead f of
          Left _ =>
            let DLam {body, _} :# Eq {ty, l, r, _} = f
                body = case p of K e => pick l r e; _ => dsub1 body p
            in
            whnf $ body :# dsub1 ty p
          Right ndlh =>
            Element (f :% p) $ fnf `orNo` ndlh

      whnf (s :# a) =
        let Element s snf = whnf s
            Element a anf = whnf a
        in
        case nchoose $ isE s of
          Left  _  => let E e = s in Element e $ noOr2 snf
          Right ne => Element (s :# a) $ ne `orNo` snf `orNo` anf

      whnf (CloE  el th) = whnf $ pushSubstsWith' id th el
      whnf (DCloE el th) = whnf $ pushSubstsWith' th id el

    namespace Term
      export covering
      whnf : Term q d n -> NonRedexTerm q d n g
      whnf (TYPE l)           = Element (TYPE l)           Ah
      whnf (Pi qty x arg res) = Element (Pi qty x arg res) Ah
      whnf (Lam x body)       = Element (Lam x body)       Ah
      whnf (Eq i ty l r)      = Element (Eq i ty l r)      Ah
      whnf (DLam i body)      = Element (DLam i body)      Ah

      whnf (E e) =
        let Element e enf = whnf e in
        case nchoose $ isAnn e of
          Left  _  => let tm :# _ = e in Element tm $ noOr1 $ noOr2 enf
          Right na => Element (E e) $ na `orNo` enf

      whnf (CloT  tm th) = whnf $ pushSubstsWith' id th tm
      whnf (DCloT tm th) = whnf $ pushSubstsWith' th id tm