module Quox.Syntax.Term.Tighten

import Quox.Syntax.Term.Base
import Quox.Syntax.Subst
import public Quox.OPE

%default total


export
Tighten (Shift from) where
  -- `OPE m n` is a spicy `m ≤ n`,
  -- and `Shift from n` is a (different) spicy `from ≤ n`
  -- so the value is `from ≤ m` (as a `Shift`), if that is the case
  tighten _        SZ      = Nothing
  tighten Id       by      = Just by
  tighten (Drop p) (SS by) = tighten p by
  tighten (Keep p) (SS by) = [|SS $ tighten p by|]


export
Tighten Dim where
  tighten p (K e loc) = pure $ K e loc
  tighten p (B i loc) = B <$> tighten p i <*> pure loc


export
tightenSub : (forall m, n. OPE m n -> env n -> Maybe (env m)) ->
             OPE to1 to2 -> Subst env from to2 -> Maybe (Subst env from to1)
tightenSub f p (Shift by) = [|Shift $ tighten p by|]
tightenSub f p (t ::: th) = [|f p t !::: tightenSub f p th|]

export
Tighten env => Tighten (Subst env from) where
  tighten p th = tightenSub tighten p th


export
tightenScope : (forall m, n. OPE m n -> f n -> Maybe (f m)) ->
               {s : Nat} -> OPE m n -> Scoped s f n -> Maybe (Scoped s f m)
tightenScope f p (S names (Y body)) = SY names <$> f (keepN s p) body
tightenScope f p (S names (N body)) = S names . N <$> f p body

export
tightenDScope : {0 f : Nat -> Nat -> Type} ->
                (forall m, n, k. OPE m n -> f n k -> Maybe (f m k)) ->
                OPE m n -> Scoped s (f n) k -> Maybe (Scoped s (f m) k)
tightenDScope f p (S names (Y body)) = SY names <$> f p body
tightenDScope f p (S names (N body)) = S names . N <$> f p body


mutual
  private
  tightenT : OPE n1 n2 -> Term d n2 -> Maybe (Term d n1)
  tightenT p (TYPE l loc) = pure $ TYPE l loc
  tightenT p (Pi qty arg res loc) =
    Pi qty <$> tightenT p arg <*> tightenS p res <*> pure loc
  tightenT p (Lam body loc) =
    Lam <$> tightenS p body <*> pure loc
  tightenT p (Sig fst snd loc) =
    Sig <$> tightenT p fst <*> tightenS p snd <*> pure loc
  tightenT p (Pair fst snd loc) =
    Pair <$> tightenT p fst <*> tightenT p snd <*> pure loc
  tightenT p (Enum cases loc) =
    pure $ Enum cases loc
  tightenT p (Tag tag loc) =
    pure $ Tag tag loc
  tightenT p (Eq ty l r loc) =
    Eq <$> tightenDS p ty <*> tightenT p l <*> tightenT p r <*> pure loc
  tightenT p (DLam body loc) =
    DLam <$> tightenDS p body <*> pure loc
  tightenT p (Nat loc) =
    pure $ Nat loc
  tightenT p (Zero loc) =
    pure $ Zero loc
  tightenT p (Succ s loc) =
    Succ <$> tightenT p s <*> pure loc
  tightenT p (BOX qty ty loc) =
    BOX qty <$> tightenT p ty <*> pure loc
  tightenT p (Box val loc) =
    Box <$> tightenT p val <*> pure loc
  tightenT p (E e) =
    assert_total $ E <$> tightenE p e
  tightenT p (CloT (Sub tm th)) = do
    th <- assert_total $ tightenSub tightenE p th
    pure $ CloT $ Sub tm th
  tightenT p (DCloT (Sub tm th)) = do
    tm <- tightenT p tm
    pure $ DCloT $ Sub tm th

  private
  tightenE : OPE n1 n2 -> Elim d n2 -> Maybe (Elim d n1)
  tightenE p (F x u loc) =
    pure $ F x u loc
  tightenE p (B i loc) =
    B <$> tighten p i <*> pure loc
  tightenE p (App fun arg loc) =
    App <$> tightenE p fun <*> tightenT p arg <*> pure loc
  tightenE p (CasePair qty pair ret body loc) =
    CasePair qty <$> tightenE p pair
                 <*> tightenS p ret
                 <*> tightenS p body
                 <*> pure loc
  tightenE p (CaseEnum qty tag ret arms loc) =
    CaseEnum qty <$> tightenE p tag
                 <*> tightenS p ret
                 <*> traverse (tightenT p) arms
                 <*> pure loc
  tightenE p (CaseNat qty qtyIH nat ret zero succ loc) =
    CaseNat qty qtyIH
      <$> tightenE p nat
      <*> tightenS p ret
      <*> tightenT p zero
      <*> tightenS p succ
      <*> pure loc
  tightenE p (CaseBox qty box ret body loc) =
    CaseBox qty <$> tightenE p box
                <*> tightenS p ret
                <*> tightenS p body
                <*> pure loc
  tightenE p (DApp fun arg loc) =
    DApp <$> tightenE p fun <*> pure arg <*> pure loc
  tightenE p (Ann tm ty loc) =
    Ann <$> tightenT p tm <*> tightenT p ty <*> pure loc
  tightenE p (Coe ty q0 q1 val loc) =
    Coe <$> tightenDS p ty
        <*> pure q0 <*> pure q1
        <*> tightenT p val
        <*> pure loc
  tightenE p (Comp ty q0 q1 val r zero one loc) =
    Comp <$> tightenT p ty
         <*> pure q0 <*> pure q1
         <*> tightenT p val
         <*> pure r
         <*> tightenDS p zero
         <*> tightenDS p one
         <*> pure loc
  tightenE p (TypeCase ty ret arms def loc) =
    TypeCase <$> tightenE p ty
             <*> tightenT p ret
             <*> traverse (tightenS p) arms
             <*> tightenT p def
             <*> pure loc
  tightenE p (CloE (Sub el th)) = do
    th <- assert_total $ tightenSub tightenE p th
    pure $ CloE $ Sub el th
  tightenE p (DCloE (Sub el th)) = do
    el <- tightenE p el
    pure $ DCloE $ Sub el th

  export
  tightenS : {s : Nat} -> OPE m n ->
             ScopeTermN s f n -> Maybe (ScopeTermN s f m)
  tightenS = assert_total $ tightenScope tightenT

  export
  tightenDS : OPE m n -> DScopeTermN s f n -> Maybe (DScopeTermN s f m)
  tightenDS = assert_total $ tightenDScope tightenT {f = \n, d => Term d n}

export Tighten (Elim d) where tighten p e = tightenE p e
export Tighten (Term d) where tighten p t = tightenT p t


mutual
  export
  dtightenT : OPE d1 d2 -> Term d2 n -> Maybe (Term d1 n)
  dtightenT p (TYPE l loc) =
    pure $ TYPE l loc
  dtightenT p (Pi qty arg res loc) =
    Pi qty <$> dtightenT p arg <*> dtightenS p res <*> pure loc
  dtightenT p (Lam body loc) =
    Lam <$> dtightenS p body <*> pure loc
  dtightenT p (Sig fst snd loc) =
    Sig <$> dtightenT p fst <*> dtightenS p snd <*> pure loc
  dtightenT p (Pair fst snd loc) =
    Pair <$> dtightenT p fst <*> dtightenT p snd <*> pure loc
  dtightenT p (Enum cases loc) =
    pure $ Enum cases loc
  dtightenT p (Tag tag loc) =
    pure $ Tag tag loc
  dtightenT p (Eq ty l r loc) =
    Eq <$> dtightenDS p ty <*> dtightenT p l <*> dtightenT p r <*> pure loc
  dtightenT p (DLam body loc) =
    DLam <$> dtightenDS p body <*> pure loc
  dtightenT p (Nat loc) =
    pure $ Nat loc
  dtightenT p (Zero loc) =
    pure $ Zero loc
  dtightenT p (Succ s loc) =
    Succ <$> dtightenT p s <*> pure loc
  dtightenT p (BOX qty ty loc) =
    BOX qty <$> dtightenT p ty <*> pure loc
  dtightenT p (Box val loc) =
    Box <$> dtightenT p val <*> pure loc
  dtightenT p (E e) =
    assert_total $ E <$> dtightenE p e
  dtightenT p (CloT (Sub tm th)) = do
    tm <- dtightenT p tm
    th <- assert_total $ traverse (dtightenE p) th
    pure $ CloT $ Sub tm th
  dtightenT p (DCloT (Sub tm th)) = do
    th <- tighten p th
    pure $ DCloT $ Sub tm th

  export
  dtightenE : OPE d1 d2 -> Elim d2 n -> Maybe (Elim d1 n)
  dtightenE p (F x u loc) =
    pure $ F x u loc
  dtightenE p (B i loc) =
    pure $ B i loc
  dtightenE p (App fun arg loc) =
    App <$> dtightenE p fun <*> dtightenT p arg <*> pure loc
  dtightenE p (CasePair qty pair ret body loc) =
    CasePair qty <$> dtightenE p pair
                 <*> dtightenS p ret
                 <*> dtightenS p body
                 <*> pure loc
  dtightenE p (CaseEnum qty tag ret arms loc) =
    CaseEnum qty <$> dtightenE p tag
                 <*> dtightenS p ret
                 <*> traverse (dtightenT p) arms
                 <*> pure loc
  dtightenE p (CaseNat qty qtyIH nat ret zero succ loc) =
    CaseNat qty qtyIH
      <$> dtightenE p nat
      <*> dtightenS p ret
      <*> dtightenT p zero
      <*> dtightenS p succ
      <*> pure loc
  dtightenE p (CaseBox qty box ret body loc) =
    CaseBox qty <$> dtightenE p box
                <*> dtightenS p ret
                <*> dtightenS p body
                <*> pure loc
  dtightenE p (DApp fun arg loc) =
    DApp <$> dtightenE p fun <*> tighten p arg <*> pure loc
  dtightenE p (Ann tm ty loc) =
    Ann <$> dtightenT p tm <*> dtightenT p ty <*> pure loc
  dtightenE p (Coe ty q0 q1 val loc) =
    [|Coe (dtightenDS p ty) (tighten p q0) (tighten p q1) (dtightenT p val)
          (pure loc)|]
  dtightenE p (Comp ty q0 q1 val r zero one loc) =
    [|Comp (dtightenT p ty) (tighten p q0) (tighten p q1)
           (dtightenT p val) (tighten p r)
           (dtightenDS p zero) (dtightenDS p one) (pure loc)|]
  dtightenE p (TypeCase ty ret arms def loc) =
    [|TypeCase (dtightenE p ty) (dtightenT p ret)
               (traverse (dtightenS p) arms) (dtightenT p def) (pure loc)|]
  dtightenE p (CloE (Sub el th)) = do
    el <- dtightenE p el
    th <- assert_total $ traverse (dtightenE p) th
    pure $ CloE $ Sub el th
  dtightenE p (DCloE (Sub el th)) = do
    th <- tighten p th
    pure $ DCloE $ Sub el th

  export
  dtightenS : OPE d1 d2 -> ScopeTermN s d2 n -> Maybe (ScopeTermN s d1 n)
  dtightenS = assert_total $ tightenDScope dtightenT {f = Term}

  export
  dtightenDS : {s : Nat} -> OPE d1 d2 ->
               DScopeTermN s d2 n -> Maybe (DScopeTermN s d1 n)
  dtightenDS = assert_total $ tightenScope dtightenT


export [TermD] Tighten (\d => Term d n) where tighten p t = dtightenT p t
export [ElimD] Tighten (\d => Elim d n) where tighten p e = dtightenE p e


-- versions of SY, etc, that try to tighten and use SN automatically

public export
ST : Tighten f => {s : Nat} -> BContext s -> f (s + n) -> Scoped s f n
ST names body =
  case tightenN s body of
    Just body => S names $ N body
    Nothing   => S names $ Y body

public export
DST : {s : Nat} -> BContext s -> Term (s + d) n -> DScopeTermN s d n
DST names body =
  case tightenN @{TermD} s body of
    Just body => S names $ N body
    Nothing   => S names $ Y body

public export %inline
PiT : (qty : Qty) -> (x : BindName) ->
      (arg : Term d n) -> (res : Term d (S n)) -> (loc : Loc) -> Term d n
PiT {qty, x, arg, res, loc} = Pi {qty, arg, res = ST [< x] res, loc}

public export %inline
LamT : (x : BindName) -> (body : Term d (S n)) -> (loc : Loc) -> Term d n
LamT {x, body, loc} = Lam {body = ST [< x] body, loc}

public export %inline
SigT : (x : BindName) -> (fst : Term d n) ->
       (snd : Term d (S n)) -> (loc : Loc) -> Term d n
SigT {x, fst, snd, loc} = Sig {fst, snd = ST [< x] snd, loc}

public export %inline
EqT : (i : BindName) -> (ty : Term (S d) n) ->
      (l, r : Term d n) -> (loc : Loc) -> Term d n
EqT {i, ty, l, r, loc} = Eq {ty = DST [< i] ty, l, r, loc}

public export %inline
DLamT : (i : BindName) -> (body : Term (S d) n) -> (loc : Loc) -> Term d n
DLamT {i, body, loc} = DLam {body = DST [< i] body, loc}

public export %inline
CoeT : (i : BindName) -> (ty : Term (S d) n) ->
       (p, q : Dim d) -> (val : Term d n) -> (loc : Loc) -> Elim d n
CoeT {i, ty, p, q, val, loc} = Coe {ty = DST [< i] ty, p, q, val, loc}

public export %inline
typeCase1T : Elim d n -> Term d n ->
             (k : TyConKind) -> BContext (arity k) -> Term d (arity k + n) ->
             (loc : Loc) ->
             {default (Nat loc) def : Term d n} ->
             Elim d n
typeCase1T ty ret k ns body loc {def} =
  typeCase ty ret [(k ** ST ns body)] def loc


export
squeeze : {s : Nat} -> ScopeTermN s d n -> ScopeTermN s d n
squeeze (S names (Y body)) = S names $ maybe (Y body) N $ tightenN s body
squeeze (S names (N body)) = S names $ N body

export
dsqueeze : {s : Nat} -> DScopeTermN s d n -> DScopeTermN s d n
dsqueeze (S names (Y body)) =
  S names $ maybe (Y body) N $ tightenN s body @{TermD}
dsqueeze (S names (N body)) = S names $ N body