{-# OPTIONS --cubical --no-positivity-check --no-termination-check #-}

open import Prelude
open import Algebra
open import Algebra.Monus
open import Relation.Binary
open import Data.Maybe
open import Data.List using (List; _∷_; []; foldr)

module Control.Monad.HeapT
  {β„“} {𝑆 : Type β„“} {𝐹 : 𝑆 β†’ Type β„“ β†’ Type β„“}
  (monoid : Monoid 𝑆)
  (gmon : GradedMonad monoid 𝐹)
  where

open Monoid monoid
open GradedMonad gmon

private
  variable
    w : 𝑆

infixr 5 _∷_
infixr 6 _β‹Š_

mutual
  record Root (A : Type β„“) : Type β„“ where
    coinductive
    constructor _β‹Š_
    field
      weight : 𝑆
      step : 𝐹 weight (Branch A)

  data Node (A : Type β„“) : Type β„“ where
    ⌊_βŒ‹ : A β†’ Node A
    ⌈_βŒ‰ : Root A β†’ Node A

  data Branch (A : Type β„“) : Type β„“ where
    []  : Branch A
    _∷_ : Node A β†’ 𝐹 Ξ΅ (Branch A) β†’ Branch A
open Root public

Heap : Type β„“ β†’ Type β„“
Heap A = 𝐹 Ρ (Branch A)

infixr 5 _++_
_++_ : 𝐹 w (Branch A) β†’ 𝐹 Ξ΅ (Branch A) β†’ 𝐹 w (Branch A)
xs ++ ys =
  xs >>=Ξ΅ Ξ»  {  []       β†’ ys
             ;  (x ∷ xs) β†’ return (x ∷ xs ++ ys) }

infixr 1 _>>=α΄Ί_ _>>=α΄΄_
mutual
  _>>=α΄Ί_ : Node A β†’ (A β†’ Heap B) β†’ Heap B
  ⌊ x βŒ‹  >>=α΄Ί f = f x
  ⌈ x βŒ‰  >>=α΄Ί f = return (⌈ weight x β‹Š (step x >>=α΄΄ f) βŒ‰ ∷ return [])

  _>>=α΄΄_ : 𝐹 w (Branch A) β†’ (A β†’ Heap B) β†’ 𝐹 w (Branch B)
  xs >>=α΄΄ f =
    xs >>=Ξ΅ Ξ»  {  []        β†’ return []
               ;  (x ∷ xs)  β†’ (x >>=α΄Ί f) ++ (xs >>=α΄΄ f) }

returnα΄΄ : A β†’ Heap A
returnα΄΄ x = return (⌊ x βŒ‹ ∷ return [])

liftα΄΄ : 𝐹 w A β†’ Heap A
liftα΄΄ xs = return (⌈ _ β‹Š map (Ξ» x β†’ ⌊ x βŒ‹ ∷ return []) xs βŒ‰ ∷ return [])

flatten : 𝐹 w (Branch A) β†’ 𝐹 w (List A Γ— List (Root A))
flatten xs =
  xs >>=Ξ΅ Ξ»  {  []            β†’ return ([] , [])
             ;  (⌊ x βŒ‹ ∷ xs)  β†’ map (map₁ (x ∷_)) (flatten xs)
             ;  (⌈ x βŒ‰ ∷ xs)  β†’ map (mapβ‚‚ (x ∷_)) (flatten xs) }

module PopMin
  (_≀|β‰₯_ : Total (Ξ» x y β†’ βˆƒ z Γ— (y ≑ x βˆ™ z)))
  (decomp : βˆ€ {A B w₁ wβ‚‚ w₃} β†’ 𝐹 (w₁ βˆ™ wβ‚‚) A β†’ 𝐹 (w₁ βˆ™ w₃) B β†’ 𝐹 w₁ (𝐹 wβ‚‚ A Γ— 𝐹 w₃ B)) where

  _βˆͺ_ : Root A β†’ Root A β†’ Root A
  xs βˆͺ ys with weight xs ≀|β‰₯ weight ys
  ... | inl (k , wʸ≑wΛ£βˆ™k) = weight xs β‹Š map (Ξ» { (xs , ys) β†’ ⌈ k β‹Š ys βŒ‰ ∷ xs }) (decomp (subst (flip 𝐹 _) (sym (βˆ™Ξ΅ _)) (step xs)) (subst (flip 𝐹 _) wʸ≑wΛ£βˆ™k (step ys)))
  ... | inr (k , wˣ≑wΚΈβˆ™k) = weight ys β‹Š map (Ξ» { (ys , xs) β†’ ⌈ k β‹Š xs βŒ‰ ∷ ys }) (decomp (subst (flip 𝐹 _) (sym (βˆ™Ξ΅ _)) (step ys)) (subst (flip 𝐹 _) wˣ≑wΚΈβˆ™k (step xs)))

  ⋃⁺ : Root A β†’ List (Root A) β†’ Root A
  ⋃⁺ x₁ []             = x₁
  ⋃⁺ x₁ (xβ‚‚ ∷ [])      = x₁ βˆͺ xβ‚‚
  ⋃⁺ x₁ (xβ‚‚ ∷ x₃ ∷ xs) = (x₁ βˆͺ xβ‚‚) βˆͺ ⋃⁺ x₃ xs

  ⋃ : List (Root A) β†’ Maybe (Root A)
  ⋃ []       = nothing
  ⋃ (x ∷ xs) = just (⋃⁺ x xs)

  popMin : 𝐹 w (Branch A) β†’ 𝐹 w (List A Γ— Maybe (Root A))
  popMin = map (mapβ‚‚ ⋃) ∘ flatten