module CBPV.Family where

open import Data.List as List using (List; []; _∷_; map)
open import Data.Unit using (⊤)
open import Function.Base using (_∘_; _$_; _⟨_⟩_; flip; case_return_of_)
open import Reflection hiding (name; Name; _≟_)
open import Reflection.Argument using (_⟨∷⟩_)
open import Reflection.Term as Term using (_⋯⟅∷⟆_)
open import Relation.Binary.PropositionalEquality
open import Relation.Nullary using (yes; no)

open import CBPV.Context
open import CBPV.Type

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

infix 4 _⇾_ _⇾ᵛ_ _⇾ᶜ_

Family : Set₁
Family = Context → Set

ValFamily : Set₁
ValFamily = ValType → Family

CompFamily : Set₁
CompFamily = CompType → Family

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

_⇾ᵛ_ : ValFamily → ValFamily → Set
V ⇾ᵛ W = {A : ValType} → V A ⇾ W A

_⇾ᶜ_ : CompFamily → CompFamily → Set
C ⇾ᶜ D = {B : CompType} → C B ⇾ D B

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

private
  variable
    Γ Δ Θ Π : Context
    A A′ : ValType
    x y : Name
    V W X : ValFamily

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

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

-- 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 [<]) []) =
  typeError (termErr name ∷ strErr " not in context." ∷ [])
searchCtx wkn hole name (con (quote _:<_) (Γ ⟨∷⟩ con (quote _:-_) (meta x _ ⟨∷⟩ A ∷ []) ⟨∷⟩ [])) =
  blockOnMeta x
searchCtx wkn hole name (con (quote _:<_) (Γ ⟨∷⟩ con (quote _:-_) (x ⟨∷⟩ A ∷ []) ⟨∷⟩ [])) =
  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 (def (quote _++_) (Γ ⟨∷⟩ Δ ⟨∷⟩ [])) =
  catchTC
    (searchCtx wkn hole name Δ)
    (searchCtx (wkn ∘ weaklTerm Δ) hole name Γ)
searchCtx wkn hole name Γ =
  typeError (termErr name ∷ strErr " not in context " ∷ termErr Γ ∷ [])
-- searchCtx (con (quote [<]) []) y = return []
-- searchCtx (con (quote _:<_) (Γ ⟨∷⟩ con (quote _:-_) (x ⟨∷⟩ A ∷ []) ⟨∷⟩ [])) y =
-- --   with x Term.≟ y
-- -- ... | yes _ =
--   do
--     let head = quote Here ⟨ con ⟩ 3 ⋯⟅∷⟆ []
--     tail ← searchCtx Γ y
--     return (head ∷ map wknTerm tail)
-- -- ... | no  _ =
-- --   do
-- --     tail ← searchCtx Γ y
-- --     return (map wknTerm tail)
-- searchCtx (def (quote _++_) (Γ ⟨∷⟩ Δ ⟨∷⟩ [])) y =
--   do
--     right ← searchCtx Δ y
--     left ← searchCtx Γ y
--     return (right List.++ map (weaklTerm Δ) left)

-- searchCtx (meta m _) y = blockOnMeta m
-- searchCtx Γ y =
--   do
--     debugPrint "squid" 10 (strErr "stuck matching term " ∷ termErr Γ ∷ [])
--     return []

-- tryEach : List Term → Term → TC ⊤
-- tryEach [] hole = typeError (strErr "not in context" ∷ [])
-- tryEach (t ∷ []) hole =
--   do
--     debugPrint "squid" 10 (strErr "trying term " ∷ termErr t ∷ [])
--     unify t hole
--     debugPrint "squid" 10 (strErr "success!" ∷ [])
-- tryEach (t ∷ ts@(_ ∷ _)) hole =
--   do
--     debugPrint "squid" 10 (strErr "trying term " ∷ termErr t ∷ [])
--     catchTC
--       (do
--         unify t hole
--         debugPrint "squid" 10 (strErr "success!" ∷ []))
--       (tryEach ts hole)

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

-- searchVarPos : Context → Name → Term → TC ⊤
-- searchVarPos Γ x hole =
--   do
--     Γ ← quoteTC Γ
--     (case Γ return (λ _ → TC ⊤) of λ
--       { (meta m _) → blockOnMeta m
--       ; _ → return _
--       })
--     x ← quoteTC x
--     debugPrint "squid" 10 (strErr "Γ = " ∷ termErr Γ ∷ [])
--     debugPrint "squid" 10 (strErr "x = " ∷ termErr x ∷ [])
--     guesses ← searchCtx Γ x
--     tryEach guesses hole

-- Back to business as normal

record Var (A : ValType) (Γ : Context) : Set where
  constructor %%_
  field
    name : Name
    @(tactic searchVarPos Γ name) {pos} : VarPos name A Γ

open Var public

toVar : VarPos x A Γ → Var A Γ
toVar pos = (%% _) {pos}

ThereVar : Var A Γ → Var A (Γ :< (y :- A′))
ThereVar i = toVar $ There (i .pos)

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

infixl 9 _⨾_ [_]_⨾_
infixr 8 _＠_
infixl 5 _:<_↦_
infix 4 Subst _↝_

data Subst (V : ValFamily) (Δ : Context) : Context → Set where
  [<]     : Subst V Δ [<]
  _:<_↦_ : Subst V Δ Γ → (x : Name) → V A Δ → Subst V Δ (Γ :< (x :- A))

syntax Subst V Δ Γ = Γ ~[ V ]↝ Δ

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

_⨾_ : Γ ~[ V ]↝ Δ → (∀ {A} → V A Δ → W A Θ) → Γ ~[ W ]↝ Θ
[<] ⨾ f = [<]
(σ :< x ↦ v) ⨾ f = (σ ⨾ f) :< x ↦ f v

[_]_⨾_ : (V : ValFamily) → Γ ~[ V ]↝ Δ → (∀ {A} → V A Δ → W A Θ) → Γ ~[ W ]↝ Θ
[ V ] σ ⨾ f = σ ⨾ f

tabulate : (∀ {A} → Var A Γ → V A Δ) → Γ ~[ V ]↝ Δ
tabulate {Γ = [<]} f = [<]
tabulate {Γ = Γ :< (x :- A)} f = tabulate (f ∘ ThereVar) :< x ↦ f (toVar Here)

lookup : Γ ~[ V ]↝ Δ → VarPos x A Γ → V A Δ
lookup (σ :< x ↦ v) Here = v
lookup (σ :< x ↦ v) (There i) = lookup σ i

_＠_ : Γ ~[ V ]↝ Δ → Var A Γ → V A Δ
σ ＠ i = lookup σ (i .pos)

[_]_＠_ : (V : ValFamily) → Γ ~[ V ]↝ Δ → Var A Γ → V A Δ
[ V ] σ ＠ i = σ ＠ i

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

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

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

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

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

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

copair : Γ ~[ V ]↝ Θ → Δ ~[ V ]↝ Θ → Γ ++ Δ ~[ V ]↝ Θ
copair σ [<] = σ
copair σ (ς :< x ↦ v) = copair σ ς :< x ↦ v

pull : (Δ Θ : Context) → Γ ++ Δ ++ Θ ↝ Γ ++ Θ ++ Δ
pull Δ Θ =
  copair {Δ = Θ}
    (copair {Δ = Δ} (weakl Θ ⨾ (weakl Δ ＠_)) weakr)
    (weakr ⨾ (weakl Δ ＠_))

-- Properties

subst-ext : (σ ς : Γ ~[ V ]↝ Δ) → (∀ {A} → (i : Var A Γ) → σ ＠ i ≡ ς ＠ i) → σ ≡ ς
subst-ext [<] [<] ext = refl
subst-ext (σ :< x ↦ v) (ς :< .x ↦ w) ext =
  cong₂ (_:< x ↦_) (subst-ext σ ς $ λ i → ext $ ThereVar i) (ext $ %% x)

⨾-assoc :
  (σ : Γ ~[ V ]↝ Δ) (f : ∀ {A} → V A Δ → W A Θ) (g : ∀ {A} → W A Θ → X A Π) →
  _⨾_ {V = W} {W = X} (σ ⨾ f) g ≡ σ ⨾ (g ∘ f)
⨾-assoc [<] f g = refl
⨾-assoc (σ :< x ↦ v) f g = cong (_:< x ↦ g (f v)) (⨾-assoc σ f g)

⨾-cong :
  (σ : Γ ~[ V ]↝ Δ) {f g : ∀ {A} → V A Δ → W A Θ} →
  (∀ {A} → (v : V A Δ) → f v ≡ g v) →
  _⨾_ {W = W} σ f ≡ σ ⨾ g
⨾-cong [<] ext = refl
⨾-cong (σ :< x ↦ v) ext = cong₂ (_:< x ↦_) (⨾-cong σ ext) (ext v)

lookup-⨾ :
  (σ : Γ ~[ V ]↝ Δ) (f : ∀ {A} → V A Δ → W A Θ) (i : VarPos x A Γ) →
  lookup {V = W} (σ ⨾ f) i ≡ f (lookup σ i)
lookup-⨾ (σ :< x ↦ v) f Here = refl
lookup-⨾ (σ :< x ↦ v) f (There i) = lookup-⨾ σ f i

lookup-tabulate :
  (V : ValFamily) (f : ∀ {A} → Var A Γ → V A Δ) (i : VarPos x A Γ) →
  lookup {V = V} (tabulate f) i ≡ f (toVar i)
lookup-tabulate V f Here = refl
lookup-tabulate V f (There i) = lookup-tabulate V (f ∘ ThereVar) i

＠-⨾ :
  (σ : Γ ~[ V ]↝ Δ) (f : ∀ {A} → V A Δ → W A Θ) (i : Var A Γ) →
  _＠_ {V = W} (σ ⨾ f) i ≡ f (σ ＠ i)
＠-⨾ σ f i = lookup-⨾ σ f (i .pos)

＠-tabulate :
  (V : ValFamily) (f : ∀ {A} → Var A Γ → V A Δ) (i : Var A Γ) →
  _＠_ {V = V} (tabulate f) i ≡ f i
＠-tabulate V f i = lookup-tabulate V f (i .pos)

tabulate-cong :
  (V : ValFamily) →
  {f g : ∀ {A} → Var A Γ → V A Δ} →
  (∀ {A} → (v : Var A Γ) → f v ≡ g v) →
  tabulate {V = V} f ≡ tabulate g
tabulate-cong {Γ = [<]} V ext = refl
tabulate-cong {Γ = Γ :< (x :- A)} V ext =
  cong₂ (_:< x ↦_) (tabulate-cong V (ext ∘ ThereVar)) (ext $ %% x)

tabulate-⨾ :
  (f : ∀ {A} → Var A Γ → V A Δ) (g : ∀ {A} → V A Δ → W A Θ) →
  tabulate {V = V} f ⨾ g ≡ tabulate {V = W} (g ∘ f)
tabulate-⨾ {Γ = [<]} f g = refl
tabulate-⨾ {Γ = Γ :< (x :- A)} f g = cong (_:< x ↦ g (f $ toVar Here)) (tabulate-⨾ (f ∘ ThereVar) g)

tabulate-＠ : (σ : Γ ~[ V ]↝ Δ) → tabulate (σ ＠_) ≡ σ
tabulate-＠ [<] = refl
tabulate-＠ (σ :< x ↦ v) = cong (_:< x ↦ v) (tabulate-＠ σ)

copair-⨾ :
  (σ : Γ ~[ V ]↝ Θ) (ς : Δ ~[ V ]↝ Θ) (f : ∀ {A} → V A Θ → W A Π) →
  copair {V = W} (σ ⨾ f) (ς ⨾ f) ≡ copair {V = V} σ ς ⨾ f
copair-⨾ σ [<] f = refl
copair-⨾ σ (ς :< x ↦ v) f = cong (_:< x ↦ f v) (copair-⨾ σ ς f)

copair-left :
  (σ : Γ ~[ V ]↝ Θ) (ς : Δ ~[ V ]↝ Θ) (i : VarPos x A Γ) →
  lookup (copair σ ς) (weaklPos Δ i) ≡ lookup σ i
copair-left σ [<] i = refl
copair-left σ (ς :< x ↦ v) i = copair-left σ ς i

pull-left :
  (Δ Θ : Context) (i : VarPos x A Γ) →
  lookup (pull Δ Θ) (weaklPos Θ $ weaklPos Δ i) ≡ toVar (weaklPos Δ $ weaklPos Θ i)
pull-left {Γ = Γ} Δ Θ i =
  begin
    lookup (pull Δ Θ) (weaklPos Θ $ weaklPos Δ i)
  ≡⟨ copair-left (copair {Δ = Δ} (weakl Θ ⨾ (weakl Δ ＠_)) weakr) (weakr {Γ = Γ} ⨾ (weakl Δ ＠_)) (weaklPos Δ i) ⟩
    lookup (copair {Δ = Δ} (weakl Θ ⨾ (weakl Δ ＠_)) weakr) (weaklPos Δ i)
  ≡⟨ copair-left {Δ = Δ} (weakl Θ ⨾ (weakl Δ ＠_)) weakr i ⟩
    lookup (weakl Θ ⨾ (weakl Δ ＠_)) i
  ≡⟨ lookup-⨾ (weakl Θ) (weakl Δ ＠_) i ⟩
    weakl Δ ＠ lookup (weakl Θ) i
  ≡⟨ cong (weakl Δ ＠_) (lookup-tabulate Var (weaklF Θ) i) ⟩
    lookup (weakl Δ) (weaklPos Θ i)
  ≡⟨ lookup-tabulate Var (weaklF Δ) (weaklPos Θ i) ⟩
    toVar (weaklPos Δ $ weaklPos Θ i)
  ∎
  where open ≡-Reasoning

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

infixr 5 _⇒_

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

_^_ : Family → ValFamily → Family
(X ^ V) Γ = {Δ : Context} → Γ ~[ V ]↝ Δ → X Δ

_^ᵛ_ : ValFamily → ValFamily → ValFamily
(W ^ᵛ V) A = W A ^ V

_^ᶜ_ : CompFamily → ValFamily → CompFamily
(C ^ᶜ V) B = C B ^ V

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

□ᵛ_ : ValFamily → ValFamily
□ᵛ_ = _^ᵛ Var

□ᶜ_ : CompFamily → CompFamily
□ᶜ_ = _^ᶜ Var

δ : Context → Family → Family
δ Δ X Γ = X (Γ ++ Δ)

δᵛ : Context → ValFamily → ValFamily
δᵛ Δ V A = δ Δ (V A)

δᶜ : Context → CompFamily → CompFamily
δᶜ Δ C B = δ Δ (C B)