{-# OPTIONS --universe-polymorphism #-}
module NaPa.Derived where

open import Level
import Category.Functor as Cat
import Category.Functor.Extras as Cat
import Category.Applicative as Cat
import Category.Monad as Cat
import Category.Monad.Identity as Id
open import Relation.Binary
open import Relation.Binary.Extras
open import Relation.Binary.PropositionalEquality as PropEq
open import Coinduction
open import Function
open import Data.Bool
open import Data.Indexed
open import Data.Nat using ( ℕ ; suc ; zero ) renaming (_≟_ to _≟ℕ_; _+_ to _+ℕ_)
open import Data.Maybe as Maybe
open import Data.Product
open import Data.Empty
open import Data.Sum using (_⊎_;inj₁;inj₂;[_,_]′) renaming (map to ⊎-map)
open import Data.Star.NP as Star hiding (_>>=_)
open import Relation.Nullary.Decidable as Dec
open import Relation.Nullary
open import NaPa

module Monad = Cat.RawMonad

maybeAppli : ∀ {f} → Cat.RawApplicative {f} Maybe
maybeAppli = rawIApplicative
  where open Cat.RawMonad Maybe.monad

idAppli : ∀ {f} → Cat.RawApplicative {f} id
idAppli = rawIApplicative
  where open Cat.RawMonad Id.IdentityMonad

Shift : (World → Set) → Set
Shift F = ∀ {α β} k → α +[ k ] ⊆ β → F α → F β

Shifted : (World → Set) → (World → Set)
Shifted F α = ∀ {β} k → α +[ k ] ⊆ β → F β

SynAbs : (World → Set) → World → Set
SynAbs F α = F (α ↑1)

FunAbs : (World → Set) → World → Set
FunAbs F = Shifted (F |→| F)

shiftFunAbs : ∀ {F} → Shift (FunAbs F)
shiftFunAbs {_} {α} {β} k ⊆w f {γ} k' ⊆w'
   = f (k' +ℕ k) (α +[ k' +ℕ k ]   ⊆⟨ ⊆-assoc-+′ refl k k' ⟩
                  α +[ k ] +[ k' ] ⊆⟨ ⊆-cong-+ ⊆w k' ⟩
                  β +[ k' ]        ⊆⟨ ⊆w' ⟩∎
                  γ ∎)
  where open ⊆-Reasoning

Partial : (World → Set) → World → World → Set
Partial F β α = F β → Maybe (F α)

module NamePack {β} (x : Name β) where
  nameOf      : Name β
  nameOf      = x

  Nameø-elim'  : ∀ {A : Set} → β ≡ ø → A
  Nameø-elim' eq = go eq x
    where go : ∀ {α} {A : Set} → (α ≡ ø) → Name α → A
          go refl = Nameø-elim

  _==nameOf_ : (y : Name β) → Bool
  _==nameOf_ = _==nm_ x

module WorldRelPack α β where
  outerOf  : World
  outerOf  = α
  innerOf  : World
  innerOf  = β

data Carry {I : Set} (F : I → I) {a} (A : Set a) : Rel I a where
  carry : ∀ {i} → A → Carry F A i (F i)

record EnvPack ℓ : Set (suc ℓ) where
  field
    Env        : Set ℓ → World → World → Set ℓ
    emptyEnv   : ∀ {A α} → Env A α α
    lookupEnv  : ∀ {A α β} → Env A α β → Name β → Name α ⊎ A
    _,,_       : ∀ {α β A} → Env A α β → A → Env A α (β ↑1)
    mapEnv     : ∀ {α β A B} → (A → B) → Env A α β → Env B α β

record ShiftableEnvPack ℓ : Set (suc ℓ) where
  field
    envPack : EnvPack ℓ
  open EnvPack envPack
  field
    shiftEnv : ∀ {α β γ A} k → α +[ k ] ⊆ β → Env A α γ → Env A β γ
  open EnvPack envPack public

-- CEnv is for closed Env
record CEnvPack ℓ : Set (suc ℓ) where
  field
    CEnv        : Set ℓ → World → Set ℓ
    emptyCEnv   : ∀ {A} → CEnv A ø
    lookupCEnv  : ∀ {A α} → CEnv A α → Name α → A
    _,,_        : ∀ {α A} → CEnv A α → A → CEnv A (α ↑1)
    mapCEnv     : ∀ {α A B} → (A → B) → CEnv A α → CEnv B α

lookupStar : ∀ {α β ℓ} {_↝_ : Rel World ℓ} → (∀ {γ δ} → γ ↝ δ → Name γ → Maybe (Name δ))
             → Star _↝_ α β → Name α → Name β ⊎ ∃₂ _↝_
lookupStar _ ε y = inj₁ y
lookupStar f (x ◅ Γ) y
 with f x y
... | just y' = lookupStar f Γ y'
... | nothing = inj₂ (_ , _ , x)

-- Star does not support univ poly yet
starEnv : EnvPack zero
starEnv = record { Env = Env ; emptyEnv = ε ; lookupEnv = lookupEnv
                 ; _,,_ = λ Γ x → carry x ◅ Γ ; mapEnv = mapEnv }
   where
     Env : (A : Set) → Rel World _
     Env = Star⁻¹ ∘ Carry _↑1

     lookupEnv : ∀ {α γ A} → Env A α γ → Name γ → Name α ⊎ A
     lookupEnv ε x = inj₁ x
     lookupEnv (carry z ◅ Γ) x
       with exportName  x
     ...  | just x' = lookupEnv Γ x'
     ...  | nothing = inj₂ z
     mapEnv : ∀ {α β A B} → (A → B) → Env A α β → Env B α β
     mapEnv f ε               = ε
     mapEnv f (carry x ◅ xs)  = carry (f x) ◅ mapEnv f xs

module Env = EnvPack starEnv

closeEnv : EnvPack zero → CEnvPack zero
closeEnv env
  = record { CEnv = CEnv
           ; emptyCEnv = emptyEnv
           ; lookupCEnv = lookupCEnv
           ; mapCEnv = mapEnv
           ; _,,_ = _,,_ }
  where
    open EnvPack env

    CEnv : Set → World → Set
    CEnv A = Env A ø

    lookupCEnv : ∀ {α τ} → CEnv τ α → Name α → τ
    lookupCEnv Γ = [ Nameø-elim , id ]′ ∘ lookupEnv Γ
      where open NamePack

module CEnv = CEnvPack (closeEnv starEnv)

funEnv : ∀ {ℓ} → EnvPack ℓ
funEnv {ℓ}
  = record { Env = Env
           ; emptyEnv = inj₁
           ; _,,_ = _,,_
           ; lookupEnv = id
           ; mapEnv = mapEnv }
  where
    Env : Set ℓ → World → World → Set ℓ
    Env A α β = Name β → Name α ⊎ A

    _,,_ : ∀ {α β A} → Env A α β → A → Env A α (β ↑1)
    _,,_ Γ v = maybe′ Γ (inj₂ v) ∘ exportName 

    mapEnv : ∀ {α β A B} → (A → B) → Env A α β → Env B α β
    mapEnv f g x = [ inj₁ , inj₂ ∘ f ]′ (g x)

funCEnv : ∀ {ℓ} → CEnvPack ℓ
funCEnv {ℓ}
  = record { CEnv        = CEnv
           ; emptyCEnv   = λ {η} → Nameø-elim
           ; _,,_        = _,,_
           ; lookupCEnv  = id
           ; mapCEnv     = _∘′_ }
  where
    open NamePack

    CEnv : Set ℓ → World → Set ℓ
    CEnv A α = Name α → A

    _,,_ : ∀ {α A} → CEnv A α → A → CEnv A (α ↑1)
    _,,_ Γ v = maybe′ Γ v ∘ exportName 

shiftableFunEnv : ∀ {ℓ} → ShiftableEnvPack ℓ
shiftableFunEnv = record { envPack = funEnv ; shiftEnv = shiftEnv }
  where
    open EnvPack funEnv
    shiftEnv : ∀ {α β γ A} k → α +[ k ] ⊆ β → Env A α γ → Env A β γ
    shiftEnv k α+k⊆β Γ = ⊎-map (coe α+k⊆β ∘ add k) id ∘ Γ

Su↑ : (F G : World → Set) (α β : World) → Set
Su↑ F G α β = ∀ k → F (α ↑ k) → G (β ↑ k)

Tr↑ : (F : World → Set) (M : Set → Set) (α β : World) → Set
Tr↑ F M = Su↑ F (M ∘ F)

-- Subst F G can be read as F → G in HOAS
Subst : (F G : World → Set) → Set
Subst F G = ∀ {α β} → (Name α → F β) → G α → G β

SubstVar : (F : World → Set) → Set
SubstVar F = ∀ {α β} → (Name α → F β) → Su↑ Name F α β

-- This module is more an example of how to write traverse for your data type.
module Traversable {α β M} (M-applicative : Cat.RawApplicative M)
                   (trName : Tr↑ Name M α β) where
  open Cat.RawApplicative M-applicative

  Traverse↑ : (World → Set) → Set
  Traverse↑ F = Tr↑ F M α β

  traverseName : Traverse↑ Name
  traverseName = trName

  ×-traverse : ∀ {G H} → Traverse↑ G → Traverse↑ H → Traverse↑ (G |×| H)
  ×-traverse traverseG traverseH k (x , y) = traverseG k x ⊗ traverseH k y

  traverseAbs : ∀ {G} → Traverse↑ G → Traverse↑ (SynAbs G)
  traverseAbs traverseG k t = traverseG (suc k) t

Traverse : (World → Set) → Set₁
Traverse F = ∀ {α β M} (M-applicative : Cat.RawApplicative M)
               (traverseName : Tr↑ Name M α β) → Tr↑ F M α β

Morph : (World → Set) → Set₁
Morph F = ∀ {α β M} (M-applicative : Cat.RawApplicative M)
            (θ : Name α → M (Name β))
          → F α → M (F β)

Var : (World → Set) → Set
Var F = ∀ {α} → Name α → F α

MapNames : (World → Set) → Set
MapNames F = ∀ {α β} → (Name α → Name β) → (F α → F β)

MapNames? : (World → Set) → Set
MapNames? F = ∀ {α β} → Partial Name α β → Partial F α β

Coe : (World → Set) → Set
Coe F = ∀ {α β} → α ⊆ β → F α → F β

Importø : (World → Set) → Set
Importø F = ∀ {α} → F ø → F α

Export : (World → Set) → Set
Export F = ∀ {α} k → Partial F (α ↑ k) α

Close : (World → Set) → Set
Close F = ∀ {α} → Partial F α ø

-- generalizes NaPa.lift↑1
lift↑ : ∀ {α β} k → (Name α → Name β) → (Name (α ↑ k) → Name (β ↑ k))
lift↑ k f x
    with x <nm k
...    | inj₁ x′ = x′ because⟨ ⊆-cong-↑ ⊆-ø-bottom k ⟩
...    | inj₂ x′ = add↑ k (f (x′ ∸nm k))

-- generalizes lift↑
substVar : ∀ {F} → Var F → Shift F → SubstVar F
substVar varF shiftF θ k x
    with x <nm k
...    | inj₁ x′ = varF     (coe (⊆-cong-↑ ⊆-ø-bottom k) x′)
...    | inj₂ x′ = shiftF k (⊆-+-↑ k) (θ (x′ ∸nm k))

module MapNamesGen {F} (mapNames : MapNames F) where
  shift : Shift F
  shift k ⊆w = mapNames (coe ⊆w ∘ add k)

  importø : Importø F
  importø = mapNames Nameø-elim

  coeF : Coe F
  coeF = mapNames ∘ coe

module MorphGen {F} (morph : Morph F) where
  mapNames : MapNames F
  mapNames = morph idAppli

  mapNames? : MapNames? F
  mapNames? = morph maybeAppli

  export : Export F
  export = mapNames? ∘ exportName

  close : Close F
  close = mapNames? (const nothing)

  open MapNamesGen mapNames public

module TraverseGen {F} (traverse : Traverse F) where
  morph : Morph F
  morph {M = M} appli θ = traverse appli (substVar pure shift θ) 
    where open Cat.RawApplicative appli using (pure; _<$>_)
          open MapNamesGen _<$>_        using (shift)

  open MorphGen morph public

-- This is more of a generic example
module Substitution {α β} where

  SuF : ∀ F → Set
  SuF F = Su↑ F F α β

  ×-subst : ∀ {F G} → SuF F → SuF G → SuF (F |×| G)
  ×-subst substF substG k (x , y) = (substF k x , substG k y)

  substAbs : ∀ {F} → SuF F → SuF (SynAbs F)
  substAbs suF = suF ∘ suc

αEq : (F : World → Set) (α β : World) → Set
αEq F α β = F α → F β → Bool

module αEquivalence {α β} where

  αEq↑ : (F : World → Set) → Set
  αEq↑ F = ∀ k → αEq F (α ↑ k) (β ↑ k)

  αeqName : αEq Name α β → αEq↑ Name
  αeqName Γ k x y with x <nm k | y <nm k
  ... | inj₁ x′ | inj₁ y′ = x′ ==nm y′
  ... | inj₂ x′ | inj₂ y′ = Γ (x′ ∸nm k) (y′ ∸nm k)
  ... | _       | _        = false

  -- α-equivalence on pairs is just the conjunction
  ×-αeq : ∀ {F G} → αEq↑ F → αEq↑ G → αEq↑ (F |×| G)
  ×-αeq αeqF αeqG k (x₁ , y₁) (x₂ , y₂) = αeqF k x₁ x₂ ∧ αeqG k y₁ y₂

  -- α-equivalence on abstraction is just incrementing `k`
  αeqAbs : ∀ {F} → αEq↑ F → αEq↑ (SynAbs F)
  αeqAbs αeqF = αeqF ∘ suc