{-# OPTIONS --cubical --safe #-}

open import Algebra
open import Relation.Binary
open import Algebra.Monus
open import Level

module Data.MonoidalHeap.Monad {s} {๐‘† : Type s} (monus : TMAPOM ๐‘†) where

open TMAPOM monus

open import Prelude
open import Data.List using (List; _โˆท_; []; foldr; _++_)
import Data.Nat as โ„•
import Data.Nat.Properties as โ„•

๐’ฎ : Type s
๐’ฎ = ๐‘† โ†’ ๐‘†

โŸฆ_โ‡‘โŸง : ๐‘† โ†’ ๐’ฎ
โŸฆ_โ‡‘โŸง = _โˆ™_

โŸฆ_โ‡“โŸง : ๐’ฎ โ†’ ๐‘†
โŸฆ x โ‡“โŸง = x ฮต

infixl 10 _โŠ™_
_โŠ™_ : (๐‘† โ†’ A) โ†’ ๐‘† โ†’ ๐‘† โ†’ A
f โŠ™ x = ฮป y โ†’ f (x โˆ™ y)

mutual
  data Node (V : ๐‘† โ†’ Type a) : Type (a โ„“โŠ” s) where
    leaf : V ฮต โ†’ Node V
    _โ‹Š_ : (w : ๐‘†) โ†’ Heap (V โŠ™ w) โ†’ Node V

  Heap : (๐‘† โ†’ Type a) โ†’ Type (a โ„“โŠ” s)
  Heap V = List (Node V)

private
  variable
    v : Level
    V : ๐‘† โ†’ Type v

Root : (๐‘† โ†’ Type v) โ†’ Type _
Root V = โˆƒ w ร— Heap (V โŠ™ w)

partition : Heap V โ†’ List (V ฮต) ร— List (Root V)
partition = foldr f ([] , [])
  where
    f : Node V โ†’ List (V ฮต) ร— List (โˆƒ w ร— Heap (V โŠ™ w)) โ†’ List (V ฮต) ร— List (โˆƒ w ร— Heap (V โŠ™ w))
    f (leaf x) (ls , hs) = (x โˆท ls) , hs
    f (w โ‹Š x)  (ls , hs) = ls , ((w , x) โˆท hs)


module _ {V : ๐‘† โ†’ Type v} where
  โŠ™-assoc : โˆ€ x y k โ†’ x โ‰ก y โˆ™ k โ†’ V โŠ™ x โ‰ก V โŠ™ y โŠ™ k
  โŠ™-assoc x y k xโ‰กyโˆ™k i z = V ((cong (_โˆ™ z) xโ‰กyโˆ™k อพ assoc y k z) i)

  โŠ™ฮต : V โŠ™ ฮต โ‰ก V
  โŠ™ฮต i x = V (ฮตโˆ™ x i)

  โŠ™-rassoc : โˆ€ x y โ†’ V โŠ™ (x โˆ™ y) โ‰ก V โŠ™ x โŠ™ y
  โŠ™-rassoc x y i z = V (assoc x y z i)

  merge : Root V โ†’ Root V โ†’ Root V
  merge (x , xs) (y , ys) with x โ‰ค|โ‰ฅ y
  ... | inl (k , yโ‰กxโˆ™k) = x , (k โ‹Š subst Heap (โŠ™-assoc y x k yโ‰กxโˆ™k) ys) โˆท xs
  ... | inr (k , xโ‰กyโˆ™k) = y , (k โ‹Š subst Heap (โŠ™-assoc x y k xโ‰กyโˆ™k) xs) โˆท ys

  mergesโบ : Root V โ†’ List (Root V) โ†’ Root V
  mergesโบ x [] = x
  mergesโบ xโ‚ (xโ‚‚ โˆท []) = merge xโ‚ xโ‚‚
  mergesโบ xโ‚ (xโ‚‚ โˆท xโ‚ƒ โˆท xs) = merge (merge xโ‚ xโ‚‚) (mergesโบ xโ‚ƒ xs)

  merges : List (Root V) โ†’ Maybe (Root V)
  merges [] = nothing
  merges (x โˆท xs) = just (mergesโบ x xs)

  popMin : Heap V โ†’ List (V ฮต) ร— Maybe (Root V)
  popMin = mapโ‚‚ merges โˆ˜ partition

  return : V ฮต โ†’ Heap V
  return x = leaf x โˆท []

  weight : โˆƒ x ร— V x โ†’ Heap V
  weight (w , x) = (w โ‹Š (leaf (subst V (sym (โˆ™ฮต w)) x) โˆท [])) โˆท []

_>>=_ : Heap V โ†’ (โˆ€ {x} โ†’ V x โ†’ Heap (V โŠ™ x)) โ†’ Heap V
_>>=_ []             k = []
_>>=_ (leaf x โˆท xs)  k = subst Heap โŠ™ฮต (k x) ++ (xs >>= k)
_>>=_ {V = V} ((w โ‹Š x) โˆท xs) k = (w โ‹Š (x >>= (subst Heap (โŠ™-rassoc {V = V} w _) โˆ˜ k))) โˆท (xs >>= k)