open import Lattice
open import Relation.Binary.PropositionalEquality as Eq
    using (_≡_;refl; sym; trans; cong; subst)
open import Agda.Primitive using (Level) renaming (_⊔_ to _⊔ℓ_)
open import Data.List using (List; _∷_; [])

module Lattice.FiniteMap {a b : Level} {A : Set a} {B : Set b}
    {_≈₂_ : B → B → Set b}
    {_⊔₂_ : B → B → B} {_⊓₂_ : B → B → B}
    (≡-dec-A : IsDecidable (_≡_ {a} {A}))
    (lB : IsLattice B _≈₂_ _⊔₂_ _⊓₂_) where

open IsLattice lB using () renaming (_≼_ to _≼₂_)
open import Lattice.Map ≡-dec-A lB as Map
    using (Map; ⊔-equal-keys; ⊓-equal-keys)
    renaming
        ( _≈_ to _≈ᵐ_
        ; _⊔_ to _⊔ᵐ_
        ; _⊓_ to _⊓ᵐ_
        ; ≈-equiv to ≈ᵐ-equiv
        ; ≈-⊔-cong to ≈ᵐ-⊔ᵐ-cong
        ; ⊔-assoc to ⊔ᵐ-assoc
        ; ⊔-comm to ⊔ᵐ-comm
        ; ⊔-idemp to ⊔ᵐ-idemp
        ; ≈-⊓-cong to ≈ᵐ-⊓ᵐ-cong
        ; ⊓-assoc to ⊓ᵐ-assoc
        ; ⊓-comm to ⊓ᵐ-comm
        ; ⊓-idemp to ⊓ᵐ-idemp
        ; absorb-⊔-⊓ to absorb-⊔ᵐ-⊓ᵐ
        ; absorb-⊓-⊔ to absorb-⊓ᵐ-⊔ᵐ
        ; ≈-dec to ≈ᵐ-dec
        ; _[_] to _[_]ᵐ
        ; m₁≼m₂⇒m₁[k]≼m₂[k] to m₁≼m₂⇒m₁[k]ᵐ≼m₂[k]ᵐ
        ; locate to locateᵐ
        ; keys to keysᵐ
        ; _updating_via_ to _updatingᵐ_via_
        ; updating-via-keys-≡ to updatingᵐ-via-keys-≡
        ; f'-Monotonic to f'-Monotonicᵐ
        ; _≼_ to _≼ᵐ_
        ; ∈k-dec to ∈k-decᵐ
        )
open import Data.List.Membership.Propositional using () renaming (_∈_ to _∈ˡ_)
open import Data.Product using (_×_; _,_; Σ; proj₁ ; proj₂)
open import Equivalence
open import Function using (_∘_)
open import Relation.Nullary using (¬_; Dec; yes; no)
open import Utils using (Pairwise; _∷_; [])
open import Data.Empty using (⊥-elim)
open import Showable using (Showable; show)

module WithKeys (ks : List A) where
    FiniteMap : Set (a ⊔ℓ b)
    FiniteMap = Σ Map (λ m → Map.keys m ≡ ks)

    instance
        showable : {{ showableA : Showable A }} {{ showableB : Showable B }} →
                   Showable FiniteMap
        showable = record { show = λ (m₁ , _) → show m₁ }

    _≈_ : FiniteMap → FiniteMap → Set (a ⊔ℓ b)
    _≈_ (m₁ , _) (m₂ , _) = m₁ ≈ᵐ m₂

    ≈₂-dec⇒≈-dec : IsDecidable _≈₂_ → IsDecidable _≈_
    ≈₂-dec⇒≈-dec ≈₂-dec fm₁ fm₂ = ≈ᵐ-dec ≈₂-dec (proj₁ fm₁) (proj₁ fm₂)

    _⊔_ : FiniteMap → FiniteMap → FiniteMap
    _⊔_ (m₁ , km₁≡ks) (m₂ , km₂≡ks) =
        ( m₁ ⊔ᵐ m₂
        , trans (sym (⊔-equal-keys {m₁} {m₂} (trans (km₁≡ks) (sym km₂≡ks))))
                km₁≡ks
        )

    _⊓_ : FiniteMap → FiniteMap → FiniteMap
    _⊓_ (m₁ , km₁≡ks) (m₂ , km₂≡ks) =
        ( m₁ ⊓ᵐ m₂
        , trans (sym (⊓-equal-keys {m₁} {m₂} (trans (km₁≡ks) (sym km₂≡ks))))
                km₁≡ks
        )

    _∈_ : A × B → FiniteMap → Set (a ⊔ℓ b)
    _∈_ k,v (m₁ , _) = k,v ∈ˡ (proj₁ m₁)

    _∈k_ : A → FiniteMap → Set a
    _∈k_ k (m₁ , _) = k ∈ˡ (keysᵐ m₁)

    ∈k-dec = ∈k-decᵐ

    locate : ∀ {k : A} {fm : FiniteMap} → k ∈k fm → Σ B (λ v → (k , v) ∈ fm)
    locate {k} {fm = (m , _)} k∈kfm = locateᵐ {k} {m} k∈kfm

    _updating_via_ : FiniteMap → List A → (A → B) → FiniteMap
    _updating_via_ (m₁ , ksm₁≡ks) ks f =
        ( m₁ updatingᵐ ks via f
        , trans (sym (updatingᵐ-via-keys-≡ m₁ ks f)) ksm₁≡ks
        )

    _[_] : FiniteMap → List A → List B
    _[_] (m₁ , _) ks = m₁ [ ks ]ᵐ

    ≈-equiv : IsEquivalence FiniteMap _≈_
    ≈-equiv = record
        { ≈-refl =
            λ {(m , _)} → IsEquivalence.≈-refl ≈ᵐ-equiv {m}
        ; ≈-sym =
            λ {(m₁ , _)} {(m₂ , _)} → IsEquivalence.≈-sym ≈ᵐ-equiv {m₁} {m₂}
        ; ≈-trans =
            λ {(m₁ , _)} {(m₂ , _)} {(m₃ , _)} →
                IsEquivalence.≈-trans ≈ᵐ-equiv {m₁} {m₂} {m₃}
        }

    isUnionSemilattice : IsSemilattice FiniteMap _≈_ _⊔_
    isUnionSemilattice = record
        { ≈-equiv = ≈-equiv
        ; ≈-⊔-cong =
            λ {(m₁ , _)} {(m₂ , _)} {(m₃ , _)} {(m₄ , _)} m₁≈m₂ m₃≈m₄ →
                ≈ᵐ-⊔ᵐ-cong {m₁} {m₂} {m₃} {m₄} m₁≈m₂ m₃≈m₄
        ; ⊔-assoc = λ (m₁ , _) (m₂ , _) (m₃ , _) → ⊔ᵐ-assoc m₁ m₂ m₃
        ; ⊔-comm = λ (m₁ , _) (m₂ , _) → ⊔ᵐ-comm m₁ m₂
        ; ⊔-idemp = λ (m , _) → ⊔ᵐ-idemp m
        }

    isIntersectSemilattice : IsSemilattice FiniteMap _≈_ _⊓_
    isIntersectSemilattice = record
        { ≈-equiv = ≈-equiv
        ; ≈-⊔-cong =
            λ {(m₁ , _)} {(m₂ , _)} {(m₃ , _)} {(m₄ , _)} m₁≈m₂ m₃≈m₄ →
                ≈ᵐ-⊓ᵐ-cong {m₁} {m₂} {m₃} {m₄} m₁≈m₂ m₃≈m₄
        ; ⊔-assoc = λ (m₁ , _) (m₂ , _) (m₃ , _) → ⊓ᵐ-assoc m₁ m₂ m₃
        ; ⊔-comm = λ (m₁ , _) (m₂ , _) → ⊓ᵐ-comm m₁ m₂
        ; ⊔-idemp = λ (m , _) → ⊓ᵐ-idemp m
        }

    isLattice : IsLattice FiniteMap _≈_ _⊔_ _⊓_
    isLattice = record
        { joinSemilattice = isUnionSemilattice
        ; meetSemilattice = isIntersectSemilattice
        ; absorb-⊔-⊓ = λ (m₁ , _) (m₂ , _) → absorb-⊔ᵐ-⊓ᵐ m₁ m₂
        ; absorb-⊓-⊔ = λ (m₁ , _) (m₂ , _) → absorb-⊓ᵐ-⊔ᵐ m₁ m₂
        }

    open IsLattice isLattice using (_≼_; ⊔-Monotonicˡ; ⊔-Monotonicʳ) public

    lattice : Lattice FiniteMap
    lattice = record
        { _≈_ = _≈_
        ; _⊔_ = _⊔_
        ; _⊓_ = _⊓_
        ; isLattice = isLattice
        }

    m₁≼m₂⇒m₁[k]≼m₂[k] : ∀ (fm₁ fm₂ : FiniteMap) {k : A} {v₁ v₂ : B} →
                        fm₁ ≼ fm₂ → (k , v₁) ∈ fm₁ → (k , v₂) ∈ fm₂ → v₁ ≼₂ v₂
    m₁≼m₂⇒m₁[k]≼m₂[k] (m₁ , _) (m₂ , _) m₁≼m₂ k,v₁∈m₁ k,v₂∈m₂ = m₁≼m₂⇒m₁[k]ᵐ≼m₂[k]ᵐ m₁ m₂ m₁≼m₂ k,v₁∈m₁ k,v₂∈m₂

    module GeneralizedUpdate
             {l} {L : Set l}
             {_≈ˡ_ : L → L → Set l} {_⊔ˡ_ : L → L → L} {_⊓ˡ_ : L → L → L}
             (lL : IsLattice L _≈ˡ_ _⊔ˡ_ _⊓ˡ_)
             (f : L → FiniteMap) (f-Monotonic : Monotonic (IsLattice._≼_ lL) _≼_ f)
             (g : A → L → B) (g-Monotonicʳ : ∀ k → Monotonic (IsLattice._≼_ lL) _≼₂_ (g k))
             (ks : List A) where

        open IsLattice lL using () renaming (_≼_ to _≼ˡ_)

        updater : L → A → B
        updater l k = g k l

        f' : L → FiniteMap
        f' l = (f l) updating ks via (updater l)

        f'-Monotonic : Monotonic _≼ˡ_ _≼_ f'
        f'-Monotonic {l₁} {l₂} l₁≼l₂ = f'-Monotonicᵐ lL (proj₁ ∘ f) f-Monotonic g g-Monotonicʳ ks l₁≼l₂

    all-equal-keys : ∀ (fm₁ fm₂ : FiniteMap) → (Map.keys (proj₁ fm₁) ≡ Map.keys (proj₁ fm₂))
    all-equal-keys (fm₁ , km₁≡ks) (fm₂ , km₂≡ks) = trans km₁≡ks (sym km₂≡ks)

    ∈k-exclusive : ∀ (fm₁ fm₂ : FiniteMap) {k : A} → ¬ ((k ∈k fm₁) × (¬ k ∈k fm₂))
    ∈k-exclusive fm₁ fm₂ {k} (k∈kfm₁ , k∉kfm₂) =
        let
            k∈kfm₂ = subst (λ l → k ∈ˡ l) (all-equal-keys fm₁ fm₂) k∈kfm₁
        in
            k∉kfm₂ k∈kfm₂

    m₁≼m₂⇒m₁[ks]≼m₂[ks] : ∀ (fm₁ fm₂ : FiniteMap) (ks' : List A) →
                          fm₁ ≼ fm₂ → Pairwise _≼₂_ (fm₁ [ ks' ]) (fm₂ [ ks' ])
    m₁≼m₂⇒m₁[ks]≼m₂[ks] _ _ [] _ = []
    m₁≼m₂⇒m₁[ks]≼m₂[ks] fm₁@(m₁ , km₁≡ks) fm₂@(m₂ , km₂≡ks) (k ∷ ks'') m₁≼m₂
        with ∈k-decᵐ k (proj₁ m₁) | ∈k-decᵐ k (proj₁ m₂)
    ...   | yes k∈km₁ | yes k∈km₂ =
            let
                (v₁ , k,v₁∈m₁) = locateᵐ {m = m₁} k∈km₁
                (v₂ , k,v₂∈m₂) = locateᵐ {m = m₂} k∈km₂
            in
                (m₁≼m₂⇒m₁[k]ᵐ≼m₂[k]ᵐ m₁ m₂ m₁≼m₂ k,v₁∈m₁ k,v₂∈m₂) ∷ m₁≼m₂⇒m₁[ks]≼m₂[ks] fm₁ fm₂ ks'' m₁≼m₂
    ...   | no k∉km₁ | no k∉km₂ = m₁≼m₂⇒m₁[ks]≼m₂[ks] fm₁ fm₂ ks'' m₁≼m₂
    ...   | yes k∈km₁ | no k∉km₂ = ⊥-elim (∈k-exclusive fm₁ fm₂ (k∈km₁ , k∉km₂))
    ...   | no k∉km₁ | yes k∈km₂ = ⊥-elim (∈k-exclusive fm₂ fm₁ (k∈km₂ , k∉km₁))

open WithKeys public