{-# OPTIONS --safe #-}

module Problem23 where

open import Data.List
open import Data.List.Properties
open import Data.Fin using (fromℕ<)
open import Data.Fin.Properties using (fromℕ<-cong)
open import Data.Empty
open import Data.Nat
open import Data.Nat.Properties
open import Data.Product
open import Function
open import Relation.Binary.PropositionalEquality
open import Relation.Nullary
open import Relation.Nullary.Negation

open ≡-Reasoning


variable State : Set
variable f : State → State

Assertion : Set → Set₁
Assertion State = State → Set
variable ϕ ϕ′ ψ ψ′ π g : Assertion State

_∧_ : Assertion State → Assertion State → Assertion State
ϕ ∧ ψ = λ σ → ϕ σ × ψ σ

¬'_ : Assertion State → Assertion State
¬' ϕ = λ σ → (ϕ σ → ⊥)

_⇒_ : Assertion State → Assertion State → Set
ϕ ⇒ ψ = ∀ σ → ϕ σ → ψ σ

data Program (State : Set) : Set₁ where
    SKIP : Program State
    IF_THEN_ELSE_FI : Assertion State → Program State → Program State → Program State
    WHILE_DO_OD     : Assertion State → Program State → Program State
    _∶_ : Program State → Program State → Program State
    UPD : (State → State) → Program State

infixl 5 _∶_

variable P Q : Program State

data ⟨_⟩_⟨_⟩ {State : Set} : Assertion State → Program State → Assertion State → Set₁ where
    skipₕ : ⟨ ϕ ⟩ SKIP ⟨ ϕ ⟩
    ifₕ : ⟨ ϕ ∧ g ⟩ P ⟨ ψ ⟩ → ⟨ ϕ ∧ (¬' g) ⟩ Q ⟨ ψ ⟩ → ⟨ ϕ ⟩ IF g THEN P ELSE Q FI ⟨ ψ ⟩
    whileₕ : ⟨ ϕ ∧ g ⟩ P ⟨ ϕ ⟩ → ⟨ ϕ ⟩ WHILE g DO P OD ⟨ ϕ ∧ (¬' g) ⟩
    seqₕ : ⟨ ϕ ⟩ P ⟨ π ⟩ → ⟨ π ⟩ Q ⟨ ψ ⟩ → ⟨ ϕ ⟩ P ∶ Q ⟨ ψ ⟩
    updₕ : ⟨ ϕ ∘ f ⟩ UPD f ⟨ ϕ ⟩
    conseqₕ : ϕ ⇒ ϕ′ → ⟨ ϕ′ ⟩ P ⟨ ψ′ ⟩ → ψ′ ⇒ ψ → ⟨ ϕ ⟩ P ⟨ ψ ⟩


record ListSum : Set where
    field
        A : List ℕ
        s : ℕ
        i : ℕ

lookup' : List ℕ → ℕ → ℕ
lookup' V n with n <? length V
... | yes p = lookup V (fromℕ< p)
... | no p = 0

open ListSum
prog : Program ListSum
prog = UPD (λ σ → record σ { s = 0 })
     ∶ UPD (λ σ → record σ { i = 0 })
     ∶ WHILE (λ σ → σ .i ≢ length (σ .A)) DO
         UPD (λ σ → record σ { s = σ .s + lookup' (σ .A) (σ .i) })
       ∶ UPD (λ σ → record σ { i = suc (σ .i) })
       OD


goal : ∀ A₀ → ⟨ (λ σ → A₀ ≡ σ .A) ⟩ prog ⟨ (λ σ → σ .s ≡ sum A₀) ⟩
goal A₀ =
  conseqₕ
    (λ σ eq → trans (+-identityʳ (sum (σ .A))) (cong sum (sym eq)) , z≤n)
    (seqₕ (seqₕ updₕ updₕ)
      (whileₕ {ϕ = λ σ → sum (drop (σ .i) (σ .A)) + σ .s ≡ sum A₀ × σ .i ≤ length (σ .A)}
        (conseqₕ
          (λ σ ((prf , i≤length) , i≢length) →
            let
              1+i≤length = ≤∧≢⇒< i≤length i≢length

              prf′ = begin
                sum (drop (suc (σ .i)) (σ .A)) + (σ .s + lookup' (σ .A) (σ .i)) ≡˘⟨ +-assoc (sum (drop (suc (σ .i)) (σ .A))) _ _ ⟩
                sum (drop (suc (σ .i)) (σ .A)) + σ .s + lookup' (σ .A) (σ .i)   ≡⟨ +-comm _ (lookup' (σ .A) (σ .i)) ⟩
                lookup' (σ .A) (σ .i) + (sum (drop (suc (σ .i)) (σ .A)) + σ .s) ≡˘⟨ +-assoc (lookup' (σ .A) (σ .i)) _ _ ⟩
                lookup' (σ .A) (σ .i) + sum (drop (suc (σ .i)) (σ .A)) + σ .s   ≡⟨ cong (_+ σ .s) (sum-drop-suc (σ .i) (σ .A) 1+i≤length) ⟩
                sum (drop (σ .i) (σ .A)) + σ .s                                 ≡⟨ prf ⟩
                sum A₀                                                          ∎
            in
            prf′ , 1+i≤length)
          (seqₕ updₕ updₕ)
          (λ σ → id))))
    (λ σ ((prf , i≤length) , ¬i≢length) →
      let
        i≡length : σ .i ≡ length (σ .A)
        i≡length = decidable-stable (σ .i ≟ length (σ .A)) ¬i≢length

        dropᵢ[A]≡[] : drop (σ .i) (σ .A) ≡ []
        dropᵢ[A]≡[] = length⁻¹ (begin
          length (drop (σ .i) (σ .A)) ≡⟨  length-drop (σ .i) (σ .A) ⟩
          length (σ .A) ∸ σ .i        ≡˘⟨ cong (_∸ σ .i) i≡length ⟩
          σ .i ∸ σ .i                 ≡⟨  n∸n≡0 (σ .i) ⟩
          0                           ∎)
      in
      begin
        σ .s                            ≡⟨⟩
        sum [] + σ .s                   ≡˘⟨ cong (λ ◌ → sum ◌ + σ .s) dropᵢ[A]≡[] ⟩
        sum (drop (σ .i) (σ .A)) + σ .s ≡⟨ prf ⟩
        sum A₀                          ∎)
    where
    length⁻¹ : ∀ {xs : List ℕ} → length xs ≡ 0 → xs ≡ []
    length⁻¹ {[]} prf = refl

    lookup'-suc : (n x : ℕ) (xs : List ℕ) → lookup' (x ∷ xs) (suc n) ≡ lookup' xs n
    lookup'-suc n x xs with n <? length xs | suc n <? suc (length xs)
    ... | yes p₁ | yes p₂ = cong (lookup xs) (fromℕ<-cong n n refl (≤-pred p₂) p₁)
    ... | yes p₁ | no  p₂ = ⊥-elim (p₂ (s≤s p₁))
    ... | no  p₁ | yes p₂ = ⊥-elim (p₁ (≤-pred p₂))
    ... | no  p₁ | no  p₂ = refl

    sum-drop-suc :
      (n : ℕ) (xs : List ℕ) →
      .(n < length xs) →
      lookup' xs n + sum (drop (suc n) xs) ≡ sum (drop n xs)
    sum-drop-suc zero    (x ∷ xs) _        = refl
    sum-drop-suc (suc n) (x ∷ xs) n<length = begin
      lookup' (x ∷ xs) (suc n) + sum (drop (suc n) xs) ≡⟨ cong (_+ sum (drop (suc n) xs)) (lookup'-suc n x xs) ⟩
      lookup' xs n + sum (drop (suc n) xs)             ≡⟨ sum-drop-suc n xs (≤-pred n<length) ⟩
      sum (drop n xs)                                  ∎