{-# OPTIONS --safe #-}
{-# OPTIONS --with-K #-}
module Pat.Matching where

open import Data.Maybe.Base using (Maybe; just; nothing; _>>=_)
open import Data.Product.Base as P using (_×_)
open import Data.Sum.Base using (_⊎_; inj₁; inj₂)
open import Relation.Binary.PropositionalEquality.Core using (_≡_; refl; sym; trans; cong)
open import Relation.Nullary.Decidable.Core using (Dec; yes; no)

open import Pat.Context
open import Pat.Context.Equality
open import Pat.Term
open import Pat.Term.Equality



-- Pattern Matching

private module _ where
  ++-↾≥-++ : ∀ (Δ′ : Context m) m≥n → Γ ↾≥ m≥n ++ (Δ ↾≥ m≥n ++ Δ′) ≡ Γ ++ Δ ↾≥ m≥n ++ Δ′
  ++-↾≥-++ Δ′ m≥n = trans (++-assoc Δ′) (cong (_++ Δ′) (sym (++-↾≥ _ m≥n)))

  ⊩↝-castʳ : Π₁ ≡ Π₂ → Γ ∣ Δ ⊩ Δ′ ↝ Π₁ → Γ ∣ Δ ⊩ Δ′ ↝ Π₂
  ⊩↝-castʳ refl π = π

match : Pattern n Γ Δ A Π → Term n (Γ ++ Δ) A → Maybe (Subst Π Γ)
matchSubst : Γ ∣ Δ ⊩ Δ′ ↝ Π → Γ ++ Δ ⊩ Δ′ → Maybe (Subst Π Γ)

private module _ where
  match' : Γ ++ Δ ≡ Γ′ → Pattern n Γ Δ A Π → Term n Γ′ A → Maybe (Subst Π Γ)
  match' refl = match

match `pat e = just (∅ , term e)
match (`_`with₁_ {Δ′ = Δ₁} x₁ σ₁) (`_`with_ {Δ = Δ₂} x₂ σ₂) with Δ₁ ≟ᶜ Δ₂
... | no _ = nothing
... | yes refl with ∋-++⁺ˡ _ x₁ ≟ˣ x₂
... | no _ = nothing
... | yes refl = viewˢ σ₁ σ₂ >>= λ _ → just ∅
match (`_`with₂_ {Δ′ = Δ₁} x₁ π₁) (`_`with_ {Δ = Δ₂} x₂ σ₂) with Δ₁ ≟ᶜ Δ₂
... | no _ = nothing
... | yes refl with ∋-++⁺ʳ _ x₁ ≟ˣ x₂
... | no _ = nothing
... | yes refl = matchSubst π₁ σ₂
match `true `true = just ∅
match `false `false = just ∅
match (`if p₁ `then p₂ `else p₃) (`if e₁ `then e₂ `else e₃) =
  match p₁ e₁ >>= λ σ₁ →
  match p₂ e₂ >>= λ σ₂ →
  match p₃ e₃ >>= λ σ₃ →
  just (σ₁ ++ˢ σ₂ ++ˢ σ₃)
match (`λ⟨ Δ ⟩ p) (`λ⟨ Δ ⟩ e) = match p e
-- require the argument type to be the same
match (_·_ {Δ′ = Δ₁} {A = A₁} p₁ p₂) (_·_ {Δ = Δ₂} {A = A₂} e₁ e₂) with Δ₁ ≟ᶜ Δ₂ | A₁ ≟ᵗ A₂
...| no _ | _ = nothing
...| yes _ | no _ = nothing
... | yes refl | yes refl =
  match p₁ e₁ >>= λ σ₁ →
  match' (++-assoc (inject₁ Δ₁)) p₂ e₂ >>= λ σ₂ →
  just (σ₁ ++ˢ σ₂)
match {Δ = Δ} (⟨ p ⟩) ⟨ e ⟩ =
  match' (sym (++-↾ Δ)) p e >>= λ σ →
  just (⊩-↾⁻ σ)
match (`let⟨_⟩ {A = A₁} Δ₁ p₁ p₂) (`let⟨_⟩ {A = A₂} Δ₂ e₁ e₂) with Δ₁ ≟ᶜ Δ₂ | A₁ ≟ᵗ A₂
... | no _ | _ = nothing
... | yes _ | no _ = nothing
... | yes refl | yes refl =
  match' (++-assoc (inject₁ Δ₁)) p₁ e₁ >>= λ σ₁ →
  match p₂ e₂ >>= λ σ₂ →
  just (σ₁ ++ˢ σ₂)
match (`wrap Δ₁ p) (`wrap Δ₂ e) with Δ₁ ≟ᶜ Δ₂
... | no _ = nothing
... | yes refl = match' (++-assoc (inject₁ Δ₁)) p e
match (`letwrap {A = A₁} Δ₁ p₁ p₂) (`letwrap {A = A₂} Δ₂ e₁ e₂) with Δ₁ ≟ᶜ Δ₂ | A₁ ≟ᵗ A₂
... | no _ | _ = nothing
... | yes _ | no _ = nothing
... | yes refl | yes refl =
  match p₁ e₁ >>= λ σ₁ →
  match p₂ e₂ >>= λ σ₂ →
  just (σ₁ ++ˢ σ₂)
match {Δ = Δ} (`iflet⟨_⟩ {A = A₁} {Π = Π₁} Δ₁ p p₁ p₂ p₃) (`iflet⟨_⟩ {A = A₂} {Π = Π₂} Δ₂ p' e₁ e₂ e₃) with Δ₁ ≟ᶜ Δ₂ | A₁ ≟ᵗ A₂ | Π₁ ≟ᶜ Π₂
... | no _ | _ | _ = nothing
... | yes _ | no _ | _ = nothing
... | yes refl | yes refl | no _ = nothing
... | yes refl | yes refl | yes refl =
  viewᵖ (transportContextₚ (sym (++-↾ Δ)) p) p' >>= λ _ → -- check if the two patterns are equal
  match' (++-assoc (inject₁ Δ₁)) p₁ e₁ >>= λ σ₁ →
  match' (++-assoc (inject₁ Π₁)) p₂ e₂ >>= λ σ₂ →
  match p₃ e₃ >>= λ σ₃ →
  just (σ₁ ++ˢ σ₂ ++ˢ σ₃)
match {Δ = Δ} (`rewritȩ {A = A₁} {Π = Π₁} p p₁ p₂) (`rewritȩ {A = A₂} {Π = Π₂} p' e₁ e₂) with A₁ ≟ᵗ A₂ | Π₁ ≟ᶜ Π₂
... | no _ | _ = nothing
... | yes _ | no _ = nothing
... | yes refl | yes refl =
  viewᵖ (transportContextₚ (sym (++-↾ Δ)) p) p' >>= λ _ → -- check if the two patterns are equal
  match' (++-assoc (inject₁ Π₁)) p₁ e₁ >>= λ σ₁ →
  match p₂ e₂ >>= λ σ₂ →
  just (σ₁ ++ˢ σ₂)
match _ _ = nothing

matchSubst ∅ ∅ = just ∅
matchSubst (_,_ {Δ′ = Δ′} {m≥n = m≥n} π (term p)) (σ , term e) =
  matchSubst π σ >>= λ σ₁ →
  match' (++-↾≥-++ Δ′ m≥n) p e >>= λ σ₂ →
  just (σ₁ ++ˢ ⊩-↾≥⁻ m≥n σ₂)
matchSubst (_,_ {m≥n = m≥n} π (var x₁)) (σ , var x₂) with x₁ ≟ˣ x₂
... | no _ = nothing
... | yes refl = matchSubst (⊩↝-castʳ (cong (_ ++_) (sym (inject≥-∅ m≥n))) π) σ
matchSubst _ _ = nothing

toPattern : Term n (Γ ++ Δ) A → Pattern n Γ Δ A ∅
toSubstPattern : Γ ++ Δ ⊩ Δ′ → Γ ∣ Δ ⊩ Δ′ ↝ ∅

private module _ where
  toPattern' :  Γ ++ Δ ≡ Γ′ → Term n Γ′ A → Pattern n Γ Δ A ∅
  toPattern' refl = toPattern

toPattern {Δ = Δ} (` x `with σ) with ∋-++⁻ Δ x
... | inj₁ xΓ = ` xΓ `with₁ σ
... | inj₂ xΔ = ` xΔ `with₂ toSubstPattern σ
toPattern `true = `true
toPattern `false = `false
toPattern (`if e `then e₁ `else e₂) = `if (toPattern e) `then (toPattern e₁) `else (toPattern e₂)
toPattern (`λ⟨ Δ′ ⟩ e) = `λ⟨ Δ′ ⟩ toPattern e
toPattern (_·_ {Δ = Δ′} e₁  e₂) = toPattern e₁ · toPattern' (++-assoc (inject₁ Δ′)) e₂
toPattern {Δ = Δ} ⟨ e ⟩ = ⟨ toPattern' (sym (++-↾ Δ)) e ⟩
toPattern (`let⟨ Δ′ ⟩ e₁ e₂) = `let⟨ Δ′ ⟩ (toPattern' (++-assoc (inject₁ Δ′)) e₁) (toPattern e₂)
toPattern (`wrap Δ e) = `wrap Δ (toPattern' (++-assoc (inject₁ Δ)) e)
toPattern (`letwrap Δ e₁ e₂) = `letwrap Δ (toPattern e₁) (toPattern e₂)
toPattern {Δ = Δ} (`iflet⟨_⟩ {Π = Π} Δ′ p e₁ e₂ e₃) = `iflet⟨ Δ′ ⟩ (transportContextₚ (++-↾ Δ) p) (toPattern' (++-assoc (inject₁ Δ′)) e₁) (toPattern' (++-assoc (inject₁ Π)) e₂) (toPattern e₃)
toPattern {Δ = Δ} (`rewritȩ {Π = Π} p e₁ e₂) = `rewritȩ (transportContextₚ (++-↾ Δ) p) (toPattern' (++-assoc (inject₁ Π)) e₁) (toPattern e₂)

toSubstPattern ∅ = ∅
toSubstPattern (_,_ {Δ = Δ′} {m≥n = m≥n} σ (term e)) = ⊩↝-castʳ (cong (∅ ++_) (inject≥-∅ m≥n)) (toSubstPattern σ , term (toPattern' (++-↾≥-++ Δ′ m≥n) e))
toSubstPattern (_,_ {m≥n = m≥n} σ (var x)) = ⊩↝-castʳ (cong (∅ ++_) (inject≥-∅ m≥n)) (toSubstPattern σ , var x)



private module _ where
  splitSubst : ∀ Π₂ → Subst (Π₁ ++ Π₂) Γ → Subst Π₁ Γ × Subst Π₂ Γ
  splitSubst ∅ σ = σ P., ∅
  splitSubst (Π₂ ,[ _ ⊢ _ ^ _ ]) (σ , x) = let (σ₁ P., σ₂) = splitSubst Π₂ σ in σ₁ P., (σ₂ , x)

  splitSubst₃ : ∀ Π₂ Π₃ → Subst (Π₁ ++ Π₂ ++ Π₃) Γ → Subst Π₁ Γ × Subst Π₂ Γ × Subst Π₃ Γ
  splitSubst₃ Π₂ Π₃ σ =
    let (σ₁₂ P., σ₃) = splitSubst Π₃ σ
    in let (σ₁ P., σ₂) = splitSubst Π₂ σ₁₂
    in σ₁ P., σ₂ P., σ₃

-- TODO: Perhaps this should be the spec and match outputs Dec (∃[ σ ] (match⁻¹ p σ ≡ e))
--       Also toPattern should invert this function
match⁻¹ : Pattern n Γ Δ A Π → Subst Π Γ → Term n (Γ ++ Δ) A
matchSubst⁻¹ : Γ ∣ Δ ⊩ Δ′ ↝ Π → Subst Π Γ → Γ ++ Δ ⊩ Δ′ 

match⁻¹ `pat (∅ , term e) = e
-- The following case can increse the nesting depth by 1
-- due to turning a variable into a term, which is not so great.
-- However, this case should not be here in the first place, as the output of match only contains term substitutions.
-- A better solution would be to restrict the input of match⁻¹ to only term substitutions.
-- Right now, we keep it here to make the function total.
match⁻¹ `pat (∅ , var x) = ` ∋-++⁺ˡ _ x `with lift (∋-++⁺ʳ _)
match⁻¹ (` x `with₁ σ) ∅ = ` ∋-++⁺ˡ _ x `with σ
match⁻¹ (` x `with₂ π) σ = ` ∋-++⁺ʳ _ x `with matchSubst⁻¹ π σ
match⁻¹ `true ∅ = `true
match⁻¹ `false ∅ = `false
match⁻¹ (`if_`then_`else_ {Π₂ = Π₂} {Π₃ = Π₃} p₁ p₂ p₃) σ =
  let (σ₁ P., σ₂ P., σ₃) = splitSubst₃ Π₂ Π₃ σ
  in `if (match⁻¹ p₁ σ₁) `then (match⁻¹ p₂ σ₂) `else (match⁻¹ p₃ σ₃)
match⁻¹ (`λ⟨ Δ′ ⟩ p) σ = `λ⟨ Δ′ ⟩ (match⁻¹ p σ)
match⁻¹ (_·_  {Δ′ = Δ′} {Π₂ = Π₂} p₁ p₂) σ =
  let (σ₁ P., σ₂) = splitSubst Π₂ σ
  in match⁻¹ p₁ σ₁ · transportContext (++-assoc (inject₁ Δ′)) (match⁻¹ p₂ σ₂)
match⁻¹ {Δ = Δ} (⟨ p ⟩) σ = ⟨ transportContext (sym (++-↾ Δ)) (match⁻¹ p (⊩-↾⁺ σ)) ⟩
match⁻¹ (`wrap Δ p) σ = `wrap Δ (transportContext (++-assoc (inject₁ Δ)) (match⁻¹ p σ))
match⁻¹ (`letwrap {Π₂ = Π₂} Δ′ p₁ p₂) σ =
  let (σ₁ P., σ₂) = splitSubst Π₂ σ
  in `letwrap Δ′ (match⁻¹ p₁ σ₁) (match⁻¹ p₂ σ₂)
match⁻¹ (`let⟨_⟩ {Π₂ = Π₂} Δ′ p₁ p₂) σ =
  let (σ₁ P., σ₂) = splitSubst Π₂ σ
  in `let⟨ Δ′ ⟩ (transportContext (++-assoc (inject₁ Δ′)) (match⁻¹ p₁ σ₁)) (match⁻¹ p₂ σ₂)
match⁻¹ {Δ = Δ} (`iflet⟨_⟩ {Π = Π} {Π₂ = Π₂} {Π₃ = Π₃} Δ′ p p₁ p₂ p₃) σ =
  let (σ₁ P., σ₂ P., σ₃) = splitSubst₃ Π₂ Π₃ σ in
  `iflet⟨ Δ′ ⟩
    (transportContextₚ (sym (++-↾ Δ)) p)
    (transportContext (++-assoc (inject₁ Δ′)) (match⁻¹ p₁ σ₁))
    (transportContext ((++-assoc (inject₁ Π))) (match⁻¹ p₂ σ₂))
    (match⁻¹ p₃ σ₃)
match⁻¹ {Δ = Δ} (`rewritȩ {Π = Π} {Π₂ = Π₂} p p₁ p₂) σ =
   let (σ₁ P., σ₂) = splitSubst Π₂ σ in
  `rewritȩ (transportContextₚ (sym (++-↾ Δ)) p) (transportContext ((++-assoc (inject₁ Π))) (match⁻¹ p₁ σ₁)) (match⁻¹ p₂ σ₂)

matchSubst⁻¹ ∅ ∅ = ∅ 
matchSubst⁻¹ (_,_  {Δ′ = Δ′} {m≥n = m≥n} {Π₂ = Π₂} π (term p)) σ =
  let (σ₁ P., σ₂) = splitSubst (inject≥ m≥n Π₂) σ in
  matchSubst⁻¹ π σ₁ , term (transportContext (++-↾≥-++ Δ′ m≥n) (match⁻¹ p (⊩-↾≥⁺ m≥n σ₂)))
matchSubst⁻¹ (_,_ {m≥n = m≥n} π (var x)) σ rewrite inject≥-∅ m≥n = matchSubst⁻¹ π σ , var x