module Quox.Syntax.Term.Split

import Quox.Syntax.Term.Base
import Quox.Syntax.Term.Subst
import Quox.Syntax.Term.Tighten
import Quox.Context

import public Quox.No
import public Data.Vect

%default total


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

public export
0 NotLam : Pred $ Term {}
NotLam = No . isLam


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

public export
0 NotDLam : Pred $ Term {}
NotDLam = No . isDLam


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

public export
0 NotPair : Pred $ Term {}
NotPair = No . isPair


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

public export
0 NotApp : Pred $ Elim {}
NotApp = No . isApp


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

public export
0 NotDApp : Pred $ Elim {}
NotDApp = No . isDApp


-- infixl 9 :@@
-- ||| apply multiple arguments at once
-- public export %inline
-- (:@@) : Elim d n -> List (Term d n) -> Elim d n
-- f :@@ ss = foldl app f ss where
--   app : Elim d n -> Term d n -> Elim d n
--   app f s = App f s (f.loc `extend'` s.loc.bounds)

public export
record GetArgs d n where
  constructor GotArgs
  fun      : Elim d n
  args     : List (Term d n)
  0 notApp : NotApp fun

mutual
  export %inline
  getArgs' : Elim d n -> List (Term d n) -> GetArgs d n
  getArgs' fun0 args =
    let Element fun nc = pushSubsts fun0 in
    getArgsNc (assert_smaller fun0 fun) args

  private
  getArgsNc : (e : Elim d n) -> (0 nc : NotClo e) =>
              List (Term d n) -> GetArgs d n
  getArgsNc fun args = case nchoose $ isApp fun of
    Left  y => let App f a _ = fun in getArgs' f (a :: args)
    Right n => GotArgs {fun, args, notApp = n}

||| splits an application into its head and arguments. if it's not an
||| application then the list is just empty.
||| looks through substitutions for applications.
export %inline
getArgs : Elim d n -> GetArgs d n
getArgs e = getArgs' e []


-- infixl 9 :%%
-- ||| apply multiple dimension arguments at once
-- public export %inline
-- (:%%) : Elim d n -> List (Dim d) -> Elim d n
-- f :%% ss = foldl dapp f ss where
--   dapp : Elim d n -> Dim d -> Elim d n
--   dapp f p = DApp f p (f.loc `extend'` p.loc.bounds)

public export
record GetDArgs d n where
  constructor GotDArgs
  fun       : Elim d n
  args      : List (Dim d)
  0 notDApp : NotDApp fun

mutual
  export %inline
  getDArgs' : Elim d n -> List (Dim d) -> GetDArgs d n
  getDArgs' fun0 args =
    let Element fun nc = pushSubsts fun0 in
    getDArgsNc (assert_smaller fun0 fun) args

  private
  getDArgsNc : (e : Elim d n) -> (0 nc : NotClo e) =>
               List (Dim d) -> GetDArgs d n
  getDArgsNc fun args = case nchoose $ isDApp fun of
    Left  y => let DApp f d _ = fun in getDArgs' f (d :: args)
    Right n => GotDArgs {fun, args, notDApp = n}

||| splits a dimension application into its head and arguments. if it's not an
||| d application then the list is just empty
export %inline
getDArgs : Elim d n -> GetDArgs d n
getDArgs e = getDArgs' e []


-- infixr 1 :\\
-- public export
-- absN : BContext m -> Term d (m + n) -> Term d n
-- absN [<]       s = s
-- absN (xs :< x) s = absN xs $ Lam (ST [< x] s) ?absloc

-- public export %inline
-- (:\\) : BContext m -> Term d (m + n) -> Term d n
-- (:\\) = absN


-- infixr 1 :\\%
-- public export
-- dabsN : BContext m -> Term (m + d) n -> Term d n
-- dabsN [<]       s = s
-- dabsN (xs :< x) s = dabsN xs $ DLam (DST [< x] s) ?dabsLoc

-- public export %inline
-- (:\\%) : BContext m -> Term (m + d) n -> Term d n
-- (:\\%) = dabsN


public export
record GetLams d n where
  constructor GotLams
  {0 lams, rest : Nat}
  names    : BContext lams
  body     : Term d rest
  0 eq     : lams + n = rest
  0 notLam : NotLam body

mutual
  export %inline
  getLams' : forall lams, rest.
             BContext lams -> Term d rest -> (0 eq : lams + n = rest) ->
             GetLams d n
  getLams' xs s0 eq =
    let Element s nc = pushSubsts s0 in
    getLamsNc xs (assert_smaller s0 s) eq

  private
  getLamsNc : forall lams, rest.
              BContext lams ->
              (t : Term d rest) -> (0 nc : NotClo t) =>
              (0 eq : lams + n = rest) ->
              GetLams d n
  getLamsNc xs s Refl = case nchoose $ isLam s of
    Left  y => let Lam (S [< x] body) _ = s in
               getLams' (xs :< x) (assert_smaller s body.term) Refl
    Right n => GotLams xs s Refl n

export %inline
getLams : Term d n -> GetLams d n
getLams s = getLams' [<] s Refl


public export
record GetDLams d n where
  constructor GotDLams
  {0 lams, rest : Nat}
  names     : BContext lams
  body      : Term rest n
  0 eq      : lams + d = rest
  0 notDLam : NotDLam body

mutual
  export %inline
  getDLams' : forall lams, rest.
              BContext lams -> Term rest n -> (0 eq : lams + d = rest) ->
              GetDLams d n
  getDLams' xs s0 eq =
    let Element s nc = pushSubsts s0 in
    getDLamsNc xs (assert_smaller s0 s) eq

  private
  getDLamsNc : forall lams, rest.
               BContext lams ->
               (t : Term rest n) -> (0 nc : NotClo t) =>
               (0 eq : lams + d = rest) ->
               GetDLams d n
  getDLamsNc is s Refl = case nchoose $ isDLam s of
    Left  y => let DLam (S [< i] body) _ = s in
               getDLams' (is :< i) (assert_smaller s body.term) Refl
    Right n => GotDLams is s Refl n

export %inline
getDLams : Term d n -> GetDLams d n
getDLams s = getDLams' [<] s Refl


public export
record GetPairs d n where
  constructor GotPairs
  init : SnocList $ Term d n
  last : Term d n
  notPair : NotPair last

mutual
  export %inline
  getPairs' : SnocList (Term d n) -> Term d n -> GetPairs d n
  getPairs' ss t0 =
    let Element t nc = pushSubsts t0 in getPairsNc ss (assert_smaller t0 t)

  private
  getPairsNc : SnocList (Term d n) ->
               (t : Term d n) -> (0 nc : NotClo t) =>
               GetPairs d n
  getPairsNc ss t = case nchoose $ isPair t of
    Left  y => let Pair s t _ = t in
               getPairs' (ss :< s) t
    Right n => GotPairs ss t n

export
getPairs : Term d n -> GetPairs d n
getPairs = getPairs' [<]