module CBPV.Family where

open import Data.List as List using (List; []; _∷_; foldr)
open import Data.List.Relation.Unary.All using (All; []; _∷_)
open import Data.Sum as ⊎ using (_⊎_; inj₁; inj₂; [_,_]; [_,_]′)
open import Data.Unit using (⊤)
open import Function.Base using (_∘_; _$_; _⟨_⟩_; flip; case_returning_of_)
open import Reflection hiding (name; Name; _≟_)
open import Reflection.AST.Argument using (_⟨∷⟩_; _⟅∷⟆_)
open import Reflection.AST.Term as Term using (_⋯⟅∷⟆_)
open import Relation.Binary.PropositionalEquality.Core
open import Relation.Binary.PropositionalEquality using (module ≡-Reasoning)
open import Relation.Nullary using (yes; no)

import Data.Sum.Properties as ⊎

open import CBPV.Context hiding (foldr)
open import CBPV.Type

-- Families and Morphisms -----------------------------------------------------

infix 4 _⇾_ _⇾ᵗ_

_-Family : Set → Set₁
A -Family = GenContext A → Set

_-Sorted-_-Family : Set → Set → Set₁
A -Sorted- B -Family = A → B -Family

_-SortedFamily : Set → Set₁
A -SortedFamily = A -Sorted- A -Family

ValFamily : Set₁
ValFamily = ValType -SortedFamily

CompFamily : Set₁
CompFamily = CompType -Sorted- ValType -Family

_⇾_ : {A : Set} → A -Family → A -Family → Set
X ⇾ Y = {Γ : _} → X Γ → Y Γ

_⇾ᵗ_ : {A B : Set} → A -Sorted- B -Family → A -Sorted- B -Family → Set
V ⇾ᵗ W = {τ : _} → V τ ⇾ W τ

-- Variables -------------------------------------------------------------------

private
  variable
    x y : Name
    A B C : Set
    Γ Δ Θ Π : GenContext A
    τ τ′ : A
    X Y : A -Sorted- B -Family

data VarPos (x : Name) (τ : A) : GenContext A → Set where
  Here  : VarPos x τ (Γ :< (x :- τ))
  There : VarPos x τ Γ → VarPos x τ (Γ :< (y :- τ′))

record Var (τ : A) (Γ : GenContext A) : Set where
  constructor MkVar
  field
    name : Name
    pos : VarPos name τ Γ

open Var public

toVar : VarPos x τ Γ → Var τ Γ
toVar pos = MkVar _ pos

-- Substitutions --------------------------------------------------------------

infixl 9 _⨾_ [_]_⨾_ _⨾′_ _≈⨾_ _≈⨾′_ _⨾≈_
infixl 4 copair _∥ˡ_ _∥ʳ_
infixl 5 _:<_↦_ subst-step
infix 4 _~[_]↝_ _↝_ [_]_≈_ _≈_

record _~[_]↝_ (Γ : GenContext A) (X : A -Sorted- B -Family) (Δ : GenContext B) : Set where
  constructor tabulate
  infixr 8 _＠_
  field
    _＠_ : Var τ Γ → X τ Δ

open _~[_]↝_ public

_↝_ : GenContext A → GenContext A → Set
Γ ↝ Δ = Γ ~[ Var ]↝ Δ

[_]_＠_ : (X : A -Sorted- B -Family) → Γ ~[ X ]↝ Δ → Var τ Γ → X τ Δ
[ X ] σ ＠ i = σ ＠ i

-- Equality

record _≈_ (σ ς : Γ ~[ X ]↝ Δ) : Set where
  constructor tabulate≈
  infixr 8 _＠_
  field
    _＠_ : (i : Var τ Γ) → σ ＠ i ≡ ς ＠ i

open _≈_ public

[_]_≈_ : (X : A -SortedFamily) → Γ ~[ X ]↝ Δ → Γ ~[ X ]↝ Δ → Set
[ V ] σ ≈ ς = σ ≈ ς

trans≈ : {σ₁ σ₂ σ₃ : Γ ~[ X ]↝ Δ} → σ₁ ≈ σ₂ → σ₂ ≈ σ₃ → σ₁ ≈ σ₃
trans≈ σ₁≈σ₂ σ₂≈σ₃ ＠ i = trans (σ₁≈σ₂ ＠ i) (σ₂≈σ₃ ＠ i)

-- Combinators

_⨾_ : Γ ~[ X ]↝ Δ → (∀ {τ} → X τ Δ → Y τ Θ) → Γ ~[ Y ]↝ Θ
(σ ⨾ f) ＠ i = f (σ ＠ i)

_⨾′_ : Γ ↝ Δ → Δ ~[ X ]↝ Θ → Γ ~[ X ]↝ Θ
(ρ ⨾′ f) ＠ i = f ＠ (ρ ＠ i)

[_]_⨾_ : (X : A -SortedFamily) → Γ ~[ X ]↝ Δ → (∀ {τ} → X τ Δ → Y τ Θ) → Γ ~[ Y ]↝ Θ
[ V ] σ ⨾ f = σ ⨾ f

_≈⨾_ : {σ ς : Γ ~[ X ]↝ Δ} → σ ≈ ς → (f : ∀ {τ} → X τ Δ → Y τ Θ) → [ Y ] (σ ⨾ f) ≈ (ς ⨾ f)
(σ≈ς ≈⨾ f) ＠ i = cong f (σ≈ς ＠ i)

_≈⨾′_ : {ρ ϱ : Γ ↝ Δ} → ρ ≈ ϱ → (σ : Δ ~[ X ]↝ Θ) → ρ ⨾′ σ ≈ ϱ ⨾′ σ
(ρ≈ϱ ≈⨾′ σ) ＠ i = cong (σ ＠_) (ρ≈ϱ ＠ i)

_⨾≈_ : (ρ : Γ ↝ Δ) → {σ ς : Δ ~[ X ]↝ Θ} → σ ≈ ς → ρ ⨾′ σ ≈ ρ ⨾′ ς
(ρ ⨾≈ σ≈ς) ＠ i = σ≈ς ＠ ρ ＠ i

-- Particular Substitutions

-- id

id : Γ ↝ Γ
id = tabulate (λ x → x)

-- empty

empty : [<] ~[ X ]↝ Γ
empty = tabulate (λ ())

-- weakr

weakrPos : VarPos x τ Δ → VarPos x τ (Γ ++ Δ)
weakrPos Here = Here
weakrPos (There i) = There (weakrPos i)

weakrF : Var τ Δ → Var τ (Γ ++ Δ)
weakrF i = toVar $ weakrPos $ i .pos

weakr : Δ ↝ Γ ++ Δ
weakr = tabulate weakrF

-- weakl

weaklPos : (Δ : GenContext A) → VarPos x τ Γ → VarPos x τ (Γ ++ Δ)
weaklPos [<] i = i
weaklPos (Δ :< (x :- τ)) i = There (weaklPos Δ i)

weaklF : (Δ : GenContext A) → Var τ Γ → Var τ (Γ ++ Δ)
weaklF Δ i = toVar $ weaklPos Δ $ i .pos

weakl : (Δ : GenContext A) → Γ ↝ Γ ++ Δ
weakl Δ = tabulate (weaklF Δ)

-- copair

splitPos : (Δ : GenContext A) → VarPos x τ (Γ ++ Δ) → VarPos x τ Γ ⊎ VarPos x τ Δ
splitPos [<] i = inj₁ i
splitPos (Δ :< _) Here = inj₂ Here
splitPos (Δ :< _) (There i) = ⊎.map₂ There (splitPos Δ i)

splitF : (Δ : GenContext A) → Var τ (Γ ++ Δ) → Var τ Γ ⊎ Var τ Δ
splitF Δ i = ⊎.map toVar toVar $ splitPos Δ $ i .pos

copair : Γ ~[ X ]↝ Θ → Δ ~[ X ]↝ Θ → Γ ++ Δ ~[ X ]↝ Θ
copair σ ς ＠ i = [ σ ＠_ , ς ＠_ ] $ splitF _ i

syntax copair σ ς = σ ∥ ς

copairₙ : {Γs : List (GenContext A)} → All (_~[ X ]↝ Θ) Γs → foldr _++_ [<] Γs ~[ X ]↝ Θ
copairₙ [] ＠ ()
copairₙ (σ ∷ σs) = σ ∥ copairₙ σs

-- singleton

singleton : (x : Name) → X τ Γ → [< x :- τ ] ~[ X ]↝ Γ
singleton x v ＠ (MkVar .x Here) = v

-- snoc

_:<_↦_ : Γ ~[ X ]↝ Δ → (x : Name) → X τ Δ → Γ :< (x :- τ) ~[ X ]↝ Δ
σ :< x ↦ v = σ ∥ singleton x v

-- weakening

data Path {A : Set} : GenContext A → GenContext A → Set where
  ⌞_⌟ : (Γ : GenContext A) → Path Γ Γ
  _∥ˡ_ : Path Γ Δ → (Θ : GenContext A) → Path Γ (Δ ++ Θ)
  _∥ʳ_ : (Θ : GenContext A) → Path Γ Δ → Path Γ (Θ ++ Δ)

wkn : Path Γ Δ → Γ ↝ Δ
wkn ⌞ Γ ⌟ = id
wkn (path ∥ˡ Θ) = wkn path ⨾′ weakl Θ
wkn (Θ ∥ʳ path) = wkn path ⨾′ weakr

-- lift

lift′ : (Θ : GenContext A) → Γ ↝ Δ → Γ ++ Θ ↝ Δ ++ Θ
lift′ Θ ρ = ρ ⨾ weaklF Θ ∥ weakr

-- weakening variables

_↑_ : Var τ Γ → (Δ : GenContext A) → Var τ (Γ ++ Δ)
i ↑ Δ = weakl Δ ＠ i

_↑[_] : Var τ Γ → (x : Name) → Var τ (Γ :< (x :- τ′))
i ↑[ x ] = i ↑ [< x :- _ ]

_↑ : Var τ Γ → Var τ (Γ :< (x :- τ′))
i ↑ = i ↑[ _ ]

-- interaction with map

map⁺Pos : {f : A → B} → Γ ~[ X ∘ f ]↝ Δ → VarPos x τ (map f Γ) → X τ Δ
map⁺Pos {Γ = Γ :< (x :- τ)} σ Here = σ ＠ toVar Here
map⁺Pos {Γ = Γ :< (x :- τ)} {X = X} σ (There i) = map⁺Pos {X = X} (tabulate (λ i → σ ＠ (i ↑))) i

map⁺ : {f : A → B} → Γ ~[ X ∘ f ]↝ Δ → map f Γ ~[ X ]↝ Δ
map⁺ σ ＠ i = map⁺Pos σ (i .pos)

varPosMap⁺ : {f : A → B} → VarPos x τ Γ → VarPos x (f τ) (map f Γ)
varPosMap⁺ Here = Here
varPosMap⁺ (There i) = There (varPosMap⁺ i)

varMap⁺ : {f : A → B} → Var τ Γ → Var (f τ) (map f Γ)
varMap⁺ v = toVar (varPosMap⁺ $ v .pos)

map⁻ : {f : A → B} → map f Γ ~[ X ]↝ Δ → Γ ~[ X ∘ f ]↝ Δ
map⁻ σ ＠ i = σ ＠ varMap⁺ i

-- Properties

subst-step :
  {x : Name} {σ ς : Γ :< (x :- τ) ~[ X ]↝ Δ} →
  weakl [< x :- τ ] ⨾′ σ ≈ weakl [< x :- τ ] ⨾′ ς →
  σ ＠ MkVar x Here ≡ ς ＠ MkVar x Here →
  σ ≈ ς
subst-step σ↑≈ς↑ σ₀≡ς₀ ＠ MkVar x Here = σ₀≡ς₀
subst-step σ↑≈ς↑ σ₀≡ς₀ ＠ MkVar x (There pos) = σ↑≈ς↑ ＠ MkVar x pos

syntax subst-step {x = x} σ≈ς v≡w = σ≈ς :< x ⇒ v≡w

-- weakl [<] is id

weakl[<]≈id : weakl {Γ = Γ} [<] ≈ id
weakl[<]≈id = tabulate≈ (λ i → refl)

-- copair is coslice coproduct

splitPos-weaklPos : (Δ : GenContext A) → (i : VarPos x τ Γ) → splitPos Δ (weaklPos Δ i) ≡ inj₁ i
splitPos-weaklPos [<] i = refl
splitPos-weaklPos (Δ :< _) i = cong (⊎.map₂ There) (splitPos-weaklPos Δ i)

splitF-weaklF : (Δ : GenContext A) (i : Var τ Γ) → splitF Δ (weaklF Δ i) ≡ inj₁ i
splitF-weaklF Δ i = cong (⊎.map toVar toVar) (splitPos-weaklPos Δ $ i .pos)

copair-weakl :
  (σ : Γ ~[ X ]↝ Θ) (ς : Δ ~[ X ]↝ Θ) →
  [ X ] weakl Δ ⨾′ (σ ∥ ς) ≈ σ
copair-weakl σ ς ＠ i = cong [ σ ＠_ , ς ＠_ ] (splitF-weaklF _ i)

splitPos-weakrPos :
  (Γ : GenContext A) → (i : VarPos x τ Δ) →
  splitPos {Γ = Γ} Δ (weakrPos i) ≡ inj₂ i
splitPos-weakrPos Γ Here = refl
splitPos-weakrPos Γ (There i) = cong (⊎.map₂ There) (splitPos-weakrPos Γ i)

splitF-weakrF : (Γ : GenContext A) (i : Var τ Δ) → splitF Δ (weakrF i) ≡ inj₂ i
splitF-weakrF Γ i = cong (⊎.map toVar toVar) (splitPos-weakrPos Γ $ i .pos)

copair-weakr :
  (σ : Γ ~[ X ]↝ Θ) (ς : Δ ~[ X ]↝ Θ) →
  [ X ] weakr {Γ = Γ} ⨾′ (σ ∥ ς) ≈ ς
copair-weakr σ ς ＠ i = cong [ σ ＠_ , ς ＠_ ] (splitF-weakrF _ i)

weakPos-splitPos :
  (Δ : GenContext A) (i : VarPos x τ (Γ ++ Δ)) →
  [ weaklPos Δ , weakrPos ] (splitPos Δ i) ≡ i
weakPos-splitPos [<] i = refl
weakPos-splitPos (Δ :< _) Here = refl
weakPos-splitPos (Δ :< _) (There i) = begin
  [ There ∘ weaklPos Δ , weakrPos ]′ (⊎.map₂ There $ splitPos Δ i) ≡⟨ ⊎.[,]-map (splitPos Δ i) ⟩
  [ There ∘ weaklPos Δ , weakrPos ∘ There ]′ (splitPos Δ i)        ≡⟨⟩
  [ There ∘ weaklPos Δ , There ∘ weakrPos ]′ (splitPos Δ i)        ≡˘⟨ ⊎.[,]-∘ There (splitPos Δ i) ⟩
  There ([ weaklPos Δ , weakrPos ]′ $ splitPos Δ i)                ≡⟨ cong There (weakPos-splitPos Δ i) ⟩
  There i                                                          ∎
  where open ≡-Reasoning

weakF-splitF : (Δ : GenContext A) (i : Var τ (Γ ++ Δ)) → [ weaklF Δ , weakrF ] (splitF Δ i) ≡ i
weakF-splitF Δ i = begin
  [ weaklF Δ , weakrF ] (splitF Δ i)                              ≡⟨ ⊎.[,]-map (splitPos Δ $ i .pos) ⟩
  [ weaklF Δ ∘ toVar , weakrF ∘ toVar ] (splitPos Δ $ i .pos)     ≡⟨⟩
  [ toVar ∘ weaklPos Δ , toVar ∘ weakrPos ] (splitPos Δ $ i .pos) ≡˘⟨ ⊎.[,]-∘ toVar (splitPos Δ $ i .pos) ⟩
  toVar ([ weaklPos Δ , weakrPos ]′ $ splitPos Δ $ i .pos)        ≡⟨ cong toVar (weakPos-splitPos Δ (i .pos)) ⟩
  toVar (i .pos)                                                  ≡⟨⟩
  i                                                               ∎
  where open ≡-Reasoning

copair-unique :
  {X : A -Sorted- B -Family} →
  (σ : Γ ~[ X ]↝ Θ) (ς : Δ ~[ X ]↝ Θ) (s : Γ ++ Δ ~[ X ]↝ Θ) →
  (weakl Δ ⨾′ s ≈ σ) → (weakr ⨾′ s ≈ ς) →
  (σ ∥ ς) ≈ s
copair-unique {Δ = Δ} σ ς s wkl⨾s≈σ wkr⨾s≈ς ＠ i = begin
  [ σ ＠_ , ς ＠_ ] (splitF Δ i)                         ≡˘⟨ ⊎.[,]-cong (wkl⨾s≈σ ＠_) (wkr⨾s≈ς ＠_) (splitF Δ i) ⟩
  [ (weakl Δ ⨾′ s) ＠_ , (weakr ⨾′ s) ＠_ ] (splitF Δ i) ≡˘⟨ ⊎.[,]-∘ (s ＠_) (splitF Δ i) ⟩
  s ＠ ([ weaklF Δ , weakrF ] $ splitF Δ i)              ≡⟨ cong (s ＠_) (weakF-splitF Δ i) ⟩
  s ＠ i                                                 ∎
  where open ≡-Reasoning

-- Special Families -----------------------------------------------------------

infixr 5 _⇒_

_⇒_ : A -Family → A -Family → A -Family
(X ⇒ Y) Γ = X Γ → Y Γ

_^_ : A -Family → B -Sorted- A -Family → B -Family
(X ^ V) Γ = {Δ : GenContext _} → Γ ~[ V ]↝ Δ → X Δ

_^ᵗ_ : A -Sorted- B -Family → C -Sorted- B -Family → A -Sorted- C -Family
(W ^ᵗ V) τ = W τ ^ V

□_ : A -Family → A -Family
□_ = _^ Var

□ᵗ_ : A -Sorted- B -Family → A -Sorted- B -Family
□ᵗ_ = _^ᵗ Var

δ : GenContext A → A -Family → A -Family
δ Δ X Γ = X (Γ ++ Δ)

δᵗ : GenContext B → A -Sorted- B -Family → A -Sorted- B -Family
δᵗ Δ V τ = δ Δ (V τ)

-- Reflection ------------------------------------------------------------------

wknTerm : Term → Term
wknTerm t = quote There ⟨ con ⟩ 5 ⋯⟅∷⟆ t ⟨∷⟩ []

weaklTerm : Term → Term → Term
weaklTerm Δ t = quote weaklPos ⟨ def ⟩ 3 ⋯⟅∷⟆ Δ ⟨∷⟩ t ⟨∷⟩ []

searchCtx :
  (wkn : Term → Term) →
  (hole : Term) →
  (name : Term) →
  (Γ : Term) →
  TC ⊤
searchCtx wkn hole name (meta m _) = blockOnMeta m
searchCtx wkn hole name (con (quote _:<_) (_ ⟅∷⟆ Γ ⟨∷⟩ con (quote _:-_) (_ ⟅∷⟆ meta x _ ⟨∷⟩ ty ∷ []) ⟨∷⟩ [])) =
  blockOnMeta x
searchCtx wkn hole name (con (quote _:<_) (_ ⟅∷⟆ Γ ⟨∷⟩ con (quote _:-_) (_ ⟅∷⟆ x ⟨∷⟩ ty ∷ []) ⟨∷⟩ [])) =
  catchTC
    (do
      debugPrint "squid" 10 (strErr "checking name " ∷ termErr x ∷ strErr "..." ∷ [])
      unify name x
      let soln = wkn (quote Here ⟨ con ⟩ 3 ⋯⟅∷⟆ [])
      debugPrint "squid" 10 (strErr "testing solution " ∷ termErr soln ∷ strErr "..." ∷ [])
      unify soln hole
      debugPrint "squid" 10 (strErr "success!" ∷ []))
    (searchCtx (wkn ∘ wknTerm) hole name Γ)
searchCtx wkn hole name (con (quote [<]) (_ ⟅∷⟆ [])) =
  typeError (termErr name ∷ strErr " not in context." ∷ [])
searchCtx wkn hole name (def (quote _++_) (_ ⟅∷⟆ Γ ⟨∷⟩ Δ ⟨∷⟩ [])) =
  catchTC
    (searchCtx wkn hole name Δ)
    (searchCtx (wkn ∘ weaklTerm Δ) hole name Γ)
searchCtx wkn hole name Γ =
  typeError (termErr name ∷ strErr " not in context " ∷ termErr Γ ∷ [])

searchVarPos : {A : Set} → GenContext A → Name → Term → TC ⊤
searchVarPos Γ x hole =
  do
    Γ ← quoteTC Γ
    x ← quoteTC x
    case x returning (λ _ → TC ⊤) of λ
      { (meta x _) → blockOnMeta x
      ; _ → pure _
      }
    debugPrint "squid" 10 (strErr "Γ = " ∷ termErr Γ ∷ [])
    debugPrint "squid" 10 (strErr "x = " ∷ termErr x ∷ [])
    searchCtx (λ t → t) hole x Γ

%%_ :
  {A : Set} → {ty : A} → {Γ : GenContext A} →
  (x : Name) → {@(tactic searchVarPos Γ x) pos : VarPos x ty Γ} → Var ty Γ
%%_ x {pos} = MkVar x pos