-- 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 Data.Product using (Σ; proj₁; proj₂; _×_)
open import Data.Empty using (⊥-elim)
open import Utils using (Unique; push; empty; All¬-¬Any)
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; locate; forget; _∈_; Map-functional)
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ᵘ; ≈-equiv to ≈ᵘ-equiv)
    open import Lattice.IterProd _≈₂_ _≈ᵘ_ _⊔₂_ _⊔ᵘ_ _⊓₂_ _⊓ᵘ_ lB isLatticeᵘ as IP using (IterProd)
    open IsLattice lB using () renaming (≈-trans to ≈₂-trans; ≈-sym to ≈₂-sym)

    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₁ ks₂ : List A} → FiniteMap ks₁ → FiniteMap ks₂ → Set
        _⊆ᵐ_ fm₁ fm₂ = subset-impl (proj₁ (proj₁ fm₁)) (proj₁ (proj₁ fm₂))

        _≈ⁱᵖ_ : ∀ {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'₁))

    private
        first-key-in-map : ∀ {k : A} {ks : List A} (fm : FiniteMap (k ∷ ks)) → Σ B (λ v → (k , v) ∈ proj₁ fm)
        first-key-in-map (((k , v) ∷ _ , _) , refl) = (v , here refl)

        from-first-value : ∀ {k : A} {ks : List A} (fm : FiniteMap (k ∷ ks)) → proj₁ (from fm) ≈₂ proj₁ (first-key-in-map fm)
        from-first-value {k} {ks} (((k , v) ∷ _ , push _ _) , refl) = IsLattice.≈-refl lB

        pop : ∀ {k : A} {ks : List A} → FiniteMap (k ∷ ks) → FiniteMap ks
        pop (((_ ∷ kvs') , push _ ukvs') , refl) = ((kvs' , ukvs') , refl)

        pop-≈ :  ∀ {k : A} {ks : List A} (fm₁ fm₂ : FiniteMap (k ∷ ks)) → fm₁ ≈ᵐ fm₂ → pop fm₁ ≈ᵐ pop fm₂
        pop-≈ {k} {ks} fm₁ fm₂ (fm₁⊆fm₂ , fm₂⊆fm₁) = (narrow fm₁⊆fm₂ , narrow fm₂⊆fm₁)
            where
                narrow₁ : ∀ {fm₁ fm₂ : FiniteMap (k ∷ ks)} → fm₁ ⊆ᵐ fm₂  → pop fm₁ ⊆ᵐ fm₂
                narrow₁ {(_ ∷ _ , push _ _) , refl} kvs₁⊆kvs₂ k' v' k',v'∈kvs'₁ = kvs₁⊆kvs₂ k' v' (there k',v'∈kvs'₁)

                narrow₂ : ∀ {fm₁ : FiniteMap ks} {fm₂ : FiniteMap (k ∷ ks)} → fm₁ ⊆ᵐ fm₂  → fm₁ ⊆ᵐ pop fm₂
                narrow₂ {fm₁} {fm₂ = (_ ∷ kvs'₂ , push k≢ks _) , kvs≡ks@refl} kvs₁⊆kvs₂ k' v' k',v'∈kvs'₁
                    with kvs₁⊆kvs₂ k' v' k',v'∈kvs'₁
                ...   | (v'' , (v'≈v'' , here refl)) rewrite sym (proj₂ fm₁) = ⊥-elim (All¬-¬Any k≢ks (forget {m = proj₁ fm₁} k',v'∈kvs'₁))
                ...   | (v'' , (v'≈v'' , there k',v'∈kvs'₂)) = (v'' , (v'≈v'' , k',v'∈kvs'₂))

                narrow : ∀ {fm₁ fm₂ : FiniteMap (k ∷ ks)} → fm₁ ⊆ᵐ fm₂  → pop fm₁ ⊆ᵐ pop fm₂
                narrow {fm₁} {fm₂} x = narrow₂ {pop fm₁} (narrow₁ {fm₂ = fm₂} x)

        from-rest : ∀ {k : A} {ks : List A} (fm : FiniteMap (k ∷ ks)) → proj₂ (from fm) ≡ from (pop fm)
        from-rest (((_ ∷ kvs') , push _ ukvs') , refl) = refl

    from-preserves-≈ : ∀ {ks : List A} → (fm₁ fm₂ : FiniteMap ks) → fm₁ ≈ᵐ fm₂ → (_≈ⁱᵖ_ {ks}) (from fm₁) (from fm₂)
    from-preserves-≈ {[]} (([] , _) , _) (([] , _) , _) _ = IsEquivalence.≈-refl ≈ᵘ-equiv
    from-preserves-≈ {k ∷ ks'} fm₁@(m₁ , _) fm₂@(m₂ , _) fm₁≈fm₂@(kvs₁⊆kvs₂ , kvs₂⊆kvs₁)
        with first-key-in-map fm₁ | first-key-in-map fm₂ | from-first-value fm₁ | from-first-value fm₂
    ...   | (v₁ , k,v₁∈fm₁) | (v₂ , k,v₂∈fm₂) | fv₁≈v₁ | fv₂≈v₂
            with kvs₁⊆kvs₂ _ _ k,v₁∈fm₁
    ...       | (v₁' , (v₁≈v₁' , k,v₁'∈fm₂))
                rewrite Map-functional {m = m₂} k,v₂∈fm₂ k,v₁'∈fm₂
                rewrite from-rest fm₁ rewrite from-rest fm₂
                = (≈₂-trans fv₁≈v₁ (≈₂-trans v₁≈v₁' (≈₂-sym fv₂≈v₂)) , from-preserves-≈ (pop fm₁) (pop fm₂) (pop-≈ fm₁ fm₂ fm₁≈fm₂))