module linear.Typing.Substitution where

open import Data.Nat
open import Data.Fin
open import Data.Vec hiding (map ; [_] ; _++_)
open import Data.Product hiding (swap)
open import Function
open import Relation.Binary.PropositionalEquality hiding ([_])

open import linear.Scope as Sc hiding (Weakening ; weakFin ; Substituting ; substFin ; copys ; Env ; withFreshVars)
open import linear.Type
open import linear.Context as C hiding (copys ; _++_)
open import linear.Language as L hiding (weakInfer ; weakCheck ; Env ; substInfer ; substCheck)
open import linear.Usage
open import linear.Usage.Functional
open import linear.Usage.Consumption hiding (_++_)
open import linear.Typing
open import linear.Typing.Functional
open import linear.Typing.Consumption

mutual

  weakInfer : Weakening Infer L.weakInfer TInfer
  weakInfer π“œ (`var k)                     = `var (weakFin π“œ k)
  weakInfer π“œ (`app t u)                   = `app (weakInfer π“œ t) (weakCheck π“œ u)
  weakInfer π“œ (`fst t)                     = `fst (weakInfer π“œ t)
  weakInfer π“œ (`snd t)                     = `snd (weakInfer π“œ t)
  weakInfer π“œ (`case t return Οƒ of l %% r) = `case weakInfer π“œ t return Οƒ
                                                of weakCheck (copy π“œ) l
                                                %% weakCheck (copy π“œ) r
  weakInfer π“œ (`exfalso Οƒ t)               = `exfalso Οƒ (weakInfer π“œ t)
  weakInfer π“œ (`cut t)                     = `cut (weakCheck π“œ t)

  weakCheck : Weakening Check L.weakCheck TCheck
  weakCheck π“œ (`lam t)            = `lam weakCheck (copy π“œ) t
  weakCheck {m = m} π“œ (`let_∷=_`in_ {Οƒ} {Ο„} {o} {rp} {Ξ΄} {rt} {Ξ”} {ΞΈ} {ru} p t u) =
    let P    = Ξ» {Ξ³} (Ξ“ Ξ“β€² : Usages Ξ³) β†’ Ξ“ ⊒ Ο„ βˆ‹ L.weakCheck (Sc.copys o m) ru ⊠ Ξ“β€²
        ih   = weakCheck (copys o π“œ) u
        cast = ++copys-elimβ‚‚ P [[ Ξ΄ ]] ]] Ξ΄ [[ Ξ” ΞΈ π“œ
    in `let p ∷= weakInfer π“œ t `in cast ih
  weakCheck π“œ `unit               = `unit
  weakCheck π“œ (`prdβŠ— t u)         = `prdβŠ— (weakCheck π“œ t) (weakCheck π“œ u)
  weakCheck π“œ (`prd& t u)         = `prd& (weakCheck π“œ t) (weakCheck π“œ u)
  weakCheck π“œ (`inl t)            = `inl weakCheck π“œ t
  weakCheck π“œ (`inr t)            = `inr weakCheck π“œ t
  weakCheck π“œ (`neu t)            = `neu weakInfer π“œ t

substFin : 
  {k l : β„•} {Ξ³ : Context k} {Ξ“ Ξ” : Usages Ξ³} {Οƒ : Type} {v : Fin k} {ρ : Sc.Env Infer k l}
  {ΞΈ : Context l} {΀₁ Ξ€β‚‚ : Usages ΞΈ} β†’
  Env TInfer ΀₁ ρ Ξ€β‚‚ Ξ“ β†’ Ξ“ ⊒ v ∈[ Οƒ ]⊠ Ξ” β†’
  βˆƒ Ξ» ΀₃ β†’ ΀₁ ⊒ Sc.substFin L.fresheyInfer ρ v ∈ Οƒ ⊠ ΀₃ Γ— Env TInfer ΀₃ ρ Ξ€β‚‚ Ξ”
substFin (t ∷ ρ)  z     = , t , β”€βˆ· ρ
substFin ([v]∷ ρ) z     = , `var z , ]v[∷ ρ
substFin (T ∷ ρ)  (s v) =
  let (ΞΈ , val , ρ′) = substFin ρ v
      (_ , c₁ , cβ‚‚)  = swap (consumptionInfer T) (consumptionInfer val)
  in , framingInfer cβ‚‚ val , framingInfer c₁ T ∷ ρ′
substFin ([v]∷ ρ) (s v) = map ([ _ ] ∷_) (map (weakInfer (insert _ finish)) [v]∷_) $ substFin ρ v
substFin (]v[∷ ρ) (s v) = map (] _ [ ∷_) (map (weakInfer (insert _ finish)) ]v[∷_) $ substFin ρ v
substFin (β”€βˆ· ρ)   (s v) = map id (map id β”€βˆ·_) $ substFin ρ v

substLam :
  {k l : β„•} {Ξ³ : Context k} {Ξ” : Usages Ξ³} {ρ : Sc.Env Infer k l}
  {ΞΈ : Context l} {΀₁ Ξ€β‚‚ : Usages ΞΈ} {Οƒ Ο„ : Type} {b : Check (suc k)} β†’
  Ξ£[ T₃ ∈ Usages (Οƒ ∷ ΞΈ) ] [ Οƒ ] ∷ ΀₁ ⊒ Ο„ βˆ‹ L.substCheck (v∷ ρ) b ⊠ T₃
                         Γ— Env TInfer T₃ (v∷ ρ) (] Οƒ [ ∷ Ξ€β‚‚) (] Οƒ [ ∷ Ξ”) β†’
  Ξ£[ ΀₃ ∈ Usages ΞΈ       ] ΀₁ ⊒ Οƒ ─o Ο„ βˆ‹ L.substCheck ρ (`lam b) ⊠ ΀₃
                         Γ— Env TInfer ΀₃ ρ Ξ€β‚‚ Ξ”
substLam (._ , bρ , ]v[∷ ρ′) = , `lam bρ , ρ′

substCase :
  {k l : β„•} {Ξ³ : Context k} {Ξ” : Usages Ξ³} {ρ : Sc.Env Infer k l}
  {ΞΈ : Context l} {΀₁ Ξ€β‚‚ Ξ€β‚„ : Usages ΞΈ} {σ₁ Οƒβ‚‚ Ο„ : Type} (t : Infer k) {l r : Check (suc k)} β†’
  ΀₁ ⊒ L.substInfer ρ t ∈ σ₁ βŠ• Οƒβ‚‚ ⊠ Ξ€β‚‚ β†’
  Ξ£[ T₃ ∈ Usages (σ₁ ∷ ΞΈ) ] [ σ₁ ] ∷ Ξ€β‚‚ ⊒ Ο„ βˆ‹ L.substCheck (v∷ ρ) l ⊠ T₃
                         Γ— Env TInfer T₃ (v∷ ρ) (] σ₁ [ ∷ Ξ€β‚„) (] σ₁ [ ∷ Ξ”) β†’
  Ξ£[ T₃ ∈ Usages (Οƒβ‚‚ ∷ ΞΈ) ] [ Οƒβ‚‚ ] ∷ Ξ€β‚‚ ⊒ Ο„ βˆ‹ L.substCheck (v∷ ρ) r ⊠ T₃
                         Γ— Env TInfer T₃ (v∷ ρ) (] Οƒβ‚‚ [ ∷ Ξ€β‚„) (] Οƒβ‚‚ [ ∷ Ξ”) β†’
  Ξ£[ ΀₃ ∈ Usages ΞΈ       ] ΀₁ ⊒ L.substInfer ρ (`case t return Ο„ of l %% r) ∈ Ο„ ⊠ ΀₃
                         Γ— Env TInfer ΀₃ ρ Ξ€β‚„ Ξ”
substCase t tρ (._ , lρ , (]v[∷ ρ₁′)) (._ , rρ , (]v[∷ ρ₂′))
  rewrite sym (functionalEnvPre functionalInferPre _ ρ₁′ ρ₂′) =
  , `case tρ return _ of lρ %% rρ , ρ₁′

-- idea: generalise with a function f applied to each side!
substLet :
  {k l o : β„•} {Ξ³ : Context k} {Ξ” : Usages Ξ³} {ρ : Sc.Env Infer k l} (Ξ΄ : Context o)
  {ΞΈ : Context l} {΀₃ : Usages ΞΈ} β†’
  Ξ£[ Tβ‚‚ ∈ Usages (Ξ΄ C.++ ΞΈ) ] Env TInfer Tβ‚‚ (Sc.withFreshVars o ρ) (]] Ξ΄ [[ ++ ΀₃) (]] Ξ΄ [[ ++ Ξ”) β†’
  Ξ£[ Ξ€β‚‚ ∈ Usages ΞΈ ] Env TInfer Ξ€β‚‚ ρ ΀₃ Ξ”
substLet []      (Ξ€β‚‚ , ρ′) = , ρ′
substLet (a ∷ Ξ΄) (._ , (]v[∷ ρ′)) = substLet Ξ΄ (, ρ′)


mutual

  substInfer : Substituting Infer Infer L.substInfer TInfer TInfer
  substInfer ρ (`var x)                     = substFin ρ x
  substInfer ρ (`app t u)                   =
    let (θ₁ , tρ , ρ₁) = substInfer ρ t
        (ΞΈβ‚‚ , uρ , ρ₂) = substCheck ρ₁ u
    in ΞΈβ‚‚ , `app tρ uρ , ρ₂
  substInfer ρ (`fst t)                     =
    let (θ₁ , tρ , ρ₁) = substInfer ρ t
    in θ₁ , `fst tρ , ρ₁
  substInfer ρ (`snd t)                     =
    let (θ₁ , tρ , ρ₁) = substInfer ρ t
    in θ₁ , `snd tρ , ρ₁
  substInfer {t = `case rt return .Οƒ of rl %% rr} ρ (`case t return Οƒ of l %% r) =
    let (θ₁ , tρ , ρ₁) = substInfer ρ t
    in substCase rt tρ (substCheck ([v]∷ ρ₁) l) (substCheck ([v]∷ ρ₁) r)
  substInfer ρ (`exfalso Οƒ t) =
    let (Ξ˜β‚ , tρ , ρ₁) = substInfer ρ t
    in Ξ˜β‚ , `exfalso Οƒ tρ , ρ₁
  substInfer ρ (`cut t)                     =
    let (θ₁ , tρ , ρ₁) = substCheck ρ t
    in θ₁ , `cut tρ , ρ₁

  substCheck : Substituting Infer Check L.substCheck TInfer TCheck
  substCheck ρ (`lam t) = substLam (substCheck ([v]∷ ρ) t) 
  substCheck {t = `let _ ∷= rt `in ru} ρ (`let p ∷= t `in u) =
    let Ξ΄              = patternContext p
        (θ₁ , tρ , ρ₁) = substInfer ρ t
        (ΞΈβ‚‚ , uρ , ρ₂) = substCheck (withFreshVars Ξ΄ ρ₁) u
        (θ₃ , ρ)       = substLet Ξ΄ (ΞΈβ‚‚ , ρ₂)
        eq             = functionalEnvPre functionalInferPre _ ρ₂ (withStaleVars (patternContext p) ρ)
    in , `let p ∷= tρ `in subst (TCheck _ _ _) eq uρ , ρ
  substCheck ρ `unit   = , `unit , ρ
  substCheck ρ (`prdβŠ— a b) =
    let (θ₁ , aρ , ρ₁) = substCheck ρ a
        (ΞΈβ‚‚ , bρ , ρ₂) = substCheck ρ₁ b
    in ΞΈβ‚‚ , `prdβŠ— aρ bρ , ρ₂
  substCheck ρ (`prd& a b) =
    let (θ₁ , aρ , ρ₁) = substCheck ρ a
        (ΞΈβ‚‚ , bρ , ρ₂) = substCheck ρ b
        eq             = functionalEnvPre functionalInferPre _ ρ₂ ρ₁
    in , `prd& aρ (subst (TCheck _ _ _) eq bρ) , ρ₁
  substCheck ρ (`inl t) =
    let (θ₁ , tρ , ρ₁) = substCheck ρ t
    in θ₁ , `inl tρ , ρ₁
  substCheck ρ (`inr t) =
    let (θ₁ , tρ , ρ₁) = substCheck ρ t
    in θ₁ , `inr tρ , ρ₁
  substCheck ρ (`neu t) =
    let (θ₁ , tρ , ρ₁) = substInfer ρ t
    in θ₁ , `neu tρ , ρ₁