module Language where

open import Data.Nat using (ℕ; suc; pred; _≤_) renaming (_+_ to _+ⁿ_)
open import Data.Nat.Properties using (m≤n⇒m≤n+o; ≤-reflexive)
open import Data.Integer using (ℤ; +_) renaming (_+_ to _+ᶻ_; _-_ to _-ᶻ_)
open import Data.String using (String) renaming (_≟_ to _≟ˢ_)
open import Data.Product using (_×_; Σ; _,_; proj₁; proj₂)
open import Data.Vec using (Vec; foldr; lookup; _∷_; []; _++_)
open import Data.List using ([]; _∷_; List) renaming (foldr to foldrˡ; map to mapˡ)
open import Data.List.Membership.Propositional as MemProp using () renaming (_∈_ to _∈ˡ_)
open import Data.List.Relation.Unary.All using (All; []; _∷_)
open import Data.List.Relation.Unary.Any as RelAny using ()
open import Data.Fin using (Fin; suc; zero; fromℕ; inject₁; inject≤; _↑ʳ_) renaming (_≟_ to _≟ᶠ_)
open import Data.Fin.Properties using (suc-injective)
open import Relation.Binary.PropositionalEquality using (subst; cong; _≡_; refl)
open import Relation.Nullary using (¬_)
open import Function using (_∘_)

open import Lattice
open import Utils using (Unique; Unique-map; empty; push; x∈xs⇒fx∈fxs; _⊗_; _,_)

data Expr : Set where
    _+_ : Expr → Expr → Expr
    _-_ : Expr → Expr → Expr
    `_ : String → Expr
    #_ : ℕ → Expr

data BasicStmt : Set where
    _←_ : String → Expr → BasicStmt
    noop : BasicStmt

data Stmt : Set where
    ⟨_⟩ : BasicStmt → Stmt
    _then_ : Stmt → Stmt → Stmt
    if_then_else_ : Expr → Stmt → Stmt → Stmt
    while_repeat_ : Expr → Stmt → Stmt

module Semantics where
    data Value : Set where
        ↑ᶻ : ℤ → Value

    Env : Set
    Env = List (String × Value)

    data _∈_ : (String × Value) → Env → Set where
        here : ∀ (s : String) (v : Value) (ρ : Env) → (s , v) ∈ ((s , v) ∷ ρ)
        there : ∀ (s s' : String) (v v' : Value) (ρ : Env) → ¬ (s ≡ s') → (s , v) ∈ ρ → (s , v) ∈ ((s' , v') ∷ ρ)

    data _,_⇒ᵉ_ : Env → Expr → Value → Set where
        ⇒ᵉ-ℕ : ∀ (ρ : Env) (n : ℕ) → ρ , (# n) ⇒ᵉ (↑ᶻ (+ n))
        ⇒ᵉ-Var : ∀ (ρ : Env) (x : String) (v : Value) → (x , v) ∈ ρ → ρ , (` x) ⇒ᵉ v
        ⇒ᵉ-+ : ∀ (ρ : Env) (e₁ e₂ : Expr) (z₁ z₂ : ℤ) →
               ρ , e₁ ⇒ᵉ (↑ᶻ z₁) → ρ , e₂ ⇒ᵉ (↑ᶻ z₂) →
               ρ , (e₁ + e₂) ⇒ᵉ (↑ᶻ (z₁ +ᶻ z₂))
        ⇒ᵉ-- : ∀ (ρ : Env) (e₁ e₂ : Expr) (z₁ z₂ : ℤ) →
               ρ , e₁ ⇒ᵉ (↑ᶻ z₁) → ρ , e₂ ⇒ᵉ (↑ᶻ z₂) →
               ρ , (e₁ - e₂) ⇒ᵉ (↑ᶻ (z₁ -ᶻ z₂))

    data _,_⇒ᵇ_ : Env → BasicStmt → Env → Set where
        ⇒ᵇ-noop : ∀ (ρ : Env) → ρ , noop ⇒ᵇ ρ
        ⇒ᵇ-← : ∀ (ρ : Env) (x : String) (e : Expr) (v : Value) →
               ρ , e ⇒ᵉ v → ρ , (x ← e) ⇒ᵇ ((x , v) ∷ ρ)

    data _,_⇒ˢ_ : Env → Stmt → Env → Set where
        ⇒ˢ-⟨⟩ : ∀ (ρ₁ ρ₂ : Env) (bs : BasicStmt) →
                ρ₁ , bs ⇒ᵇ ρ₂ → ρ₁ , ⟨ bs ⟩ ⇒ˢ ρ₂
        ⇒ˢ-then : ∀ (ρ₁ ρ₂ ρ₃ : Env) (s₁ s₂ : Stmt) →
                  ρ₁ , s₁ ⇒ˢ ρ₂ → ρ₂ , s₂ ⇒ˢ ρ₃ →
                  ρ₁ , (s₁ then s₂) ⇒ˢ ρ₃
        ⇒ˢ-if-true : ∀ (ρ₁ ρ₂ : Env) (e : Expr) (z : ℤ) (s₁ s₂ : Stmt) →
            ρ₁ , e ⇒ᵉ (↑ᶻ z) → ¬ z ≡ (+ 0) → ρ₁ , s₁ ⇒ˢ ρ₂ →
            ρ₁ , (if e then s₁ else s₂) ⇒ˢ ρ₂
        ⇒ˢ-if-false : ∀ (ρ₁ ρ₂ : Env) (e : Expr) (s₁ s₂ : Stmt) →
            ρ₁ , e ⇒ᵉ (↑ᶻ (+ 0)) → ρ₁ , s₂ ⇒ˢ ρ₂ →
            ρ₁ , (if e then s₁ else s₂) ⇒ˢ ρ₂
        ⇒ˢ-while-true : ∀ (ρ₁ ρ₂ ρ₃ : Env) (e : Expr) (z : ℤ) (s : Stmt) →
            ρ₁ , e ⇒ᵉ (↑ᶻ z) → ¬ z ≡ (+ 0) → ρ₁ , s ⇒ˢ ρ₂ → ρ₂ , (while e repeat s) ⇒ˢ ρ₃ →
            ρ₁ , (while e repeat s) ⇒ˢ ρ₃
        ⇒ˢ-while-false : ∀ (ρ : Env) (e : Expr) (s : Stmt) →
            ρ , e ⇒ᵉ (↑ᶻ (+ 0)) →
            ρ , (while e repeat s) ⇒ˢ ρ

module Graphs where
    open Semantics

    record Graph : Set where
        field
            size : ℕ

        Index : Set
        Index = Fin size

        Edge : Set
        Edge = Index × Index

        field
            nodes : Vec (List BasicStmt) size
            edges : List Edge

    _[_] : ∀ (g : Graph) → Graph.Index g → List BasicStmt
    _[_] g idx = lookup (Graph.nodes g) idx

    _⊆_ : Graph → Graph → Set
    _⊆_ g₁ g₂ =
        Σ (Graph.size g₁ ≤ Graph.size g₂) (λ n₁≤n₂ →
            ( ∀ (idx : Graph.Index g₁) → g₁ [ idx ] ≡ g₂ [ inject≤ idx n₁≤n₂ ]
            × ∀ (idx₁ idx₂ : Graph.Index g₁) → (idx₁ , idx₂) ∈ˡ (Graph.edges g₁) →
                (inject≤ idx₁ n₁≤n₂ , inject≤ idx₂ n₁≤n₂) ∈ˡ (Graph.edges g₂)
            ))

    -- Note: inject≤ doesn't seem to work as nicely with vector lookups.
    --       The ↑ˡ and ↑ʳ operators are way nicer. Can we reformulate the
    --       ⊆ property to use them?

    n≤n+m : ∀ (n m : ℕ) → n ≤ n +ⁿ m
    n≤n+m n m = m≤n⇒m≤n+o m (≤-reflexive (refl {x = n}))

    lookup-++ˡ : ∀ {a} {A : Set a} {n m : ℕ} (xs : Vec A n) (ys : Vec A m)
                   (idx : Fin n) → lookup xs idx ≡ lookup (xs ++ ys) (inject≤ idx (n≤n+m n m))
    lookup-++ˡ = {!!}

    pushBasicBlock : List BasicStmt → (g₁ : Graph) → Σ Graph (λ g₂ → Graph.Index g₂ × g₁ ⊆ g₂)
    pushBasicBlock bss g₁ =
        let
            size' = Graph.size g₁ +ⁿ 1
            size≤size' = n≤n+m (Graph.size g₁) 1
            inject-Edge = λ (idx₁ , idx₂) → (inject≤ idx₁ size≤size' , inject≤ idx₂ size≤size')
        in
            ( record
                { size = size'
                ; nodes = Graph.nodes g₁ ++ (bss ∷ [])
                ; edges = mapˡ inject-Edge (Graph.edges g₁)
                }
            , ( (Graph.size g₁) ↑ʳ zero
              , ( size≤size'
                , λ idx → lookup-++ˡ (Graph.nodes g₁) (bss ∷ []) idx
                , λ idx₁ idx₂ e∈es → x∈xs⇒fx∈fxs inject-Edge e∈es
                )
              )
            )

    record Relaxable (T : Graph → Set) : Set where
        field relax : ∀ {g₁ g₂ : Graph} → g₁ ⊆ g₂ → T g₁ → T g₂

    instance
        IndexRelaxable : Relaxable Graph.Index
        IndexRelaxable = record
            { relax = λ g₁⊆g₂ idx → inject≤ idx (proj₁ g₁⊆g₂)
            }

        EdgeRelaxable : Relaxable Graph.Edge
        EdgeRelaxable = record
            { relax = λ {g₁} {g₂} g₁⊆g₂ (idx₁ , idx₂) →
                ( Relaxable.relax IndexRelaxable {g₁} {g₂} g₁⊆g₂ idx₁
                , Relaxable.relax IndexRelaxable {g₁} {g₂} g₁⊆g₂ idx₂
                )
            }

        ProdRelaxable : ∀ {P : Graph → Set} {Q : Graph → Set} →
                        {{ PRelaxable : Relaxable P }} → {{ QRelaxable : Relaxable Q }} →
                        Relaxable (P ⊗ Q)
        ProdRelaxable {{pr}} {{qr}} = record
            { relax = (λ { g₁⊆g₂ (p , q) →
                ( Relaxable.relax pr g₁⊆g₂ p
                , Relaxable.relax qr g₁⊆g₂ q) }
                )
            }

    open Relaxable {{...}} public

open import Lattice.MapSet _≟ˢ_
    renaming
        ( MapSet to StringSet
        ; insert to insertˢ
        ; to-List to to-Listˢ
        ; empty to emptyˢ
        ; singleton to singletonˢ
        ; _⊔_ to _⊔ˢ_
        ; `_ to `ˢ_
        ; _∈_ to _∈ˢ_
        ; ⊔-preserves-∈k₁ to ⊔ˢ-preserves-∈k₁
        ; ⊔-preserves-∈k₂ to ⊔ˢ-preserves-∈k₂
        )

data _∈ᵉ_ : String → Expr → Set where
    in⁺₁ : ∀ {e₁ e₂ : Expr} {k : String} → k ∈ᵉ e₁ → k ∈ᵉ (e₁ + e₂)
    in⁺₂ : ∀ {e₁ e₂ : Expr} {k : String} → k ∈ᵉ e₂ → k ∈ᵉ (e₁ + e₂)
    in⁻₁ : ∀ {e₁ e₂ : Expr} {k : String} → k ∈ᵉ e₁ → k ∈ᵉ (e₁ - e₂)
    in⁻₂ : ∀ {e₁ e₂ : Expr} {k : String} → k ∈ᵉ e₂ → k ∈ᵉ (e₁ - e₂)
    here : ∀ {k : String} → k ∈ᵉ (` k)

data _∈ᵇ_ : String → BasicStmt → Set where
    in←₁ : ∀ {k : String} {e : Expr} → k ∈ᵇ (k ← e)
    in←₂ : ∀ {k k' : String} {e : Expr} → k ∈ᵉ e → k ∈ᵇ (k' ← e)

private
    Expr-vars : Expr → StringSet
    Expr-vars (l + r) = Expr-vars l ⊔ˢ Expr-vars r
    Expr-vars (l - r) = Expr-vars l ⊔ˢ Expr-vars r
    Expr-vars (` s) = singletonˢ s
    Expr-vars (# _) = emptyˢ

    -- ∈-Expr-vars⇒∈ : ∀ {k : String} (e : Expr) → k ∈ˢ (Expr-vars e) → k ∈ᵉ e
    -- ∈-Expr-vars⇒∈ {k} (e₁ + e₂) k∈vs
    --     with Expr-Provenance k ((`ˢ (Expr-vars e₁)) ∪ (`ˢ (Expr-vars e₂))) k∈vs
    -- ...   | in₁ (single k,tt∈vs₁) _ = (in⁺₁ (∈-Expr-vars⇒∈ e₁ (forget k,tt∈vs₁)))
    -- ...   | in₂ _ (single k,tt∈vs₂) = (in⁺₂ (∈-Expr-vars⇒∈ e₂ (forget k,tt∈vs₂)))
    -- ...   | bothᵘ (single k,tt∈vs₁) _ = (in⁺₁ (∈-Expr-vars⇒∈ e₁ (forget k,tt∈vs₁)))
    -- ∈-Expr-vars⇒∈ {k} (e₁ - e₂) k∈vs
    --     with Expr-Provenance k ((`ˢ (Expr-vars e₁)) ∪ (`ˢ (Expr-vars e₂))) k∈vs
    -- ...   | in₁ (single k,tt∈vs₁) _ = (in⁻₁ (∈-Expr-vars⇒∈ e₁ (forget k,tt∈vs₁)))
    -- ...   | in₂ _ (single k,tt∈vs₂) = (in⁻₂ (∈-Expr-vars⇒∈ e₂ (forget k,tt∈vs₂)))
    -- ...   | bothᵘ (single k,tt∈vs₁) _ = (in⁻₁ (∈-Expr-vars⇒∈ e₁ (forget k,tt∈vs₁)))
    -- ∈-Expr-vars⇒∈ {k} (` k) (RelAny.here refl) = here

    -- ∈⇒∈-Expr-vars : ∀ {k : String} {e : Expr} → k ∈ᵉ e → k ∈ˢ (Expr-vars e)
    -- ∈⇒∈-Expr-vars {k} {e₁ + e₂} (in⁺₁ k∈e₁) =
    --     ⊔ˢ-preserves-∈k₁ {m₁ = Expr-vars e₁}
    --                      {m₂ = Expr-vars e₂}
    --                      (∈⇒∈-Expr-vars k∈e₁)
    -- ∈⇒∈-Expr-vars {k} {e₁ + e₂} (in⁺₂ k∈e₂) =
    --     ⊔ˢ-preserves-∈k₂ {m₁ = Expr-vars e₁}
    --                      {m₂ = Expr-vars e₂}
    --                      (∈⇒∈-Expr-vars k∈e₂)
    -- ∈⇒∈-Expr-vars {k} {e₁ - e₂} (in⁻₁ k∈e₁) =
    --     ⊔ˢ-preserves-∈k₁ {m₁ = Expr-vars e₁}
    --                      {m₂ = Expr-vars e₂}
    --                      (∈⇒∈-Expr-vars k∈e₁)
    -- ∈⇒∈-Expr-vars {k} {e₁ - e₂} (in⁻₂ k∈e₂) =
    --     ⊔ˢ-preserves-∈k₂ {m₁ = Expr-vars e₁}
    --                      {m₂ = Expr-vars e₂}
    --                      (∈⇒∈-Expr-vars k∈e₂)
    -- ∈⇒∈-Expr-vars here = RelAny.here refl

    BasicStmt-vars : BasicStmt → StringSet
    BasicStmt-vars (x ← e) = (singletonˢ x) ⊔ˢ (Expr-vars e)
    BasicStmt-vars noop = emptyˢ

    Stmt-vars : Stmt → StringSet
    Stmt-vars ⟨ bs ⟩ = BasicStmt-vars bs
    Stmt-vars (s₁ then s₂) = (Stmt-vars s₁) ⊔ˢ (Stmt-vars s₂)
    Stmt-vars (if e then s₁ else s₂) = ((Expr-vars e) ⊔ˢ (Stmt-vars s₁)) ⊔ˢ (Stmt-vars s₂)
    Stmt-vars (while e repeat s) = (Expr-vars e) ⊔ˢ (Stmt-vars s)

    -- ∈-Stmt-vars⇒∈ : ∀ {k : String} (s : Stmt) → k ∈ˢ (Stmt-vars s) → k ∈ᵇ s
    -- ∈-Stmt-vars⇒∈ {k} (k' ← e) k∈vs
    --     with Expr-Provenance k ((`ˢ (singletonˢ k')) ∪ (`ˢ (Expr-vars e))) k∈vs
    -- ...   | in₁ (single (RelAny.here refl)) _ = in←₁
    -- ...   | in₂ _ (single k,tt∈vs') = in←₂ (∈-Expr-vars⇒∈ e (forget k,tt∈vs'))
    -- ...   | bothᵘ (single (RelAny.here refl)) _ = in←₁

    -- ∈⇒∈-Stmt-vars : ∀ {k : String} {s : Stmt} → k ∈ᵇ s → k ∈ˢ (Stmt-vars s)
    -- ∈⇒∈-Stmt-vars {k} {k ← e} in←₁ =
    --     ⊔ˢ-preserves-∈k₁ {m₁ = singletonˢ k}
    --                      {m₂ = Expr-vars e}
    --                      (RelAny.here refl)
    -- ∈⇒∈-Stmt-vars {k} {k' ← e} (in←₂ k∈e) =
    --     ⊔ˢ-preserves-∈k₂ {m₁ = singletonˢ k'}
    --                      {m₂ = Expr-vars e}
    --                      (∈⇒∈-Expr-vars k∈e)

    Stmts-vars : ∀ {n : ℕ} → Vec Stmt n → StringSet
    Stmts-vars = foldr (λ n → StringSet)
                       (λ {k} stmt acc → (Stmt-vars stmt) ⊔ˢ acc) emptyˢ

    -- ∈-Stmts-vars⇒∈ : ∀ {n : ℕ} {k : String} (ss : Vec Stmt n) →
    --                  k ∈ˢ (Stmts-vars ss) → Σ (Fin n) (λ f → k ∈ᵇ lookup ss f)
    -- ∈-Stmts-vars⇒∈ {suc n'} {k} (s ∷ ss') k∈vss
    --     with Expr-Provenance k ((`ˢ (Stmt-vars s)) ∪ (`ˢ (Stmts-vars ss'))) k∈vss
    -- ...   | in₁ (single k,tt∈vs) _ = (zero , ∈-Stmt-vars⇒∈ s (forget k,tt∈vs))
    -- ...   | in₂ _ (single k,tt∈vss') =
    --         let
    --             (f' , k∈s') = ∈-Stmts-vars⇒∈ ss' (forget k,tt∈vss')
    --         in
    --             (suc f' , k∈s')
    -- ...   | bothᵘ (single k,tt∈vs) _ = (zero , ∈-Stmt-vars⇒∈ s (forget k,tt∈vs))

    -- ∈⇒∈-Stmts-vars : ∀ {n : ℕ} {k : String} {ss : Vec Stmt n} {f : Fin n} →
    --                  k ∈ᵇ lookup ss f → k ∈ˢ (Stmts-vars ss)
    -- ∈⇒∈-Stmts-vars {suc n} {k} {s ∷ ss'} {zero} k∈s =
    --     ⊔ˢ-preserves-∈k₁ {m₁ = Stmt-vars s}
    --                      {m₂ = Stmts-vars ss'}
    --                      (∈⇒∈-Stmt-vars k∈s)
    -- ∈⇒∈-Stmts-vars {suc n} {k} {s ∷ ss'} {suc f'} k∈ss' =
    --     ⊔ˢ-preserves-∈k₂ {m₁ = Stmt-vars s}
    --                      {m₂ = Stmts-vars ss'}
    --                      (∈⇒∈-Stmts-vars {n} {k} {ss'} {f'} k∈ss')

    -- Creating a new number from a natural number can never create one
    -- equal to one you get from weakening the bounds on another number.
    z≢sf : ∀ {n : ℕ} (f : Fin n) → ¬ (zero ≡ suc f)
    z≢sf f ()

    z≢mapsfs : ∀ {n : ℕ} (fs : List (Fin n)) → All (λ sf → ¬ zero ≡ sf) (mapˡ suc fs)
    z≢mapsfs [] = []
    z≢mapsfs (f ∷ fs') = z≢sf f ∷ z≢mapsfs fs'

    indices : ∀ (n : ℕ) → Σ (List (Fin n)) Unique
    indices 0 = ([] , empty)
    indices (suc n') =
        let
            (inds' , unids') = indices n'
        in
            ( zero ∷ mapˡ suc inds'
            , push (z≢mapsfs inds') (Unique-map suc suc-injective unids')
            )

    indices-complete : ∀ (n : ℕ) (f : Fin n) → f ∈ˡ (proj₁ (indices n))
    indices-complete (suc n') zero = RelAny.here refl
    indices-complete (suc n') (suc f') = RelAny.there (x∈xs⇒fx∈fxs suc (indices-complete n' f'))

    -- Sketch, 'build control flow graph'

    -- -- Create new block, mark it as the current insertion point.
    -- emptyBlock : m Id

    -- currentBlock : m Id

    -- -- Create a new block, and insert the statement into it. Shold restore insertion pont.
    -- createBlock : Stmt → m (Id × Id)

    -- -- Note that the given ID is a successor / predecessor of the given
    -- -- insertion point.
    -- noteSuccessor : Id → m ()
    -- notePredecessor : Id → m ()
    -- noteEdge : Id → Id → m ()

    -- -- Insert the given statment into the current insertion point.
    -- buildCfg : Stmt → m Cfg
    -- buildCfg { bs₁ } = push bs₁
    -- buildCfg (s₁ ; s₂ ) = buildCfg s₁ >> buildCfg s₂
    -- buildCfg (if _ then s₁ else s₂) = do
    --     (b₁ , b₁') ← createBlock s₁
    --     noteSuccessor b₁

    --     (b₂ , b₂') ← createBlock s₂
    --     noteSuccessor b₂

    --     b ← emptyBlock
    --     notePredecessor b₁'
    --     notePredecessor b₂'
    -- buildCfg (while e repeat s) = do
    --     (b₁, b₁') ← createBlock s
    --     noteSuccessor b₁
    --     noteEdge b₁' b₁

    --     b ← emptyBlock
    --     notePredecessor b₁'


-- For now, just represent the program and CFG as one type, without branching.
record Program : Set where
    field
        length : ℕ
        stmts : Vec Stmt length

    private
        vars-Set : StringSet
        vars-Set = Stmts-vars stmts

    vars : List String
    vars = to-Listˢ vars-Set

    vars-Unique : Unique vars
    vars-Unique = proj₂ vars-Set

    State : Set
    State = Fin length

    states : List State
    states = proj₁ (indices length)

    states-complete : ∀ (s : State) → s ∈ˡ states
    states-complete = indices-complete length

    states-Unique : Unique states
    states-Unique = proj₂ (indices length)

    code : State → Stmt
    code = lookup stmts

    -- vars-complete : ∀ {k : String} (s : State) → k ∈ᵇ (code s) → k ∈ˡ vars
    -- vars-complete {k} s = ∈⇒∈-Stmts-vars {length} {k} {stmts} {s}

    _≟_ : IsDecidable (_≡_ {_} {State})
    _≟_ = _≟ᶠ_

    -- Computations for incoming and outgoing edges will have to change too
    -- when we support branching etc.

    incoming : State → List State
    incoming
        with length
    ...   | 0 = (λ ())
    ...   | suc n' = (λ
            { zero → []
            ; (suc f') → (inject₁ f') ∷ []
            })