{-# OPTIONS --safe --no-import-sorts #-}

open import Axiom.Set using (Theory)

module Axiom.Set.TotalMap (th : Theory) where

open import abstract-set-theory.Prelude hiding (lookup; map)

open import Data.Product.Properties     using (Σ-≡,≡→≡)
open import Axiom.Set.Map th            using (left-unique; Map ; mapWithKey-uniq ; left-unique-mapˢ)
open import Axiom.Set.Rel th            using (Rel ; dom ; dom∈)

open Theory th
open Equivalence

private variable A B : Type

-- defines a total map for a given set
total : Rel A B → Type
total R = ∀ {a} → a ∈ dom R

record TotalMap (A B : Type) : Type where
  field rel             : Set (A × B)
        left-unique-rel : left-unique rel
        total-rel       : total rel

  toMap : Map A B
  toMap = rel , left-unique-rel

  lookup : A → B
  lookup _ = proj₁ (from dom∈ total-rel)

  -- verify that lookup is what we expect
  lookup∈rel : {a : A} → (a , lookup a) ∈ rel
  lookup∈rel = proj₂ (from dom∈ total-rel)

  -- this is useful for proving equalities involving lookup
  ∈-rel⇒lookup-≡ : {a : A}{b : B} → (a , b) ∈ rel → lookup a ≡ b
  ∈-rel⇒lookup-≡ ab∈rel = sym (left-unique-rel ab∈rel (proj₂ (from dom∈ total-rel)))

open TotalMap

module Update ⦃ _ : DecEq A ⦄ where

  private
    updateFn : A × B → A → B → B
    updateFn (a , b) x y with (x ≟ a)
    ... | yes  _ = b
    ... | no   _ = y

  updateFn-id : {a : A} {b b' : B} → b ≡ updateFn (a , b) a b'
  updateFn-id {a = a} with (a ≟ a)
  ... | yes _ = refl
  ... | no ¬p = ⊥-elim (¬p refl)

  mapWithKey : {B' : Type} → (A → B → B') → TotalMap A B → TotalMap A B'
  mapWithKey f tm .rel              = map (λ{(x , y) → x , f x y}) (rel tm)
  mapWithKey _ tm .left-unique-rel  = mapWithKey-uniq (left-unique-rel tm)
  mapWithKey _ tm .total-rel        = ∈-map′ (∈-map′ (proj₂ (from dom∈ (total-rel tm))))

  update : A → B → TotalMap A B → TotalMap A B
  update a b = mapWithKey (updateFn (a , b))


module LookupUpdate {X : Set A} {a : A} {b : B} {a∈X : a ∈ X} ⦃ _ : DecEq A ⦄ where

  open Update

  ∈-rel-update : (tm : TotalMap A B) → (a , b) ∈ rel (update a b tm)
  ∈-rel-update tm = to ∈-map ((a , lookup tm a) , Σ-≡,≡→≡ (refl , updateFn-id {A = A}) , lookup∈rel tm)

  lookup-update-id : (tm : TotalMap A B) → lookup (update a b tm) a ≡ b
  lookup-update-id tm = ∈-rel⇒lookup-≡ (update _ _ tm) (∈-rel-update tm)



------------------------------------------------------
-- Correspondences between total maps and functions --

module FunTot (X : Set A) (⋁A≡X : isMaximal X) where

  Fun⇒Map : ∀ {f : A → B} (X : Set A) → Map A B
  Fun⇒Map {f = f} X = map (λ x → (x , f x)) X , left-unique-mapˢ X


  Fun⇒TotalMap : (f : A → B) → TotalMap A B
  Fun⇒TotalMap f .rel              = map (λ x → (x , f x)) X
  Fun⇒TotalMap _ .left-unique-rel  = left-unique-mapˢ X
  Fun⇒TotalMap _ .total-rel        = ∈-map′ (∈-map′ ⋁A≡X)

  Fun∈TotalMap : {f : A → B} {a : A}
               → a ∈ X → (a , f a) ∈ rel (Fun⇒TotalMap f)
  Fun∈TotalMap a∈X = ∈-map′ a∈X

  lookup∘Fun⇒TotalMap-id : {f : A → B}{a : A}
                         → lookup (Fun⇒TotalMap f) a ≡ f a
  lookup∘Fun⇒TotalMap-id {f = f} = ∈-rel⇒lookup-≡ ((Fun⇒TotalMap f)) (Fun∈TotalMap ⋁A≡X)