-- Because iterated products currently require both A and B to be of the same
-- universe, and the FiniteMap is written in a universe-polymorphic way,
-- specialize the FiniteMap module with Set-typed types only.

open import Lattice
open import Equivalence
open import Relation.Binary.PropositionalEquality as Eq using (_≡_; refl; sym; trans; cong; subst)
open import Relation.Binary.Definitions using (Decidable)
open import Agda.Primitive using (Level) renaming (_⊔_ to _⊔ℓ_)
open import Function.Definitions using (Inverseˡ; Inverseʳ)

module Lattice.FiniteValueMap (A : Set) (B : Set)
    (_≈₂_ : B → B → Set)
    (_⊔₂_ : B → B → B) (_⊓₂_ : B → B → B)
    (≈-dec-A : Decidable (_≡_ {_} {A}))
    (lB : IsLattice B _≈₂_ _⊔₂_ _⊓₂_) where

open import Data.List using (List; length; []; _∷_)
open import Utils using (Unique; push; empty)
open import Data.Product using (_,_)
open import Data.List.Properties using (∷-injectiveʳ)
open import Data.List.Relation.Unary.All using (All)
open import Data.List.Relation.Unary.Any using (Any; here; there)
open import Relation.Nullary using (¬_)

open import Lattice.Map A B _≈₂_ _⊔₂_ _⊓₂_ ≈-dec-A lB using (subset-impl)
open import Lattice.FiniteMap A B _≈₂_ _⊔₂_ _⊓₂_ ≈-dec-A lB public

module IterProdIsomorphism where
    open import Data.Unit using (⊤; tt)
    open import Lattice.Unit using () renaming (_≈_ to _≈ᵘ_; _⊔_ to _⊔ᵘ_; _⊓_ to _⊓ᵘ_; ≈-dec to ≈ᵘ-dec; isLattice to isLatticeᵘ)
    open import Lattice.IterProd _≈₂_ _≈ᵘ_ _⊔₂_ _⊔ᵘ_ _⊓₂_ _⊓ᵘ_ lB isLatticeᵘ as IP using (IterProd)

    from : ∀ {ks : List A} → FiniteMap ks → IterProd (length ks)
    from {[]} (([] , _) , _) = tt
    from {k ∷ ks'} (((k' , v) ∷ kvs' , push _ uks') , refl) = (v , from ((kvs' , uks'), refl))

    to : ∀ {ks : List A} → Unique ks → IterProd (length ks) → FiniteMap ks
    to {[]} _ ⊤ = (([] , empty) , refl)
    to {k ∷ ks'} (push k≢ks' uks') (v , rest)
        with to uks' rest
    ...   | ((kvs' , ukvs') , kvs'≡ks') =
            let
                -- This would be easier if we pattern matched on the equiality proof
                -- to get refl, but that makes it harder to reason about 'to' when
                -- the arguments are not known to be refl.
                k≢kvs' = subst (λ ks → All (λ k' → ¬ k ≡ k') ks) (sym kvs'≡ks') k≢ks'
                kvs≡ks = cong (k ∷_) kvs'≡ks'
            in
                (((k , v) ∷ kvs' , push k≢kvs' ukvs') , kvs≡ks)


    private
        _≈ᵐ_ : ∀ {ks : List A} → FiniteMap ks → FiniteMap ks → Set
        _≈ᵐ_ {ks} = _≈_ ks

        _≈ⁱᵖ_ : ∀ {ks : List A} → IterProd (length ks) → IterProd (length ks) → Set
        _≈ⁱᵖ_ {ks} = IP._≈_ (length ks)

    from-to-inverseˡ : ∀ {ks : List A} (uks : Unique ks) →
                      Inverseˡ (_≈ᵐ_ {ks}) (_≈ⁱᵖ_ {ks}) (from {ks}) (to {ks} uks) -- from (to x) = x
    from-to-inverseˡ {[]} _ _ = IsEquivalence.≈-refl (IP.≈-equiv 0)
    from-to-inverseˡ {k ∷ ks'} (push k≢ks' uks') (v , rest)
        with ((kvs' , ukvs') , refl) ← to uks' rest in p rewrite sym p =
            (IsLattice.≈-refl lB , from-to-inverseˡ {ks'} uks' rest)
            -- the rewrite here is needed because the IH is in terms of `to uks' rest`,
            -- but we end up with the 'unpacked' form (kvs', ...). So, put it back
            -- in the 'packed' form after we've performed enough inspection
            -- to know we take the cons branch of `to`.

    -- The map has its own uniqueness proof, but the call to 'to' needs a standalone
    -- uniqueness proof too. Work with both proofs as needed to thread things through.
    from-to-inverseʳ : ∀ {ks : List A} (uks : Unique ks) →
                       Inverseʳ (_≈ᵐ_ {ks}) (_≈ⁱᵖ_ {ks}) (from {ks}) (to {ks} uks) -- to (from x) = x
    from-to-inverseʳ {[]} _ (([] , empty) , kvs≡ks) rewrite kvs≡ks = ((λ k v ()) , (λ k v ()))
    from-to-inverseʳ {k ∷ ks'} uks@(push k≢ks'₁ uks'₁) fm₁@(m₁@((k , v) ∷ kvs'₁ , push k≢ks'₂ uks'₂) , refl)
        with to uks'₁ (from ((kvs'₁ , uks'₂) , refl)) | from-to-inverseʳ {ks'} uks'₁ ((kvs'₁ , uks'₂) , refl)
    ...   | ((kvs'₂ , ukvs'₂) , _) | (kvs'₂⊆kvs'₁ , kvs'₁⊆kvs'₂) = (m₂⊆m₁ , m₁⊆m₂)
                where
                    kvs₁ = (k , v) ∷ kvs'₁
                    kvs₂ = (k , v) ∷ kvs'₂

                    m₁⊆m₂ : subset-impl kvs₁ kvs₂
                    m₁⊆m₂ k' v' (here refl) = (v' , (IsLattice.≈-refl lB , here refl))
                    m₁⊆m₂ k' v' (there k',v'∈kvs'₁) =
                        let (v'' , (v'≈v'' , k',v''∈kvs'₂)) = kvs'₁⊆kvs'₂ k' v' k',v'∈kvs'₁
                        in (v'' , (v'≈v'' , there k',v''∈kvs'₂))

                    m₂⊆m₁ : subset-impl kvs₂ kvs₁
                    m₂⊆m₁ k' v' (here refl) = (v' , (IsLattice.≈-refl lB , here refl))
                    m₂⊆m₁ k' v' (there k',v'∈kvs'₂) =
                        let (v'' , (v'≈v'' , k',v''∈kvs'₁)) = kvs'₂⊆kvs'₁ k' v' k',v'∈kvs'₂
                        in (v'' , (v'≈v'' , there k',v''∈kvs'₁))