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

open import Algebra
open import Level

module Control.Monad.Weighted.Monad {β„“} {𝑅 : Type β„“} (rng : Semiring 𝑅) where

open Semiring rng

open import Level
open import Path
open import Path.Reasoning
open import HLevels
open import Data.Sigma

open import Control.Monad.Weighted.Definition rng
open import Control.Monad.Weighted.Eliminators rng
open import Control.Monad.Weighted.Union rng
open import Control.Monad.Weighted.Cond rng
open import Control.Monad.Weighted.Functor rng
open import Control.Monad.Weighted.Functor.TypeDef
\end{code}
%<*bind-alg>
\begin{code}
bind-alg :  (A β†’ Weighted B) β†’
            Ξ¦ A (Weighted B)
bind-alg f .fst (p β—ƒ x ∷ _ ⟨ xs ⟩) = (p β‹Š f x) βˆͺ xs
bind-alg f .fst [] = []
\end{code}
%</bind-alg>
\begin{code}
bind-alg f .snd .c-set _ = trunc
bind-alg f .snd .c-del x _ xs = cong (_βˆͺ xs) (0β‹Š (f x))
bind-alg f .snd .c-dup p q x _ xs =
  p β‹Š f x βˆͺ (q β‹Š f x βˆͺ xs)  β‰‘βŸ¨ βˆͺ-assoc (p β‹Š f x) (q β‹Š f x) xs ⟩
  (p β‹Š f x βˆͺ q β‹Š f x) βˆͺ xs  β‰‘βŸ¨ cong (_βˆͺ xs) (β‹Š-distribΚ³ p q (f x)) ⟩
  (p + q β‹Š f x) βˆͺ xs ∎
bind-alg f .snd .c-com p x q y _ xs =
  p β‹Š f x βˆͺ q β‹Š f y βˆͺ xs    β‰‘βŸ¨ βˆͺ-assoc (p β‹Š f x) (q β‹Š f y) xs ⟩
  (p β‹Š f x βˆͺ q β‹Š f y) βˆͺ xs  β‰‘βŸ¨ cong (_βˆͺ xs) (βˆͺ-com (p β‹Š f x) (q β‹Š f y)) ⟩
  (q β‹Š f y βˆͺ p β‹Š f x) βˆͺ xs  β‰‘Λ˜βŸ¨ βˆͺ-assoc (q β‹Š f y) (p β‹Š f x) xs ⟩
  q β‹Š f y βˆͺ p β‹Š f x βˆͺ xs ∎
\end{code}
%<*bind-impl>
\begin{code}
_>>=_ :  Weighted A β†’ (A β†’ Weighted B) β†’
         Weighted B
xs >>= f = ⟦ bind-alg f ⟧ xs
\end{code}
%</bind-impl>
%<*pure>
\begin{code}
pure : A β†’ Weighted A
pure x = 1# β—ƒ x ∷ []
\end{code}
%</pure>
\begin{code}
_>>_ : Weighted A β†’ Weighted B β†’ Weighted B
xs >> ys = do
  _ ← xs
  ys

_<*>_ : Weighted (A β†’ B) β†’ Weighted A β†’ Weighted B
fs <*> xs = do
  f ← fs
  x ← xs
  pure (f x)

>>=-distrib : (xs ys : Weighted A) (g : A β†’ Weighted B) β†’ (xs βˆͺ ys) >>= g ≑ (xs >>= g) βˆͺ (ys >>= g)
>>=-distrib xs ys g = ⟦ >>=-distribβ€² ys g ⟧ xs
  where
  >>=-distribβ€² : (ys : Weighted A) (g : A β†’ Weighted B) β†’ Ξ¨[ xs ⦂ A ] ((xs βˆͺ ys) >>= g) ≑ (xs >>= g) βˆͺ (ys >>= g)
  >>=-distribβ€² ys g .fst [] = refl
  >>=-distribβ€² ys g .fst (p β—ƒ x ∷ xs ⟨ P ⟩) =
    (((p β—ƒ x ∷ xs) βˆͺ ys) >>= g)          β‰‘βŸ¨βŸ©
    (p β—ƒ x ∷ xs βˆͺ ys) >>= g              β‰‘βŸ¨βŸ©
    p β‹Š g x βˆͺ ((xs βˆͺ ys) >>= g)          β‰‘βŸ¨ cong (p β‹Š g x βˆͺ_) P ⟩
    p β‹Š g x βˆͺ ((xs >>= g) βˆͺ (ys >>= g))  β‰‘βŸ¨ βˆͺ-assoc (p β‹Š g x) (xs >>= g) (ys >>= g) ⟩
    (p β‹Š g x βˆͺ (xs >>= g)) βˆͺ (ys >>= g)  β‰‘βŸ¨βŸ©
    ((p β—ƒ x ∷ xs) >>= g) βˆͺ (ys >>= g) ∎
  >>=-distribβ€² ys g .snd = eq-coh

β‹Š-assoc->>= : βˆ€ p (xs : Weighted A) (f : A β†’ Weighted B) β†’ (p β‹Š xs) >>= f ≑ p β‹Š (xs >>= f)
β‹Š-assoc->>= p xs f = ⟦ β‹Š-assoc->>=β€² p f ⟧ xs
  where
  β‹Š-assoc->>=β€² : βˆ€ p (f : A β†’ Weighted B) β†’ Ξ¨[ xs ⦂ A ] (p β‹Š xs) >>= f ≑ p β‹Š (xs >>= f)
  β‹Š-assoc->>=β€² p f .fst [] = refl
  β‹Š-assoc->>=β€² p f .fst (q β—ƒ x ∷ xs ⟨ P ⟩) =
    (p β‹Š (q β—ƒ x ∷ xs)) >>= f            β‰‘βŸ¨βŸ©
    (p * q β—ƒ x ∷ p β‹Š xs) >>= f          β‰‘βŸ¨βŸ©
    ((p * q) β‹Š f x) βˆͺ ((p β‹Š xs) >>= f)  β‰‘βŸ¨ cong (((p * q) β‹Š f x) βˆͺ_) P ⟩
    ((p * q) β‹Š f x) βˆͺ (p β‹Š (xs >>= f))  β‰‘βŸ¨ cong (_βˆͺ (p β‹Š (xs >>= f))) (*-assoc-β‹Š p q (f x)) ⟩
    (p β‹Š (q β‹Š f x)) βˆͺ (p β‹Š (xs >>= f))  β‰‘βŸ¨ β‹Š-distribΛ‘ p (q β‹Š f x) (xs >>= f) ⟩
    p β‹Š ((q β—ƒ x ∷ xs) >>= f) ∎
  β‹Š-assoc->>=β€² p f .snd = eq-coh

>>=-idΛ‘ : (x : A) β†’ (f : A β†’ Weighted B)
      β†’ (pure x >>= f) ≑ f x
>>=-idΛ‘ x f =
  pure x >>= f         β‰‘βŸ¨βŸ©
  (1# β—ƒ x ∷ []) >>= f  β‰‘βŸ¨βŸ©
  1# β‹Š f x βˆͺ [] >>= f  β‰‘βŸ¨βŸ©
  1# β‹Š f x βˆͺ []        β‰‘βŸ¨ βˆͺ-idr (1# β‹Š f x) ⟩
  1# β‹Š f x             β‰‘βŸ¨ 1β‹Š (f x) ⟩
  f x ∎

>>=-idΚ³ : (xs : Weighted A) β†’ xs >>= pure ≑ xs
>>=-idΚ³ = ⟦ >>=-idΚ³β€² ⟧
  where
  >>=-idΚ³β€² : Ξ¨[ xs ⦂ A ] xs >>= pure ≑ xs
  >>=-idΚ³β€² .fst [] = refl
  >>=-idΚ³β€² .fst (p β—ƒ x ∷ xs ⟨ P ⟩) =
    ((p β—ƒ x ∷ xs) >>= pure) β‰‘βŸ¨βŸ©
    p β‹Š (pure x) βˆͺ (xs >>= pure) β‰‘βŸ¨βŸ©
    p β‹Š (1# β—ƒ x ∷ []) βˆͺ (xs >>= pure) β‰‘βŸ¨βŸ©
    p * 1# β—ƒ x ∷ [] βˆͺ (xs >>= pure) β‰‘βŸ¨βŸ©
    p * 1# β—ƒ x ∷ (xs >>= pure) β‰‘βŸ¨ cong (_β—ƒ x ∷ (xs >>= pure)) (*1 p) ⟩
    p β—ƒ x ∷ xs >>= pure β‰‘βŸ¨ cong (p β—ƒ x ∷_) P ⟩
    p β—ƒ x ∷ xs ∎
  >>=-idΚ³β€² .snd = eq-coh

>>=-assoc : (xs : Weighted A) β†’ (f : A β†’ Weighted B) β†’ (g : B β†’ Weighted C)
      β†’ ((xs >>= f) >>= g) ≑ xs >>= (Ξ» x β†’ f x >>= g)
>>=-assoc xs f g = ⟦ >>=-assocβ€² f g ⟧ xs
  where
  >>=-assocβ€² : (f : A β†’ Weighted B) β†’ (g : B β†’ Weighted C) β†’ Ξ¨[ xs ⦂ A ] ((xs >>= f) >>= g) ≑ xs >>= (Ξ» x β†’ f x >>= g)
  >>=-assocβ€² f g .fst [] = refl
  >>=-assocβ€² f g .fst (p β—ƒ x ∷ xs ⟨ P ⟩) =
    (((p β—ƒ x ∷ xs) >>= f) >>= g) β‰‘βŸ¨βŸ©
    ((p β‹Š f x βˆͺ (xs >>= f)) >>= g) β‰‘βŸ¨ >>=-distrib (p β‹Š f x) (xs >>= f) g ⟩
    ((p β‹Š f x) >>= g) βˆͺ ((xs >>= f) >>= g) β‰‘βŸ¨ cong ((p β‹Š f x) >>= g βˆͺ_) P ⟩
    ((p β‹Š f x) >>= g) βˆͺ (xs >>= (Ξ» y β†’ f y >>= g)) β‰‘βŸ¨ cong (_βˆͺ (xs >>= (Ξ» y β†’ f y >>= g))) (β‹Š-assoc->>= p (f x) g) ⟩
    p β‹Š (f x >>= g) βˆͺ (xs >>= (Ξ» y β†’ f y >>= g)) β‰‘βŸ¨βŸ©
    ((p β—ƒ x ∷ xs) >>= (Ξ» y β†’ f y >>= g)) ∎
  >>=-assocβ€² f g .snd = eq-coh
\end{code}