module Quox.Syntax.Term.Tighten

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

%default total


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
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 s =
    let Element s' _ = pushSubsts s in
    tightenT' p $ assert_smaller s s'

  private
  tightenE : OPE n1 n2 -> Elim d n2 -> Maybe (Elim d n1)
  tightenE p e =
    let Element e' _ = pushSubsts e in
    tightenE' p $ assert_smaller e e'

  private
  tightenT' : OPE n1 n2 -> (t : Term d n2) -> (0 nt : NotClo t) =>
              Maybe (Term d n1)
  tightenT' p (TYPE l loc) = pure $ TYPE l loc
  tightenT' p (IOState loc) = pure $ IOState 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 (Nat n loc) =
    pure $ Nat n loc
  tightenT' p (Succ s loc) =
    Succ <$> tightenT p s <*> pure loc
  tightenT' p (STRING loc) =
    pure $ STRING loc
  tightenT' p (Str s loc) =
    pure $ Str s 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 (Let qty rhs body loc) =
    Let qty <$> assert_total tightenE p rhs <*> tightenS p body <*> pure loc
  tightenT' p (E e) =
    E <$> assert_total tightenE p e

  private
  tightenE' : OPE n1 n2 -> (e : Elim d n2) -> (0 ne : NotClo e) =>
              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 (Fst pair loc) =
    Fst <$> tightenE p pair <*> pure loc
  tightenE' p (Snd pair loc) =
    Snd <$> tightenE p pair <*> 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

  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 s =
    let Element s' _ = pushSubsts s in
    dtightenT' p $ assert_smaller s s'

  export
  dtightenE : OPE d1 d2 -> Elim d2 n -> Maybe (Elim d1 n)
  dtightenE p e =
    let Element e' _ = pushSubsts e in
    dtightenE' p $ assert_smaller e e'

  private
  dtightenT' : OPE d1 d2 -> (t : Term d2 n) -> (0 nt : NotClo t) =>
               Maybe (Term d1 n)
  dtightenT' p (TYPE l loc) =
    pure $ TYPE l loc
  dtightenT' p (IOState loc) =
    pure $ IOState 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 (Nat n loc) =
    pure $ Nat n loc
  dtightenT' p (Succ s loc) =
    Succ <$> dtightenT p s <*> pure loc
  dtightenT' p (STRING loc) =
    pure $ STRING loc
  dtightenT' p (Str s loc) =
    pure $ Str s 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 (Let qty rhs body loc) =
    Let qty <$> assert_total dtightenE p rhs <*> dtightenS p body <*> pure loc
  dtightenT' p (E e) =
    E <$> assert_total dtightenE p e

  export
  dtightenE' : OPE d1 d2 -> (e : Elim d2 n) -> (0 ne : NotClo e) =>
               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 (Fst pair loc) =
    Fst <$> dtightenE p pair <*> pure loc
  dtightenE' p (Snd pair loc) =
    Snd <$> dtightenE p pair <*> 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)|]

  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 Tighten (\d => Term d n) where tighten p t = dtightenT p t
export Tighten (\d => Elim d n) where tighten p e = dtightenE p e


parameters {auto _ : Tighten f} {s : Nat}
  export
  squeeze : Scoped s f n -> (BContext s, Either (f (s + n)) (f n))
  squeeze (S ns (N t)) = (ns, Right t)
  squeeze (S ns (Y t)) = (ns, maybe (Left t) Right $ tightenN s t)

  export
  squeeze' : Scoped s f n -> Scoped s f n
  squeeze' t = let (ns, res) = squeeze t in S ns $ either Y N res

parameters {0 f : Nat -> Nat -> Type}
           {auto tt : Tighten (\d => f d n)} {s : Nat}
  export
  dsqueeze : Scoped s (\d => f d n) d ->
             (BContext s, Either (f (s + d) n) (f d n))
  dsqueeze = squeeze

  export
  dsqueeze' : Scoped s (\d => f d n) d -> Scoped s (\d => f d n) d
  dsqueeze' = squeeze'


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

public export %inline
ST : Tighten f => {s : Nat} -> BContext s -> f (s + n) -> Scoped s f n
ST names body = squeeze' $ SY names body

public export %inline
DST : {s : Nat} -> BContext s -> Term (s + d) n -> DScopeTermN s d n
DST names body = dsqueeze' {f = Term} $ SY names 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


public export %inline
CompH' : (ty : DScopeTerm d n) -> (p, q : Dim d) -> (val : Term d n) ->
         (r : Dim d) -> (zero, one : DScopeTerm d n) -> (loc : Loc) -> Elim d n
CompH' {ty, p, q, val, r, zero, one, loc} =
  let ty' = DST ty.names $ ty.term // (B VZ ty.name.loc ::: shift 2) in
  Comp {
    ty = dsub1 ty q, p, q,
    val = E $ Coe ty p q val val.loc, r,
    zero = DST zero.names $ E $
      Coe ty' (B VZ zero.loc) (weakD 1 q) zero.term zero.loc,
    one = DST one.names $ E $
      Coe ty' (B VZ one.loc) (weakD 1 q) one.term one.loc,
    loc
  }

||| heterogeneous composition, using Comp and Coe (and subst)
|||
||| comp [i ⇒ A] @p @q s @r { 0 j ⇒ t₀; 1 j ⇒ t₁ }
||| ≔
||| comp [A‹q/i›] @p @q (coe [i ⇒ A] @p @q s) @r {
|||   0 j ⇒ coe [i ⇒ A] @j @q t₀;
|||   1 j ⇒ coe [i ⇒ A] @j @q t₁
||| }
public export %inline
CompH : (i : BindName) -> (ty : Term (S d) n) ->
        (p, q : Dim d) -> (val : Term d n) -> (r : Dim d) ->
        (j0 : BindName) -> (zero : Term (S d) n) ->
        (j1 : BindName) -> (one : Term (S d) n) ->
        (loc : Loc) ->
        Elim d n
CompH {i, ty, p, q, val, r, j0, zero, j1, one, loc} =
  CompH' {ty = DST [< i] ty, p, q, val, r,
          zero = DST [< j0] zero, one = DST [< j1] one, loc}