module Quox.Reduce

import Quox.No
import Quox.Syntax
import Quox.Definition
import Data.Vect
import Data.Maybe

%default total



public export
0 CloTest : TermLike -> Type
CloTest tm = forall q, d, n. tm q d n -> Bool

interface PushSubsts (0 tm : TermLike) (0 isClo : CloTest tm) | tm where
  pushSubstsWith : DSubst dfrom dto -> TSubst q dto from to ->
                   tm q dfrom from -> Subset (tm q dto to) (No . isClo)

public export
0 NotClo : {isClo : CloTest tm} -> PushSubsts tm isClo => Pred (tm q d n)
NotClo = No . isClo

public export
0 NonClo : (tm : TermLike) -> {isClo : CloTest tm} ->
           PushSubsts tm isClo => TermLike
NonClo tm q d n = Subset (tm q d n) NotClo

public export %inline
nclo : {isClo : CloTest tm} -> (0 _ : PushSubsts tm isClo) =>
       (t : tm q d n) -> (0 nc : NotClo t) => NonClo tm q d n
nclo t = Element t nc

parameters {0 isClo : CloTest tm} {auto _ : PushSubsts tm isClo}
  ||| if the input term has any top-level closures, push them under one layer of
  ||| syntax
  export %inline
  pushSubsts : tm q d n -> NonClo tm q d n
  pushSubsts s = pushSubstsWith id id s

  export %inline
  pushSubstsWith' : DSubst dfrom dto -> TSubst q dto from to ->
                    tm q dfrom from -> tm q dto to
  pushSubstsWith' th ph x = fst $ pushSubstsWith th ph x

mutual
  public export
  isCloT : CloTest Term
  isCloT (CloT  {}) = True
  isCloT (DCloT {}) = True
  isCloT (E e)      = isCloE e
  isCloT _          = False

  public export
  isCloE : CloTest Elim
  isCloE (CloE  {}) = True
  isCloE (DCloE {}) = True
  isCloE _          = False

mutual
  export
  PushSubsts Term Reduce.isCloT where
    pushSubstsWith th ph (TYPE l) =
      nclo $ TYPE l
    pushSubstsWith th ph (Pi qty a body) =
      nclo $ Pi qty (a // th // ph) (body // th // ph)
    pushSubstsWith th ph (Lam body) =
      nclo $ Lam (body // th // ph)
    pushSubstsWith th ph (Sig a b) =
      nclo $ Sig (a // th // ph) (b // th // ph)
    pushSubstsWith th ph (Pair s t) =
      nclo $ Pair (s // th // ph) (t // th // ph)
    pushSubstsWith th ph (Eq ty l r) =
      nclo $ Eq (ty // th // ph) (l // th // ph) (r // th // ph)
    pushSubstsWith th ph (DLam body) =
      nclo $ DLam (body // th // ph)
    pushSubstsWith th ph (E e) =
      let Element e nc = pushSubstsWith th ph e in nclo $ 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

  export
  PushSubsts Elim Reduce.isCloE where
    pushSubstsWith th ph (F x) =
      nclo $ F x
    pushSubstsWith th ph (B i) =
      let res = ph !! i in
      case nchoose $ isCloE res of
        Left  yes => assert_total pushSubsts res
        Right no  => Element res no
    pushSubstsWith th ph (f :@ s) =
      nclo $ (f // th // ph) :@ (s // th // ph)
    pushSubstsWith th ph (CasePair pi p r b) =
      nclo $ CasePair pi (p // th // ph) (r // th // ph) (b // th // ph)
    pushSubstsWith th ph (f :% d) =
      nclo $ (f // th // ph) :% (d // th)
    pushSubstsWith th ph (s :# a) =
      nclo $ (s // th // ph) :# (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



public export
0 RedexTest : TermLike -> Type
RedexTest tm = forall q, d, n, g. Definitions' q g -> tm q d n -> Bool

public export
interface Whnf (0 tm : TermLike) (0 isRedex : RedexTest tm) | tm where
  whnf : (defs : Definitions' q g) ->
         tm q d n -> Subset (tm q d n) (No . isRedex defs)

public export
0 IsRedex, NotRedex : {isRedex : RedexTest tm} -> Whnf tm isRedex =>
                      Definitions' q g -> Pred (tm q d n)
IsRedex  defs = So . isRedex defs
NotRedex defs = No . isRedex defs

public export
0 NonRedex : (tm : TermLike) -> {isRedex : RedexTest tm} -> Whnf tm isRedex =>
             (q : Type) -> (d, n : Nat) -> {g : _} ->
             (defs : Definitions' q g) -> Type
NonRedex tm q d n defs = Subset (tm q d n) (NotRedex defs)

public export %inline
nred : {0 isRedex : RedexTest tm} -> (0 _ : Whnf tm isRedex) =>
       (t : tm q d n) -> (0 nr : NotRedex defs t) =>
       NonRedex tm q d n defs
nred t = Element t nr


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
isPairHead : Elim {} -> Bool
isPairHead (Pair {} :# Sig {}) = True
isPairHead _                   = False

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

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


mutual
  public export
  isRedexE : RedexTest Elim
  isRedexE defs (F x) {d, n} =
    isJust $ lookupElim x defs {d, n}
  isRedexE _ (B _) = False
  isRedexE defs (f :@ _) =
    isRedexE defs f || isLamHead f
  isRedexE defs (CasePair {pair, _}) =
    isRedexE defs pair || isPairHead pair
  isRedexE defs (f :% _) =
    isRedexE defs f || isDLamHead f
  isRedexE defs (t :# a) =
    isE t || isRedexT defs t || isRedexT defs a
  isRedexE _ (CloE  {}) = True
  isRedexE _ (DCloE {}) = True

  public export
  isRedexT : RedexTest Term
  isRedexT _    (CloT  {}) = True
  isRedexT _    (DCloT {}) = True
  isRedexT defs (E e)      = isAnn e || isRedexE defs e
  isRedexT _    _          = False

mutual
  export covering
  Whnf Elim Reduce.isRedexE where
    whnf defs (F x) with (lookupElim x defs) proof eq
      _ | Just y  = whnf defs y
      _ | Nothing = Element (F x) $ rewrite eq in Ah

    whnf _ (B i) = nred $ B i

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

    whnf defs (CasePair pi pair ret body) =
      let Element pair pairnf = whnf defs pair in
      case nchoose $ isPairHead pair of
        Left _ =>
          let Pair {fst, snd} :# Sig {fst = tfst, snd = tsnd, _} = pair
              fst = fst :# tfst
              snd = snd :# sub1 tsnd fst
          in
          whnf defs $ subN body [fst, snd] :# sub1 ret pair
        Right np =>
          Element (CasePair pi pair r ret x y body) $ pairnf `orNo` np

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

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

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

  export covering
  Whnf Term Reduce.isRedexT where
    whnf _ t@(TYPE {}) = nred t
    whnf _ t@(Pi   {}) = nred t
    whnf _ t@(Lam  {}) = nred t
    whnf _ t@(Sig  {}) = nred t
    whnf _ t@(Pair {}) = nred t
    whnf _ t@(Eq   {}) = nred t
    whnf _ t@(DLam {}) = nred t

    whnf defs (E e) =
      let Element e enf = whnf defs 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 defs (CloT  tm th) = whnf defs $ pushSubstsWith' id th tm
    whnf defs (DCloT tm th) = whnf defs $ pushSubstsWith' th id tm