{-# OPTIONS --without-K --safe #-}

open import Algebra

module Data.FingerTree.Measures
  {r m}
  (β„³ : Monoid r m)
  where

open import Level using (_βŠ”_)

open Monoid β„³ renaming (Carrier to 𝓑)
open import Data.List as List using (List; _∷_; [])
open import Data.Product
open import Function

-- | A measure.
record Οƒ {a} (Ξ£ : Set a) : Set (a βŠ” r) where field ΞΌ : Ξ£ β†’ 𝓑
open Οƒ ⦃ ... ⦄
{-# DISPLAY Οƒ.ΞΌ _ = ΞΌ #-}

instance
  Οƒ-List : βˆ€ {a} {Ξ£ : Set a} β†’ ⦃ _ : Οƒ Ξ£ ⦄ β†’ Οƒ (List Ξ£)
  ΞΌ ⦃ Οƒ-List ⦄ = List.foldr (_βˆ™_ ∘ ΞΌ) Ξ΅


-- A "fiber" (I think) from the ΞΌ function.
--
-- μ⟨ Ξ£ βŸ©β‰ˆ 𝓂 means "There exists a Ξ£ such that ΞΌ Ξ£ β‰ˆ 𝓂"
infixl 2 _⇑[_]
record μ⟨_βŸ©β‰ˆ_ {a} (Ξ£ : Set a) ⦃ _ : Οƒ Ξ£ ⦄ (𝓂 : 𝓑) : Set (a βŠ” r βŠ” m) where
  constructor _⇑[_]
  field
    𝓒 : Ξ£
    𝒻 : ΞΌ 𝓒 β‰ˆ 𝓂
open μ⟨_βŸ©β‰ˆ_ public

-- Construct a measured value without any transformations of the measure.
infixl 2 _⇑
_⇑ : βˆ€ {a} {Ξ£ : Set a} ⦃ _ : Οƒ Ξ£ ⦄ (𝓒 : Ξ£) β†’ μ⟨ Ξ£ βŸ©β‰ˆ ΞΌ 𝓒
𝓒 (x ⇑) = x
𝒻 (x ⇑) = refl
{-# INLINE _⇑ #-}

-- These combinators allow for a kind of easoning syntax over the measures.
-- The first is used like so:
--
--   xs β‰ˆ[ assoc _ _ _ ]
--
-- Which will have the type:
--
--   μ⟨ Ξ£ βŸ©β‰ˆ x βˆ™ y βˆ™ z β†’ μ⟨ Ξ£ βŸ©β‰ˆ (x βˆ™ y) βˆ™ z
--
-- The second does the same:
--
--   xs β‰ˆ[ assoc _ _ _ ]β€²
--
-- However, when used in a chain, it typechecks after the ones used to its
-- left. This means you can call the solver on the left.
--
--   xs β‰ˆ[ β„³ β†― ] β‰ˆ[ assoc _ _ _ ]β€²
infixl 2 _β‰ˆ[_] β‰ˆ-rev
_β‰ˆ[_] : βˆ€ {a} {Ξ£ : Set a} ⦃ _ : Οƒ Ξ£ ⦄ {x : 𝓑} β†’ μ⟨ Ξ£ βŸ©β‰ˆ x β†’ βˆ€ {y} β†’ x β‰ˆ y β†’ μ⟨ Ξ£ βŸ©β‰ˆ y
𝓒 (xs β‰ˆ[ yβ‰ˆz ]) = 𝓒 xs
𝒻 (xs β‰ˆ[ yβ‰ˆz ]) = trans (𝒻 xs) yβ‰ˆz
{-# INLINE _β‰ˆ[_] #-}

β‰ˆ-rev : βˆ€ {a} {Ξ£ : Set a} ⦃ _ : Οƒ Ξ£ ⦄ {x : 𝓑} β†’ βˆ€ {y} β†’ x β‰ˆ y β†’ μ⟨ Ξ£ βŸ©β‰ˆ x β†’ μ⟨ Ξ£ βŸ©β‰ˆ y
𝓒 (β‰ˆ-rev yβ‰ˆz xs) = 𝓒 xs
𝒻 (β‰ˆ-rev yβ‰ˆz xs) = trans (𝒻 xs) yβ‰ˆz
{-# INLINE β‰ˆ-rev #-}

syntax β‰ˆ-rev yβ‰ˆz x↦y = x↦y β‰ˆ[ yβ‰ˆz ]β€²

infixr 2 β‰ˆ-right
β‰ˆ-right : βˆ€ {a} {Ξ£ : Set a} ⦃ _ : Οƒ Ξ£ ⦄ {x : 𝓑} β†’ μ⟨ Ξ£ βŸ©β‰ˆ x β†’ βˆ€ {y} β†’ x β‰ˆ y β†’ μ⟨ Ξ£ βŸ©β‰ˆ y
β‰ˆ-right (x ⇑[ xβ‰ˆy ]) yβ‰ˆz = x ⇑[ trans xβ‰ˆy yβ‰ˆz ]

syntax β‰ˆ-right x xβ‰ˆ = [ xβ‰ˆ ]β‰ˆ x

infixr 1 _ↀ_
-- A memoized application of ΞΌ
record βŸͺ_⟫ {a} (Ξ£ : Set a) ⦃ _ : Οƒ Ξ£ ⦄ : Set (a βŠ” r βŠ” m) where
  constructor _ↀ_
  field
    𝔐 : 𝓑
    𝓕 : μ⟨ Ξ£ βŸ©β‰ˆ 𝔐
open βŸͺ_⟫ public

-- Construct the memoized version
βŸͺ_β‡“βŸ« : βˆ€ {a} {Ξ£ : Set a} ⦃ _ : Οƒ Ξ£ ⦄ β†’ Ξ£ β†’ βŸͺ Ξ£ ⟫
𝔐 βŸͺ x β‡“βŸ« = ΞΌ x
𝓕 βŸͺ x β‡“βŸ« = x ⇑

instance
  Οƒ-βŸͺ⟫ : βˆ€ {a} {Ξ£ : Set a} ⦃ _ : Οƒ Ξ£ ⦄ β†’ Οƒ βŸͺ Ξ£ ⟫
  ΞΌ ⦃ Οƒ-βŸͺ⟫ ⦄ = 𝔐

open import Algebra.FunctionProperties _β‰ˆ_

-- This section allows us to use the do-notation to clean up proofs.
-- First, we construct arguments:
infixl 2 arg-syntax
record Arg {a} (Ξ£ : Set a) ⦃ _ : Οƒ Ξ£ ⦄ (𝓂 : 𝓑) (f : 𝓑 β†’ 𝓑) : Set (m βŠ” r βŠ” a) where
  constructor arg-syntax
  field
    ⟨f⟩ : Congruent₁ f
    arg : μ⟨ Ξ£ βŸ©β‰ˆ 𝓂
open Arg

syntax arg-syntax (Ξ» sz β†’ e₁) xs = xs [ e₁ ⟿ sz ]
-- This syntax is meant to be used like so:
--
--   do x ← xs [ a βˆ™> (s <βˆ™ b) ⟿ s ]
--
-- And it means "the size of the variable I'm binding here will be stored in this
-- part of the expression". See, for example, the listToTree function:
--
--   listToTree [] = empty ⇑
--   listToTree (x ∷ xs) = [ β„³ β†― ]β‰ˆ do
--     ys ← listToTree xs [ ΞΌ x βˆ™> s ⟿ s ]
--     x β—‚ ys
--
infixl 1 _>>=_
_>>=_ : βˆ€ {a b} {Σ₁ : Set a} {Ξ£β‚‚ : Set b} ⦃ _ : Οƒ Σ₁ ⦄ ⦃ _ : Οƒ Ξ£β‚‚ ⦄ {𝓂 f}
      β†’ Arg Σ₁ 𝓂 f
      β†’ ((x : Σ₁) β†’ ⦃ xβ‰ˆ : ΞΌ x β‰ˆ 𝓂 ⦄ β†’ μ⟨ Ξ£β‚‚ βŸ©β‰ˆ f (ΞΌ x))
      β†’ μ⟨ Ξ£β‚‚ βŸ©β‰ˆ f 𝓂
arg-syntax cng xs >>= k = k (𝓒 xs) ⦃ 𝒻 xs ⦄ β‰ˆ[ cng (𝒻 xs) ]
{-# INLINE _>>=_ #-}

-- Inside the lambda generated by do notation, we can only pass one argument.
-- So, to provide the proof (if needed)
_β‰ˆ?_ : βˆ€ x y β†’ ⦃ xβ‰ˆy : x β‰ˆ y ⦄ β†’ x β‰ˆ y
_β‰ˆ?_ _ _ ⦃ xβ‰ˆy ⦄ = xβ‰ˆy