module multicategory where

open import Level
open import Data.List
open import Data.Product
open import Function
open import Agda.Builtin.Equality

data Pullback {ℓ} {a b c : Set ℓ} (f : a → c) (g : b → c) : Set ℓ
  where
  equation : (x : a) → (y : b) → f x ≡ g y → Pullback f g

pb₁ : ∀ {ℓ} {a b c : Set ℓ} {f : a → c} {g : b → c} → Pullback f g → a
pb₁ (equation a _ _) = a

pb₂ : ∀ {ℓ} {a b c : Set ℓ} {f : a → c} {g : b → c} → Pullback f g → b
pb₂ (equation _ b _) = b

data HList {ℓ} : List (Set ℓ) → Set ℓ
  where
  []ₕ : HList []
  _∷ₕ_ : ∀ {x} {xs} → x → HList xs → HList (x ∷ xs)

_++ₕ_ : ∀ {ℓ} {xs ys} → HList {ℓ} xs → HList ys → HList (xs ++ ys)
[]ₕ ++ₕ xs = xs
(x ∷ₕ xs) ++ₕ ys = x ∷ₕ (xs ++ₕ ys)

------------------------------------------------------

record Functor {ℓ} (f : Set ℓ → Set ℓ) : Set (suc ℓ)
  where
  field
    fmap : ∀ {a b} → (a → b) → f a → f b

open Functor {{...}}

_<$>_ : ∀ {ℓ} {f} {a b} {{ _ : Functor {ℓ} f }} → (a → b) → f a → f b
_<$>_ = fmap

record Monad {ℓ} (m : Set ℓ → Set ℓ) : Set (suc ℓ)
  where
  field
    {{functor}} : Functor {ℓ} m
    return : ∀ {a} → a → m a
    join : ∀ {a} → m (m a) → m a

open Monad {{...}}

_>>=_ : ∀ {ℓ} {m} {a b} {{ _ : Monad {ℓ} m }} → m a → (a → m b) → m b
ma >>= amb = join (amb <$> ma)

instance
  functorList : ∀ {ℓ} → Functor {ℓ} List
  fmap {{functorList}} _ [] = []
  fmap {{functorList}} f (x ∷ xs) = f x ∷ fmap f xs

  monadList : ∀ {ℓ} → Monad {ℓ} List
  return {{monadList}} = [_]
  join {{monadList}} = Data.List.concat

record MultiCategory {ℓ} (m : ∀ {κ} → Set κ → Set κ) (c₀ c₁ : Set ℓ) : Setω
  where
  field
    {{mm}} : Monad {ℓ} m
    dom : c₁ → m c₀
    cod : c₁ → c₀
    idₘ : ∀ {a : c₀} → Σ[ f ∈ c₁ ] (dom f ≡ return a × cod f ≡ a)
    _∘ₘ_ : (f : c₁) (gs : m c₁) (_ : dom f ≡ cod <$> gs) → Σ[ r ∈ c₁ ] (dom r ≡ gs >>= dom × cod r ≡ cod f)

open MultiCategory {{...}}

data Multifunction : Set₁
  where
  MFun : {xs : List Set} {y : Set} → (HList xs → y) → Multifunction

domMFun : Multifunction → List Set
domMFun (MFun {xs} _) = xs

codMFun : Multifunction → Set
codMFun (MFun {_} {y} _) = y

idMFun : ∀ {a : Set} → Σ[ f ∈ Multifunction ] (domMFun f ≡ return a × codMFun f ≡ a)
idMFun {a} = MFun f , refl , refl
  where
  f : HList [ a ] → a
  f (x ∷ₕ []ₕ) = x

split : ∀ {ℓ} {ys} → (xs : _) → HList {ℓ} (xs ++ ys) → HList xs × HList ys
split []        vs       = []ₕ , vs
split (x ∷ xs) (v ∷ₕ vs) = case (split xs vs) of λ { (vs₁ , vs₂) → (v ∷ₕ vs₁) , vs₂ }

contramap₁ : ∀ {ℓ} {xs ys} {y r : Set ℓ} → (HList xs → y) → (HList (y ∷ ys) → r) → HList (xs ++ ys) → r
contramap₁ {xs = xs} g f vs = case (split xs vs) of λ { (vs₁ , vs₂) → f (g vs₁ ∷ₕ vs₂) }

contramap₂ : ∀ {ℓ} {xs ys zs} {r : Set ℓ} → (HList {ℓ} zs → HList ys) → (HList (xs ++ ys) → r) → HList (xs ++ zs) → r
contramap₂ {xs = []}     g f vs = f (g vs)
contramap₂ {xs = x ∷ xs} g f (v ∷ₕ vs) = f (v ∷ₕ (case (split xs vs) of λ { (vs₁ , vs₂) → vs₁ ++ₕ (g vs₂) }))

unpack : (fs : List Multifunction) → HList (fs >>= domMFun) → HList (codMFun <$> fs)
unpack [] []ₕ = []ₕ
unpack (f@(MFun f') ∷ fs) xs = case (split (domMFun f) xs) of λ { (xs₁ , xs₂) → f' xs₁ ∷ₕ unpack fs xs₂ }

_∘ₘₙ_ : (f : Multifunction) (gs : List Multifunction) (_ : domMFun f ≡ codMFun <$> gs) → Σ[ r ∈ Multifunction ] (domMFun r ≡ gs >>= domMFun × codMFun r ≡ codMFun f)
_∘ₘₙ_ (MFun f) []            refl = MFun f , refl , refl
_∘ₘₙ_ (MFun f) (MFun g ∷ gs) refl = MFun (λ { xs → contramap₂ (unpack gs) (contramap₁ g f) xs }) , refl , refl

instance
  mcatMFun : MultiCategory List Set Multifunction
  mcatMFun = record
    { dom = domMFun
    ; cod = codMFun
    ; idₘ = idMFun
    ; _∘ₘ_ = _∘ₘₙ_
    }