\begin{code}
{-# OPTIONS --safe #-}

open import Algebra
open import Prelude
open import Path.Reasoning

module Data.Weighted.Monad {β„“} {𝑆 : Type β„“} (weight : Weight 𝑆) where

open Weight weight
open import Data.Weighted.Definition βŠ•-semigroup
open import Data.Weighted.CommutativeMonoid βŠ•-semigroup
open import Data.Weighted.Eliminators βŠ•-semigroup
open import Data.Weighted.Condition βŠ•-semigroup
open import Data.Weighted.Functor



module _ {A : Type a} {B : Type b} (f : A β†’ B) where
  wmap-alg : A ⟢W B
  wmap-alg = record
    { act-w = Ξ» w x β†’ w β–Ή f x ∷ βŸ…βŸ†
    ; coh-w = Ξ» p q x β†’ dup p q (f x) βŸ…βŸ†
    }

  wmap : Weighted A β†’ Weighted B
  wmap = [ wmap-alg ]W↓

wmap-id : (xs : Weighted A) β†’ wmap id xs ≑ xs
wmap-id = ⟦ alg ⟧
  where
  alg : Ξ¨[ xs ⦂ Weighted A ] ↦ (wmap id xs ≑ xs)
  alg .fst βŸ…βŸ† = refl
  alg .fst (w β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) = cong (w β–Ή x ∷_) P⟨xs⟩
  alg .snd = eq-coh

wmap-comp : (f : B β†’ C) (g : A β†’ B) (xs : Weighted A) β†’ wmap f (wmap g xs) ≑ wmap (f ∘ g) xs
wmap-comp f g = ⟦ alg f g ⟧
  where
  alg : (f : B β†’ C) (g : A β†’ B) β†’ Ξ¨[ xs ⦂ Weighted A ] ↦ wmap f (wmap g xs) ≑ wmap (f ∘ g) xs
  alg f g .fst βŸ…βŸ† = refl
  alg f g .fst (w β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) = cong (w β–Ή _ ∷_) P⟨xs⟩
  alg f g .snd = eq-coh

functorWeighted : Functor (Weighted {a = a})
functorWeighted .Functor.map = wmap
functorWeighted .Functor.map-id = wmap-id
functorWeighted .Functor.map-comp = wmap-comp

private module Display-β‹Š where
  _βˆ™_ = _βŠ—_

  infixr 5 _β‹Š_
  _β‹Š_ : 𝑆 β†’ Weighted A β†’ Weighted A
\end{code}
%<*rtimes>
\begin{code}[number=eqn:rtimes]
  w β‹Š βŸ…βŸ†          = βŸ…βŸ†
  w β‹Š p β–Ή x ∷ xs  = w βˆ™ p β–Ή x ∷ w β‹Š xs
\end{code}
%</rtimes>
\begin{code}

  w β‹Š com p x q y xs i = com (w βŠ— p) x (w βŠ— q) y (w β‹Š xs) i
  w β‹Š dup p q x xs i = (dup (w βŠ— p) (w βŠ— q) x (w β‹Š xs) ΝΎ cong (_β–Ή x ∷ w β‹Š xs) (sym (βŠ—βŸ¨βŠ•βŸ© w p q))) i
  w β‹Š trunc xs ys p q i j = trunc (w β‹Š xs) (w β‹Š ys) (cong (w β‹Š_) p) (cong (w β‹Š_) q) i j

β‹Š-alg : 𝑆 β†’ A ⟢W A
β‹Š-alg w = cond-alg (w βŠ—_) (Ξ» p q β†’ sym (βŠ—βŸ¨βŠ•βŸ© w p q))

infixr 5 _β‹Š_
_β‹Š_ : 𝑆 β†’ Weighted A β†’ Weighted A
w β‹Š xs = cond (w βŠ—_) (Ξ» p q β†’ sym (βŠ—βŸ¨βŠ•βŸ© w p q)) xs

wmap-β‹Š : βˆ€ w (f : A β†’ B) xs β†’ w β‹Š wmap f xs ≑ wmap f (w β‹Š xs)
wmap-β‹Š w f = ⟦ alg w f ⟧
  where
  alg : βˆ€ w (f : A β†’ B) β†’ Ξ¨[ xs ⦂ Weighted A ] ↦ w β‹Š wmap f xs ≑ wmap f (w β‹Š xs)
  alg w f .snd = eq-coh
  alg w f .fst βŸ…βŸ† = refl
  alg w f .fst (p β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) = cong (w βŠ— p β–Ή f x ∷_) P⟨xs⟩

βŠ•βŸ¨β‹ŠβŸ© : βˆ€ p q (xs : Weighted A) β†’ (p β‹Š xs) βˆͺ (q β‹Š xs) ≑ (p βŠ• q) β‹Š xs
βŠ•βŸ¨β‹ŠβŸ© p q = ⟦ alg p q ⟧
  where
  alg : βˆ€ p q β†’ Ξ¨[ xs ⦂ Weighted A ] ↦ ((p β‹Š xs) βˆͺ (q β‹Š xs) ≑ (p βŠ• q) β‹Š xs)
  alg p q .fst βŸ…βŸ† = refl
  alg p q .fst (w β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) =
    (p β‹Š (w β–Ή x ∷ xs)) βˆͺ (q β‹Š (w β–Ή x ∷ xs)) β‰‘βŸ¨βŸ©
    (p βŠ— w β–Ή x ∷ p β‹Š xs) βˆͺ (q βŠ— w β–Ή x ∷ q β‹Š xs) β‰‘βŸ¨βŸ©
    p βŠ— w β–Ή x ∷ (p β‹Š xs) βˆͺ (q βŠ— w β–Ή x ∷ q β‹Š xs) β‰‘Λ˜βŸ¨ cong (p βŠ— w β–Ή x ∷_) (βˆͺ-cons (q βŠ— w) x (p β‹Š xs) _) ⟩
    p βŠ— w β–Ή x ∷ q βŠ— w β–Ή x ∷ (p β‹Š xs) βˆͺ (q β‹Š xs) β‰‘βŸ¨ dup (p βŠ— w) (q βŠ— w) x _ ⟩
    p βŠ— w βŠ• q βŠ— w β–Ή x ∷ (p β‹Š xs) βˆͺ (q β‹Š xs) β‰‘βŸ¨ congβ‚‚ (_β–Ή x ∷_) (sym (βŸ¨βŠ•βŸ©βŠ— _ _ _)) P⟨xs⟩ ⟩
    (p βŠ• q) β‹Š (w β–Ή x ∷ xs) ∎
  alg p q .snd = eq-coh

β‹ŠβŸ¨βˆͺ⟩ : βˆ€ w (x y : Weighted A) β†’ w β‹Š (x βˆͺ y) ≑ (w β‹Š x) βˆͺ (w β‹Š y)
β‹ŠβŸ¨βˆͺ⟩ w x y = ⟦ alg w y ⟧ x
  where
  alg : βˆ€ w y β†’ Ξ¨[ x ⦂ Weighted A ] ↦ (w β‹Š (x βˆͺ y) ≑ (w β‹Š x) βˆͺ (w β‹Š y))
  alg w y .snd = eq-coh
  alg w y .fst βŸ…βŸ† = refl
  alg w y .fst (u β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) = cong (w βŠ— u β–Ή x ∷_) P⟨xs⟩

private module InlineBind where
  _>>=_ : Weighted A β†’ (A β†’ Weighted B) β†’ Weighted B
\end{code}
%<*bind>
\begin{code}
  βŸ…βŸ†            >>= k = βŸ…βŸ†
  (p β–Ή x ∷ xs)  >>= k = (p β‹Š k x) βˆͺ (xs >>= k)
\end{code}
%</bind>
\begin{code}
  com p x q y xs i >>= k = begin[ i ]
    (p β‹Š k x) βˆͺ (q β‹Š k y)   βˆͺ (xs >>= k) β‰‘Λ˜βŸ¨ βˆͺ-assoc (p β‹Š k x) _ _ ⟩
    ((p β‹Š k x) βˆͺ (q β‹Š k y)) βˆͺ (xs >>= k) β‰‘βŸ¨ cong (_βˆͺ (xs >>= k)) (βˆͺ-com (p β‹Š k x) _) ⟩
    ((q β‹Š k y) βˆͺ (p β‹Š k x)) βˆͺ (xs >>= k) β‰‘βŸ¨ βˆͺ-assoc (q β‹Š k y) _ _ ⟩
    (q β‹Š k y) βˆͺ (p β‹Š k x)   βˆͺ (xs >>= k) ∎
  dup p q x xs i >>= k = begin[ i ]
    (p β‹Š k x) βˆͺ (q β‹Š k x) βˆͺ (xs >>= k) β‰‘Λ˜βŸ¨ βˆͺ-assoc (p β‹Š k x) _ _ ⟩
    ((p β‹Š k x) βˆͺ (q β‹Š k x)) βˆͺ (xs >>= k) β‰‘βŸ¨ cong (_βˆͺ (xs >>= k)) (βŠ•βŸ¨β‹ŠβŸ© p q (k x)) ⟩
    ((p βŠ• q) β‹Š k x) βˆͺ (xs >>= k) ∎
  trunc xs ys p q i j >>= k =
    trunc (xs >>= k) (ys >>= k) (cong (_>>= k) p) (cong (_>>= k) q) i j

module _ {A : Type a} {B : Type b} where
  bind-alg : (A β†’ Weighted B) β†’ A ⟢W B
  bind-alg k = record
    { act-w = Ξ» w x β†’ w β‹Š k x
    ; coh-w = Ξ» p q x β†’ βŠ•βŸ¨β‹ŠβŸ© p q (k x)
    }

  infixl 4.1 _>>=_
  _>>=_ : Weighted A β†’ (A β†’ Weighted B) β†’ Weighted B
  xs >>= k = [ bind-alg k ]W↓ xs

  >>=-βˆͺ : (xs ys : Weighted A) (k : A β†’ Weighted B) β†’ (xs βˆͺ ys) >>= k ≑ (xs >>= k) βˆͺ (ys >>= k)
  >>=-βˆͺ xs ys k = collect-βˆͺ (bind-alg k) xs ys

=<<-βˆͺ : (xs : Weighted A) (ys zs : A β†’ Weighted B) β†’ (xs >>= Ξ» x β†’ ys x βˆͺ zs x) ≑ (xs >>= ys) βˆͺ (xs >>= zs)
=<<-βˆͺ {B = B} xs ys zs = ⟦ alg ys zs ⟧ xs
  where
  lemma : (w x y z : Weighted B) β†’ (w βˆͺ x) βˆͺ (y βˆͺ z) ≑ (w βˆͺ y) βˆͺ (x βˆͺ z)
  lemma w x y z =
    (w βˆͺ x) βˆͺ (y βˆͺ z) β‰‘βŸ¨ βˆͺ-assoc w x _ ΝΎ cong (w βˆͺ_) (sym (βˆͺ-assoc x y z)) ⟩
    w βˆͺ (x βˆͺ y) βˆͺ z β‰‘βŸ¨ cong (Ξ» e β†’ w βˆͺ e βˆͺ z) (βˆͺ-com x y ) ⟩
    w βˆͺ (y βˆͺ x) βˆͺ z β‰‘Λ˜βŸ¨ βˆͺ-assoc w y _ ΝΎ cong (w βˆͺ_) (sym (βˆͺ-assoc y x z)) ⟩
    (w βˆͺ y) βˆͺ (x βˆͺ z) ∎
    

  alg : (ys zs : A β†’ Weighted B) β†’ Ξ¨[ xs ⦂ Weighted A ] ↦ (xs >>= Ξ» x β†’ ys x βˆͺ zs x) ≑ (xs >>= ys) βˆͺ (xs >>= zs)
  alg ys zs .snd = eq-coh
  alg ys zs .fst βŸ…βŸ† = refl
  alg ys zs .fst (w β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) =
    let k = Ξ» xβ€² β†’ ys xβ€² βˆͺ zs xβ€² in
    (w β–Ή x ∷ xs) >>= k β‰‘βŸ¨βŸ©
    (w β‹Š (ys x βˆͺ zs x)) βˆͺ (xs >>= k) β‰‘βŸ¨ congβ‚‚ _βˆͺ_ (β‹ŠβŸ¨βˆͺ⟩ w (ys x) (zs x)) P⟨xs⟩ ⟩ 
    ((w β‹Š ys x) βˆͺ (w β‹Š zs x)) βˆͺ ((xs >>= ys) βˆͺ (xs >>= zs))  β‰‘βŸ¨ lemma (w β‹Š ys x) (w β‹Š zs x) (xs >>= ys) (xs >>= zs) ⟩
    ((w β‹Š ys x) βˆͺ (xs >>= ys)) βˆͺ ((w β‹Š zs x) βˆͺ (xs >>= zs)) β‰‘βŸ¨βŸ©
    ((w β–Ή x ∷ xs) >>= ys) βˆͺ ((w β–Ή x ∷ xs) >>= zs) ∎


return : A β†’ Weighted A
return x = πŸ™ β–Ή x ∷ βŸ…βŸ†

β‹Š-assoc : βˆ€ w v (s : Weighted A) β†’ w β‹Š v β‹Š s ≑ (w βŠ— v) β‹Š s
β‹Š-assoc w v = ⟦ alg ⟧
  where
  alg : Ξ¨[ s ⦂ Weighted A ] ↦ w β‹Š v β‹Š s ≑ (w βŠ— v) β‹Š s
  alg .snd = eq-coh
  alg .fst βŸ…βŸ† = refl
  alg .fst (p β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) = congβ‚‚ (_β–Ή x ∷_) (sym (βŠ—-assoc w v p)) P⟨xs⟩

Ξ΅β‹Š : (x : Weighted A) β†’ πŸ™ β‹Š x ≑ x
Ξ΅β‹Š = ⟦ alg ⟧
  where
  alg : Ξ¨[ s ⦂ Weighted A ] ↦ πŸ™ β‹Š s ≑ s
  alg .snd = eq-coh
  alg .fst βŸ…βŸ† = refl
  alg .fst (p β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) = congβ‚‚ (_β–Ή x ∷_) (πŸ™βŠ— p) P⟨xs⟩

β‹ŠβŸ…βŸ† : βˆ€ w (x : Weighted A) β†’ w β‹Š x ≑ βŸ…βŸ† β†’ x ≑ βŸ…βŸ†
β‹ŠβŸ…βŸ† w = ⟦ alg ⟧
  where
  alg : Ξ¨[ x ⦂ Weighted A ] ↦ (w β‹Š x ≑ βŸ…βŸ† β†’ x ≑ βŸ…βŸ†)
  alg .fst βŸ…βŸ† _ = refl
  alg .fst (p β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) e = absurd (falseβ‰’true (cong null e))
  alg .snd = prop-coh λ _ → isProp→ (trunc _ _)

β‹Š->>= : βˆ€ p (xs : Weighted A) (k : A β†’ Weighted B) β†’ (p β‹Š xs) >>= k ≑ p β‹Š (xs >>= k)
β‹Š->>= p xs k = ⟦ alg p k ⟧ xs
  where
  alg : βˆ€ p (k : A β†’ Weighted B) β†’ Ξ¨[ xs ⦂ Weighted A ] ↦ (p β‹Š xs) >>= k ≑ p β‹Š (xs >>= k)
  alg p k .snd = eq-coh
  alg p k .fst βŸ…βŸ† = refl
  alg p k .fst (w β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) =
    ((p βŠ— w) β–Ή x ∷ (p β‹Š xs)) >>= k β‰‘βŸ¨βŸ©
    ((p βŠ— w) β‹Š k x) βˆͺ  ((p β‹Š xs) >>= k) β‰‘βŸ¨ congβ‚‚ _βˆͺ_ (sym (β‹Š-assoc p w (k x))) P⟨xs⟩ ⟩
    (p β‹Š (w β‹Š k x)) βˆͺ (p β‹Š (xs >>= k)) β‰‘Λ˜βŸ¨ β‹ŠβŸ¨βˆͺ⟩ p (w β‹Š k x) (xs >>= k) ⟩
    p β‹Š ((w β‹Š k x) βˆͺ (xs >>= k)) ∎

>>=-assoc  : (xs : Weighted A) (ys : A β†’ Weighted B) (zs : B β†’ Weighted C)
           β†’ (xs >>= ys) >>= zs ≑ xs >>= (Ξ» x β†’ ys x >>= zs)
>>=-assoc xs ys zs = ⟦ alg ys zs ⟧ xs
  where
  alg : (ys : A β†’ Weighted B) (zs : B β†’ Weighted C) β†’
        Ξ¨[ xs ⦂ Weighted A ] ↦ (xs >>= ys) >>= zs ≑ xs >>= (Ξ» x β†’ ys x >>= zs)
  alg ys zs .snd = eq-coh
  alg ys zs .fst βŸ…βŸ† = refl
  alg ys zs .fst (w β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) =
    (w β–Ή x ∷ xs) >>= ys >>= zs β‰‘βŸ¨βŸ©
    ((w β‹Š ys x) βˆͺ (xs >>= ys)) >>= zs β‰‘βŸ¨ >>=-βˆͺ (w β‹Š ys x) (xs >>= ys) zs ⟩
    ((w β‹Š ys x) >>= zs) βˆͺ ((xs >>= ys) >>= zs) β‰‘βŸ¨ congβ‚‚ _βˆͺ_ (β‹Š->>= w (ys x) zs) P⟨xs⟩ ⟩
    (w β‹Š (ys x >>= zs)) βˆͺ (xs >>= (Ξ» x β†’ ys x >>= zs)) β‰‘βŸ¨βŸ©
    (w β–Ή x ∷ xs) >>= (Ξ» x β†’ ys x >>= zs) ∎

>>=-idΛ‘ : (x : A) (k : A β†’ Weighted B) β†’ return x >>= k ≑ k x
>>=-idΛ‘ x k = βˆͺ-idΚ³ (πŸ™ β‹Š k x) ΝΎ Ξ΅β‹Š (k x)

>>=-idΚ³ : (xs : Weighted A) β†’ xs >>= return ≑ xs
>>=-idʳ = ⟦ alg ⟧
  where
  alg : Ξ¨[ xs ⦂ Weighted A ] ↦ (xs >>= return) ≑ xs
  alg .snd = eq-coh
  alg .fst βŸ…βŸ† = refl
  alg .fst (w β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) = congβ‚‚ (_β–Ή x ∷_) (βŠ—πŸ™ w) P⟨xs⟩

module _ {A : Type a} {B : Type b} where
  >>=-βŸ…βŸ† : (xs : Weighted A) β†’ xs >>= const βŸ…βŸ† ≑ (βŸ…βŸ† ⦂ Weighted B)
  >>=-βŸ…βŸ† = ⟦ alg ⟧
    where
    alg : Ξ¨[ xs ⦂ Weighted A ] ↦ (xs >>= const βŸ…βŸ† ≑ (βŸ…βŸ† ⦂ Weighted B))
    alg .snd = eq-coh
    alg .fst βŸ…βŸ† = refl
    alg .fst (w β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) = P⟨xs⟩

>>=-wmap : (f : A β†’ B) (k : B β†’ Weighted C) (xs : Weighted A) β†’ wmap f xs >>= k ≑ xs >>= (k ∘ f)
>>=-wmap f k = ⟦ alg f k ⟧
  where
  alg : (f : A β†’ B) (k : B β†’ Weighted C) β†’ Ξ¨[ xs ⦂ Weighted A ] ↦ wmap f xs >>= k ≑ xs >>= (k ∘ f)
  alg f k .snd = eq-coh
  alg f k .fst βŸ…βŸ† = refl
  alg f k .fst (w β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) = cong ((w β‹Š k (f x)) βˆͺ_) P⟨xs⟩

wmap->>= : (f : A β†’ B) (xs : Weighted A) β†’ wmap f xs ≑ xs >>= (return ∘ f)
wmap->>= f = ⟦ alg f ⟧
  where
  alg : (f : A β†’ B) β†’ Ξ¨[ xs ⦂ Weighted A ] ↦ wmap f xs ≑ xs >>= (return ∘ f)
  alg f .snd = eq-coh
  alg f .fst βŸ…βŸ† = refl
  alg f .fst (w β–Ή x ∷ xs ⟨ P⟨xs⟩ ⟩) = congβ‚‚ _β–Ή f x ∷_ (sym (βŠ—πŸ™ w)) P⟨xs⟩

\end{code}