module Analysis.Sign where

open import Data.String using (String) renaming (_≟_ to _≟ˢ_)
open import Data.Nat using (suc)
open import Data.Product using (_×_; proj₁; _,_)
open import Data.List using (List; _∷_; []; foldr; cartesianProduct; cartesianProductWith)
open import Data.List.Membership.Propositional as MemProp using () renaming (_∈_ to _∈ˡ_)
open import Relation.Binary.PropositionalEquality using (_≡_; refl; sym; trans; subst)
open import Relation.Nullary using (¬_; Dec; yes; no)
open import Data.Unit using (⊤)
open import Function using (_∘_)

open import Language
open import Lattice
open import Utils using (Pairwise)
open import Showable using (Showable; show)
import Lattice.FiniteValueMap

data Sign : Set where
    + : Sign
    - : Sign
    0ˢ : Sign

instance
    showable : Showable Sign
    showable = record
        { show = (λ
            { + → "+"
            ; - → "-"
            ; 0ˢ → "0"
            })
        }

-- g for siGn; s is used for strings and i is not very descriptive.
_≟ᵍ_ : IsDecidable (_≡_ {_} {Sign})
_≟ᵍ_ + + = yes refl
_≟ᵍ_ + - = no (λ ())
_≟ᵍ_ + 0ˢ = no (λ ())
_≟ᵍ_ - + = no (λ ())
_≟ᵍ_ - - = yes refl
_≟ᵍ_ - 0ˢ = no (λ ())
_≟ᵍ_ 0ˢ + = no (λ ())
_≟ᵍ_ 0ˢ - = no (λ ())
_≟ᵍ_ 0ˢ 0ˢ = yes refl

-- embelish 'sign' with a top and bottom element.
open import Lattice.AboveBelow Sign _≡_ (record { ≈-refl = refl; ≈-sym = sym; ≈-trans = trans }) _≟ᵍ_ as AB
    using ()
    renaming
        ( AboveBelow to SignLattice
        ; ≈-dec to ≈ᵍ-dec
        ; ⊥ to ⊥ᵍ
        ; ⊤ to ⊤ᵍ
        ; [_] to [_]ᵍ
        ; _≈_ to _≈ᵍ_
        ; ≈-⊥-⊥ to ≈ᵍ-⊥ᵍ-⊥ᵍ
        ; ≈-⊤-⊤ to ≈ᵍ-⊤ᵍ-⊤ᵍ
        ; ≈-lift to ≈ᵍ-lift
        ; ≈-refl to ≈ᵍ-refl
        )
-- 'sign' has no underlying lattice structure, so use the 'plain' above-below lattice.
open AB.Plain 0ˢ using ()
    renaming
        ( finiteHeightLattice to finiteHeightLatticeᵍ
        ; isLattice to isLatticeᵍ
        ; fixedHeight to fixedHeightᵍ
        ; _≼_ to _≼ᵍ_
        ; _⊔_ to _⊔ᵍ_
        )

open IsLattice isLatticeᵍ using ()
    renaming
        ( ≼-trans to ≼ᵍ-trans
        )

plus : SignLattice → SignLattice → SignLattice
plus ⊥ᵍ _ = ⊥ᵍ
plus _ ⊥ᵍ = ⊥ᵍ
plus ⊤ᵍ _ = ⊤ᵍ
plus _ ⊤ᵍ = ⊤ᵍ
plus [ + ]ᵍ [ + ]ᵍ = [ + ]ᵍ
plus [ + ]ᵍ [ - ]ᵍ = ⊤ᵍ
plus [ + ]ᵍ [ 0ˢ ]ᵍ = [ + ]ᵍ
plus [ - ]ᵍ [ + ]ᵍ = ⊤ᵍ
plus [ - ]ᵍ [ - ]ᵍ = [ - ]ᵍ
plus [ - ]ᵍ [ 0ˢ ]ᵍ = [ - ]ᵍ
plus [ 0ˢ ]ᵍ [ + ]ᵍ = [ + ]ᵍ
plus [ 0ˢ ]ᵍ [ - ]ᵍ = [ - ]ᵍ
plus [ 0ˢ ]ᵍ [ 0ˢ ]ᵍ = [ 0ˢ ]ᵍ

-- this is incredibly tedious: 125 cases per monotonicity proof, and tactics
-- are hard. postulate for now.
postulate plus-Monoˡ : ∀ (s₂ : SignLattice) → Monotonic _≼ᵍ_ _≼ᵍ_ (λ s₁ → plus s₁ s₂)
postulate plus-Monoʳ : ∀ (s₁ : SignLattice) → Monotonic _≼ᵍ_ _≼ᵍ_ (plus s₁)

minus : SignLattice → SignLattice → SignLattice
minus ⊥ᵍ _ = ⊥ᵍ
minus _ ⊥ᵍ = ⊥ᵍ
minus ⊤ᵍ _ = ⊤ᵍ
minus _ ⊤ᵍ = ⊤ᵍ
minus [ + ]ᵍ [ + ]ᵍ = ⊤ᵍ
minus [ + ]ᵍ [ - ]ᵍ = [ + ]ᵍ
minus [ + ]ᵍ [ 0ˢ ]ᵍ = [ + ]ᵍ
minus [ - ]ᵍ [ + ]ᵍ = [ - ]ᵍ
minus [ - ]ᵍ [ - ]ᵍ = ⊤ᵍ
minus [ - ]ᵍ [ 0ˢ ]ᵍ = [ - ]ᵍ
minus [ 0ˢ ]ᵍ [ + ]ᵍ = [ - ]ᵍ
minus [ 0ˢ ]ᵍ [ - ]ᵍ = [ + ]ᵍ
minus [ 0ˢ ]ᵍ [ 0ˢ ]ᵍ = [ 0ˢ ]ᵍ

postulate minus-Monoˡ : ∀ (s₂ : SignLattice) → Monotonic _≼ᵍ_ _≼ᵍ_ (λ s₁ → minus s₁ s₂)
postulate minus-Monoʳ : ∀ (s₁ : SignLattice) → Monotonic _≼ᵍ_ _≼ᵍ_ (minus s₁)

module WithProg (prog : Program) where
    open Program prog

    -- The variable -> sign map is a finite value-map with keys strings. Use a bundle to avoid explicitly specifying operators.
    module VariableSignsFiniteMap = Lattice.FiniteValueMap.WithKeys _≟ˢ_ isLatticeᵍ vars
    open VariableSignsFiniteMap
        using ()
        renaming
            ( FiniteMap to VariableSigns
            ; isLattice to isLatticeᵛ
            ; _≈_ to _≈ᵛ_
            ; _⊔_ to _⊔ᵛ_
            ; _≼_ to _≼ᵛ_
            ; ≈₂-dec⇒≈-dec to ≈ᵍ-dec⇒≈ᵛ-dec
            ; _∈_ to _∈ᵛ_
            ; _∈k_ to _∈kᵛ_
            ; _updating_via_ to _updatingᵛ_via_
            ; locate to locateᵛ
            ; m₁≼m₂⇒m₁[k]≼m₂[k] to m₁≼m₂⇒m₁[k]ᵛ≼m₂[k]ᵛ
            )
    open IsLattice isLatticeᵛ
        using ()
        renaming
            ( ⊔-Monotonicˡ to ⊔ᵛ-Monotonicˡ
            ; ⊔-Monotonicʳ to ⊔ᵛ-Monotonicʳ
            ; ⊔-idemp to ⊔ᵛ-idemp
            )
    open Lattice.FiniteValueMap.IterProdIsomorphism.WithUniqueKeysAndFixedHeight _≟ˢ_ isLatticeᵍ vars-Unique ≈ᵍ-dec _ fixedHeightᵍ
        using ()
        renaming
            ( isFiniteHeightLattice to isFiniteHeightLatticeᵛ
            )

    ≈ᵛ-dec = ≈ᵍ-dec⇒≈ᵛ-dec ≈ᵍ-dec
    joinSemilatticeᵛ = IsFiniteHeightLattice.joinSemilattice isFiniteHeightLatticeᵛ
    fixedHeightᵛ = IsFiniteHeightLattice.fixedHeight isFiniteHeightLatticeᵛ
    ⊥ᵛ = proj₁ (proj₁ (proj₁ fixedHeightᵛ))

    -- Finally, the map we care about is (state -> (variables -> sign)). Bring that in.
    module StateVariablesFiniteMap = Lattice.FiniteValueMap.WithKeys _≟_ isLatticeᵛ states
    open StateVariablesFiniteMap
        using (_[_]; m₁≼m₂⇒m₁[ks]≼m₂[ks])
        renaming
            ( FiniteMap to StateVariables
            ; isLattice to isLatticeᵐ
            ; _∈k_ to _∈kᵐ_
            ; locate to locateᵐ
            ; _≼_ to _≼ᵐ_
            ; ≈₂-dec⇒≈-dec to ≈ᵛ-dec⇒≈ᵐ-dec
            ; m₁≼m₂⇒m₁[k]≼m₂[k] to m₁≼m₂⇒m₁[k]ᵐ≼m₂[k]ᵐ
            )
    open Lattice.FiniteValueMap.IterProdIsomorphism.WithUniqueKeysAndFixedHeight _≟_ isLatticeᵛ states-Unique ≈ᵛ-dec _ fixedHeightᵛ
        using ()
        renaming
            ( isFiniteHeightLattice to isFiniteHeightLatticeᵐ
            )

    ≈ᵐ-dec = ≈ᵛ-dec⇒≈ᵐ-dec ≈ᵛ-dec
    fixedHeightᵐ = IsFiniteHeightLattice.fixedHeight isFiniteHeightLatticeᵐ

    -- build up the 'join' function, which follows from Exercise 4.26's
    --
    --    L₁ → (A → L₂)
    --
    -- Construction, with L₁ = (A → L₂), and f = id

    joinForKey : State → StateVariables → VariableSigns
    joinForKey k states = foldr _⊔ᵛ_ ⊥ᵛ (states [ incoming k ])

    -- The per-key join is made up of map key accesses (which are monotonic)
    -- and folds using the join operation (also monotonic)

    joinForKey-Mono : ∀ (k : State) → Monotonic _≼ᵐ_ _≼ᵛ_ (joinForKey k)
    joinForKey-Mono k {fm₁} {fm₂} fm₁≼fm₂ =
        foldr-Mono joinSemilatticeᵛ joinSemilatticeᵛ (fm₁ [ incoming k ]) (fm₂ [ incoming k ]) _⊔ᵛ_ ⊥ᵛ ⊥ᵛ
                   (m₁≼m₂⇒m₁[ks]≼m₂[ks] fm₁ fm₂ (incoming k) fm₁≼fm₂)
                   (⊔ᵛ-idemp ⊥ᵛ) ⊔ᵛ-Monotonicʳ ⊔ᵛ-Monotonicˡ

    -- The name f' comes from the formulation of Exercise 4.26.

    open StateVariablesFiniteMap.GeneralizedUpdate states isLatticeᵐ (λ x → x) (λ a₁≼a₂ → a₁≼a₂) joinForKey joinForKey-Mono states
        renaming
            ( f' to joinAll
            ; f'-Monotonic to joinAll-Mono
            )

    -- With 'join' in hand, we need to perform abstract evaluation.

    vars-in-Map : ∀ (k : String) (vs : VariableSigns) →
                  k ∈ˡ vars → k ∈kᵛ vs
    vars-in-Map k vs@(m , kvs≡vars) k∈vars rewrite kvs≡vars = k∈vars

    states-in-Map : ∀ (s : State) (sv : StateVariables) → s ∈kᵐ sv
    states-in-Map s sv@(m , ksv≡states) rewrite ksv≡states = states-complete s

    eval : ∀ (e : Expr) → (∀ k → k ∈ᵉ e → k ∈ˡ vars) → VariableSigns → SignLattice
    eval (e₁ + e₂) k∈e⇒k∈vars vs =
        plus (eval e₁ (λ k k∈e₁ → k∈e⇒k∈vars k (in⁺₁ k∈e₁)) vs)
             (eval e₂ (λ k k∈e₂ → k∈e⇒k∈vars k (in⁺₂ k∈e₂)) vs)
    eval (e₁ - e₂) k∈e⇒k∈vars vs =
        minus (eval e₁ (λ k k∈e₁ → k∈e⇒k∈vars k (in⁻₁ k∈e₁)) vs)
              (eval e₂ (λ k k∈e₂ → k∈e⇒k∈vars k (in⁻₂ k∈e₂)) vs)
    eval (` k) k∈e⇒k∈vars vs = proj₁ (locateᵛ {k} {vs} (vars-in-Map k vs (k∈e⇒k∈vars k here)))
    eval (# 0) _ _ = [ 0ˢ ]ᵍ
    eval (# (suc n')) _ _ = [ + ]ᵍ

    eval-Mono : ∀ (e : Expr) (k∈e⇒k∈vars : ∀ k → k ∈ᵉ e → k ∈ˡ vars) → Monotonic _≼ᵛ_ _≼ᵍ_ (eval e k∈e⇒k∈vars)
    eval-Mono (e₁ + e₂) k∈e⇒k∈vars {vs₁} {vs₂} vs₁≼vs₂ =
        let
            -- TODO: can this be done with less boilerplate?
            k∈e₁⇒k∈vars = λ k k∈e₁ → k∈e⇒k∈vars k (in⁺₁ k∈e₁)
            k∈e₂⇒k∈vars = λ k k∈e₂ → k∈e⇒k∈vars k (in⁺₂ k∈e₂)
            g₁vs₁ = eval e₁ k∈e₁⇒k∈vars vs₁
            g₂vs₁ = eval e₂ k∈e₂⇒k∈vars vs₁
            g₁vs₂ = eval e₁ k∈e₁⇒k∈vars vs₂
            g₂vs₂ = eval e₂ k∈e₂⇒k∈vars vs₂
        in
            ≼ᵍ-trans
                {plus g₁vs₁ g₂vs₁} {plus g₁vs₂ g₂vs₁} {plus g₁vs₂ g₂vs₂}
                (plus-Monoˡ g₂vs₁ {g₁vs₁} {g₁vs₂} (eval-Mono e₁ k∈e₁⇒k∈vars {vs₁} {vs₂} vs₁≼vs₂))
                (plus-Monoʳ g₁vs₂ {g₂vs₁} {g₂vs₂} (eval-Mono e₂ k∈e₂⇒k∈vars {vs₁} {vs₂} vs₁≼vs₂))
    eval-Mono (e₁ - e₂) k∈e⇒k∈vars {vs₁} {vs₂} vs₁≼vs₂ =
        let
            -- TODO: here too -- can this be done with less boilerplate?
            k∈e₁⇒k∈vars = λ k k∈e₁ → k∈e⇒k∈vars k (in⁻₁ k∈e₁)
            k∈e₂⇒k∈vars = λ k k∈e₂ → k∈e⇒k∈vars k (in⁻₂ k∈e₂)
            g₁vs₁ = eval e₁ k∈e₁⇒k∈vars vs₁
            g₂vs₁ = eval e₂ k∈e₂⇒k∈vars vs₁
            g₁vs₂ = eval e₁ k∈e₁⇒k∈vars vs₂
            g₂vs₂ = eval e₂ k∈e₂⇒k∈vars vs₂
        in
            ≼ᵍ-trans
                {minus g₁vs₁ g₂vs₁} {minus g₁vs₂ g₂vs₁} {minus g₁vs₂ g₂vs₂}
                (minus-Monoˡ g₂vs₁ {g₁vs₁} {g₁vs₂} (eval-Mono e₁ k∈e₁⇒k∈vars {vs₁} {vs₂} vs₁≼vs₂))
                (minus-Monoʳ g₁vs₂ {g₂vs₁} {g₂vs₂} (eval-Mono e₂ k∈e₂⇒k∈vars {vs₁} {vs₂} vs₁≼vs₂))
    eval-Mono (` k) k∈e⇒k∈vars {vs₁} {vs₂} vs₁≼vs₂ =
        let
            (v₁ , k,v₁∈vs₁) = locateᵛ {k} {vs₁} (vars-in-Map k vs₁ (k∈e⇒k∈vars k here))
            (v₂ , k,v₂∈vs₂) = locateᵛ {k} {vs₂} (vars-in-Map k vs₂ (k∈e⇒k∈vars k here))
        in
            m₁≼m₂⇒m₁[k]ᵛ≼m₂[k]ᵛ vs₁ vs₂ vs₁≼vs₂ k,v₁∈vs₁ k,v₂∈vs₂
    eval-Mono (# 0) _ _ = ≈ᵍ-refl
    eval-Mono (# (suc n')) _ _ = ≈ᵍ-refl

    private module _ (k : String) (e : Expr) (k∈e⇒k∈vars : ∀ k → k ∈ᵉ e → k ∈ˡ vars) where
        open VariableSignsFiniteMap.GeneralizedUpdate vars isLatticeᵛ (λ x → x) (λ a₁≼a₂ → a₁≼a₂) (λ _ → eval e k∈e⇒k∈vars) (λ _ {vs₁} {vs₂} vs₁≼vs₂ → eval-Mono e k∈e⇒k∈vars {vs₁} {vs₂} vs₁≼vs₂) (k ∷ [])
            renaming
                ( f' to updateVariablesFromExpression
                ; f'-Monotonic to updateVariablesFromExpression-Mono
                )
            public

    updateVariablesForState : State → StateVariables → VariableSigns
    updateVariablesForState s sv
        -- More weirdness here. Apparently, capturing the with-equality proof
        -- using 'in p' makes code that reasons about this function (below)
        -- throw ill-typed with-abstraction errors. Instead, make use of the
        -- fact that later with-clauses are generalized over earlier ones to
        -- construct a specialization of vars-complete for (code s).
        with code s | (λ k → vars-complete {k} s)
    ...   | k ← e | k∈codes⇒k∈vars =
            let
                (vs , s,vs∈sv) = locateᵐ {s} {sv} (states-in-Map s sv)
            in
                updateVariablesFromExpression k e (λ k k∈e → k∈codes⇒k∈vars k (in←₂ k∈e)) vs

    updateVariablesForState-Monoʳ : ∀ (s : State) → Monotonic _≼ᵐ_ _≼ᵛ_ (updateVariablesForState s)
    updateVariablesForState-Monoʳ s {sv₁} {sv₂} sv₁≼sv₂
        with code s | (λ k → vars-complete {k} s)
    ...   | k ← e | k∈codes⇒k∈vars =
            let
                (vs₁ , s,vs₁∈sv₁) = locateᵐ {s} {sv₁} (states-in-Map s sv₁)
                (vs₂ , s,vs₂∈sv₂) = locateᵐ {s} {sv₂} (states-in-Map s sv₂)
                vs₁≼vs₂ = m₁≼m₂⇒m₁[k]ᵐ≼m₂[k]ᵐ sv₁ sv₂ sv₁≼sv₂ s,vs₁∈sv₁ s,vs₂∈sv₂
            in
                updateVariablesFromExpression-Mono k e (λ k k∈e → k∈codes⇒k∈vars k (in←₂ k∈e)) {vs₁} {vs₂} vs₁≼vs₂

    open StateVariablesFiniteMap.GeneralizedUpdate states isLatticeᵐ (λ x → x) (λ a₁≼a₂ → a₁≼a₂) updateVariablesForState updateVariablesForState-Monoʳ states
        renaming
            ( f' to updateAll
            ; f'-Monotonic to updateAll-Mono
            )

    analyze : StateVariables → StateVariables
    analyze = updateAll ∘ joinAll

    analyze-Mono : Monotonic _≼ᵐ_ _≼ᵐ_ analyze
    analyze-Mono {sv₁} {sv₂} sv₁≼sv₂ = updateAll-Mono {joinAll sv₁} {joinAll sv₂} (joinAll-Mono {sv₁} {sv₂} sv₁≼sv₂)

    open import Fixedpoint ≈ᵐ-dec isFiniteHeightLatticeᵐ analyze (λ {m₁} {m₂} m₁≼m₂ → analyze-Mono {m₁} {m₂} m₁≼m₂)
        using ()
        renaming (aᶠ to signs)


    -- Debugging code: print the resulting map.
    output = show signs