-- FVIntmax/Balance.lean

import Mathlib

import Mathlib.Algebra.Group.Int

import FVIntmax.BalanceProof

import FVIntmax.Block

import FVIntmax.Key

import FVIntmax.Propositions

import FVIntmax.State

import FVIntmax.Transaction

import FVIntmax.Wheels

import FVIntmax.Wheels.Dictionary

namespace Intmax

noncomputable section

open Classical

section Balance

variable {Pi K₁ K₂ V C Sigma : Type}

section Extraction

section Deposit

variable [PreWithZero V]

def TransactionsInBlock_deposit
  (b : { b : Block K₁ K₂ C Sigma V // b.isDepositBlock }) : List (Τ K₁ K₂ V) :=
  match h : b.1 with
  | .deposit r v => [⟨(.Source, r, v), by admit /- proof elided -/
  ⟩]
  | .withdrawal .. | .transfer .. => by admit /- proof elided -/

end Deposit

section Transfer

variable [Finite K₁] [Finite K₂]
         [LinearOrder K₁] [LinearOrder K₂] [PreWithZero V]

def TransactionsInBlock_transfer 
  (π : BalanceProof K₁ K₂ C Pi V) (b : { b : Block K₁ K₂ C Sigma V // b.isTransferBlock }) : List (Τ K₁ K₂ V) :=
  match h : b.1 with
  | .transfer _ _ commitment S _ =>
     
    let senderRecipient : Finset (K₂ × Key K₁ K₂) := { (k₂, k) | (k₂ : K₂) (k : Key K₁ K₂) (_h : k₂ ≠ₖ k) }
    let sorted : List (K₂ × Key K₁ K₂) := senderRecipient.sort Key.lexLe 
     
    let v (s : K₂) (r : Key K₁ K₂) : Option V₊ :=
      if s ∉ S
      then .some 0
      else 
        if h : (commitment, s) ∈ π.keys
        then let (_, t) := π[(commitment, s)]
             t r
        else .none
    sorted.attach.map λ ⟨(s, r), h⟩ ↦ ⟨(s, r, v s r), by admit /- proof elided -/
    ⟩
  | .deposit .. | .withdrawal .. => by admit /- proof elided -/

end Transfer

section Withdrawal

variable [LinearOrder K₁] [Finite K₁] [PreWithZero V]

def TransactionsInBlock_withdrawal 
  (b : { b : Block K₁ K₂ C Sigma V // b.isWithdrawalBlock }) : List (Τ K₁ K₂ V) :=
  match h : b.1 with
  | .withdrawal withdrawals =>
     
    let k₁InOrder := { s | s : K₁ }.toFinset.sort (·≤·)
    k₁InOrder.attach.map λ s : K₁ ↦ ⟨(s, .Source, withdrawals s), by admit /- proof elided -/
    ⟩
  | .deposit r v | .transfer .. => by admit /- proof elided -/

end Withdrawal

variable [Finite K₁] [LinearOrder K₁]
         [Finite K₂] [LinearOrder K₂]
         [PreWithZero V]
         {b : Block K₁ K₂ C Sigma V}
         {bs : List (Block K₁ K₂ C Sigma V)}
         {π₁ π₂ : BalanceProof K₁ K₂ C Pi V}

def TransactionsInBlock (π : BalanceProof K₁ K₂ C Pi V) (b : Block K₁ K₂ C Sigma V) : List (Τ K₁ K₂ V) := 
  match h : b with
  | .deposit ..    => TransactionsInBlock_deposit ↪b
  | .transfer ..   => TransactionsInBlock_transfer π ↪b
  | .withdrawal .. => TransactionsInBlock_withdrawal ↪b

end Extraction

section e

def e (i : Kbar K₁ K₂) : Kbar K₁ K₂ → ℤ := λ j ↦ if i = j then 1 else 0

variable {i j : Kbar K₁ K₂}

end e

section WithStructuredTypes

section v'

variable [Zero V] [Lattice V] -- NB `PreWithZero V` is implied as `CompleteLattice V` gives `Preorder V`.

def v' (v : V₊) (b : S K₁ K₂ V) (s : Kbar K₁ K₂) : V₊ :=
  match h : s with
  | .Source => v
  | .key _  => ⟨v ⊓ b s, by admit /- proof elided -/
  ⟩

variable {v : V₊} {b : S K₁ K₂ V} {s : Kbar K₁ K₂}

end v'

section Fc

variable [Lattice V]
         [AddCommGroup V]
         [CovariantClass V V (· + ·) (· ≤ ·)]
         [CovariantClass V V (Function.swap (· + ·)) (· ≤ ·)]

def fc (τcXb : Τc K₁ K₂ V × S K₁ K₂ V) : S K₁ K₂ V :=
  ⟨λ k : Kbar K₁ K₂ ↦
    match τcXb with
    | ⟨⟨⟨⟨s, r, v⟩, _⟩, hτ⟩, b⟩ =>
      let v' := v' (v.get hτ) b s
      b k + (e r - e s) k • v',
  by admit /- proof elided -/
  ⟩

variable {τc : Τc K₁ K₂ V} {b : S K₁ K₂ V}

end Fc

section Order

variable [Lattice V] [AddCommGroup V]

end Order

section BoundedBelow

variable [Lattice V] [AddCommGroup V]

abbrev boundedBelow (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) :=
  { a : Τc K₁ K₂ V × S K₁ K₂ V | (T, b) ≤ (↑a.1, a.2) }

end BoundedBelow

section LGroup

variable [Lattice V] [AddCommGroup V]
         [CovariantClass V V (· + ·) (· ≤ ·)]
         [CovariantClass V V (Function.swap (· + ·)) (· ≤ ·)]

def V' (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) (k : Kbar K₁ K₂) : Set V :=
  { v : V | v ∈ (fc · k) '' boundedBelow b T }

section f

def f' (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) : S K₁ K₂ V := 
  ⟨
    λ k ↦
      match h : T with
      | ⟨(_, _, .some _), hT⟩ => fc (⟨T, by admit /- proof elided -/
      ⟩, b) k
      | ⟨(s, _, .none), _⟩ => if k = s then 0 else b k,
    by admit /- proof elided -/
  ⟩

def exists_inf (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) : { s : S K₁ K₂ V // ∀ k, IsGLB (V' b T k) (s k) } :=
  ⟨
    f' b T,
    λ k ↦
      have f'_codomain : (f' b T) k ∈ V' b T k := by admit /- proof elided -/
  ⟩

def infV (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) (k : Kbar K₁ K₂) :
  InfSet V where
    sInf := λ s ↦ if s = V' b T k
                  then (exists_inf b T).1 k
                  else 0

def f (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) : S K₁ K₂ V :=
  ⟨
    λ k ↦
      have : InfSet V := infV b T k
      ⨅ x : boundedBelow b T, fc x.1 k,
    by admit /- proof elided -/
  ⟩

end f

section fStar

variable {s : S K₁ K₂ V}

end fStar

variable [Finite K₁] [LinearOrder K₁]
         [Finite K₂] [LinearOrder K₂]

section LocalProperties

variable {σ : S K₁ K₂ V}
         {π : BalanceProof K₁ K₂ C Pi V}
         {T : Τ K₁ K₂ V}
         {b : Block K₁ K₂ C Sigma V}
         {Sigma : Type}

-- Repository definitions used:
content: local macro:max (priority := high) "↪" b:term : term => `(⟨$b, by admit /- proof elided -/
⟩)

content: inductive Block (K₁ K₂ : Type) (C Sigma : Type) (V : Type) [PreWithZero V] where
   
  | deposit (recipient : K₂) (amount : V₊)
   
  | transfer (aggregator : K₁) (extradata : ExtraDataT) (commitment : C) (senders : List K₂) (sigma : Sigma)
   
  | withdrawal (withdrawals : K₁ → V₊)

content: abbrev isWithdrawalBlock (b : Block K₁ K₂ C Sigma V) := b matches (Block.withdrawal _)

content: abbrev ExtraDataT : Type := ℕ

content: class abbrev PreWithZero (α : Type) := Preorder α, Zero α

content: def NonNeg (α : Type) [PreWithZero α] := { a : α // 0 ≤ a }

content: noncomputable def attach (α : Type) [Finite α] : UniquelyIndexed α :=
  have := Finite.exists_equiv_fin α
  this.choose_spec.some.toEmbedding

content: abbrev UniquelyIndexed (α : Type) [Finite α] : Type := α ↪ !α

content: abbrev UniqueTokenT (α : Type) [Finite α] : Type := Fin (Finite.exists_equiv_fin α |>.choose)

content: abbrev Τ (K₁ K₂ V : Type) [PreWithZero V] := { τ : Τ' K₁ K₂ V // τ.isValid }

content: abbrev Τ' (K₁ K₂ V : Type) [PreWithZero V] := Kbar K₁ K₂ × Kbar K₁ K₂ × Option V₊

content: inductive Kbar (K₁ K₂ : Type) where
  | key (k : Key K₁ K₂)
  | Source
deriving DecidableEq

content: abbrev Key (K₁ K₂ : Type) := K₁ ⊕ K₂

content: class Injective {α ω : Type} (f : α → ω) where
  h : ComputationallyInfeasible (¬ Function.Injective f)

content: abbrev BalanceProof (K₁ K₂ : Type) [Finite K₁] [Finite K₂]
                    (C Pi V : Type) [PreWithZero V] : Type :=
  Dict (C × K₂) ((Pi × ExtraDataT) × TransactionBatch K₁ K₂ V) 

content: abbrev TransactionBatch (K₁ : Type) [Finite K₁]
                        (K₂ : Type) [Finite K₂]
                        (V : Type) [PreWithZero V] :=
  Key K₁ K₂ → V₊

content: abbrev Dict (α ω : Type) : Type := α → Option ω

content: def lexLe (a b : K₂ × Key K₁ K₂) : Prop :=
  a.1 < b.1 ∨ (a.1 = b.1 ∧ a.2 ≤ b.2)

content: abbrev isTransferBlock (b : Block K₁ K₂ C Sigma V) := b matches (Block.transfer _ _ _ _ _)

content: def keys (m : Dict α ω) : Set α := { x | Dict.is_mem m x }

content: def Dict.is_mem (m : Dict α ω) (x : α) : Prop := (m x).isSome

content: abbrev keys (ct : CommitT C K Pi) := ct.dict.keys

content: structure CommitT (C K Pi : Type) where
  commitment : C
  dict       : Dict K Pi

content: abbrev keysUneq (k₂ : K₂) (k : Key K₁ K₂) : Prop :=
  match k with
  | .inl _   => True
  | .inr k₂' => k₂ ≠ k₂'

content: abbrev isDepositBlock (b : Block K₁ K₂ C Sigma V) := b matches (Block.deposit _ _) 

content: abbrev S (K₁ K₂ V : Type) [PreWithZero V] := { s : S' K₁ K₂ V // s.isValid }

content: abbrev S' (K₁ K₂ V : Type) := Kbar K₁ K₂ → V

content: abbrev Τc (K₁ K₂ V : Type) [PreWithZero V] : Type := { τ : Τ K₁ K₂ V // τ.isComplete }

content: def isComplete (τ : Τ K₁ K₂ V) :=
  match τ with | ⟨(_, _, v), _⟩ => v.isSome

content: def getWithdrawal (b : Block K₁ K₂ C Sigma V) (_h : b.isWithdrawalBlock) : K₁ → V₊ :=
  match b with | .withdrawal B => B

content: infix:50 " ≠ₖ " => Key.keysUneq 

content: prefix:max "!" => UniqueTokenT

content: postfix:max "₊" => NonNeg

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: def TransactionsInBlock_deposit
  (b : { b : Block K₁ K₂ C Sigma V // b.isDepositBlock }) : List (Τ K₁ K₂ V) :=
  match h : b.1 with
  | .deposit r v => [⟨(.Source, r, v), by admit /- proof elided -/
  ⟩]
  | .withdrawal .. | .transfer .. => by admit /- proof elided -/

content: def TransactionsInBlock_transfer 
  (π : BalanceProof K₁ K₂ C Pi V) (b : { b : Block K₁ K₂ C Sigma V // b.isTransferBlock }) : List (Τ K₁ K₂ V) :=
  match h : b.1 with
  | .transfer _ _ commitment S _ =>
     
    let senderRecipient : Finset (K₂ × Key K₁ K₂) := { (k₂, k) | (k₂ : K₂) (k : Key K₁ K₂) (_h : k₂ ≠ₖ k) }
    let sorted : List (K₂ × Key K₁ K₂) := senderRecipient.sort Key.lexLe 
     
    let v (s : K₂) (r : Key K₁ K₂) : Option V₊ :=
      if s ∉ S
      then .some 0
      else 
        if h : (commitment, s) ∈ π.keys
        then let (_, t) := π[(commitment, s)]
             t r
        else .none
    sorted.attach.map λ ⟨(s, r), h⟩ ↦ ⟨(s, r, v s r), by admit /- proof elided -/
    ⟩
  | .deposit .. | .withdrawal .. => by admit /- proof elided -/

content: def TransactionsInBlock_withdrawal 
  (b : { b : Block K₁ K₂ C Sigma V // b.isWithdrawalBlock }) : List (Τ K₁ K₂ V) :=
  match h : b.1 with
  | .withdrawal withdrawals =>
     
    let k₁InOrder := { s | s : K₁ }.toFinset.sort (·≤·)
    k₁InOrder.attach.map λ s : K₁ ↦ ⟨(s, .Source, withdrawals s), by admit /- proof elided -/
    ⟩
  | .deposit r v | .transfer .. => by admit /- proof elided -/

content: def TransactionsInBlock (π : BalanceProof K₁ K₂ C Pi V) (b : Block K₁ K₂ C Sigma V) : List (Τ K₁ K₂ V) := 
  match h : b with
  | .deposit ..    => TransactionsInBlock_deposit ↪b
  | .transfer ..   => TransactionsInBlock_transfer π ↪b
  | .withdrawal .. => TransactionsInBlock_withdrawal ↪b

content: def e (i : Kbar K₁ K₂) : Kbar K₁ K₂ → ℤ := λ j ↦ if i = j then 1 else 0

content: def v' (v : V₊) (b : S K₁ K₂ V) (s : Kbar K₁ K₂) : V₊ :=
  match h : s with
  | .Source => v
  | .key _  => ⟨v ⊓ b s, by admit /- proof elided -/
  ⟩

content: def fc (τcXb : Τc K₁ K₂ V × S K₁ K₂ V) : S K₁ K₂ V :=
  ⟨λ k : Kbar K₁ K₂ ↦
    match τcXb with
    | ⟨⟨⟨⟨s, r, v⟩, _⟩, hτ⟩, b⟩ =>
      let v' := v' (v.get hτ) b s
      b k + (e r - e s) k • v',
  by admit /- proof elided -/
  ⟩

content: abbrev boundedBelow (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) :=
  { a : Τc K₁ K₂ V × S K₁ K₂ V | (T, b) ≤ (↑a.1, a.2) }

content: def V' (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) (k : Kbar K₁ K₂) : Set V :=
  { v : V | v ∈ (fc · k) '' boundedBelow b T }

content: def f' (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) : S K₁ K₂ V := 
  ⟨
    λ k ↦
      match h : T with
      | ⟨(_, _, .some _), hT⟩ => fc (⟨T, by admit /- proof elided -/
      ⟩, b) k
      | ⟨(s, _, .none), _⟩ => if k = s then 0 else b k,
    by admit /- proof elided -/
  ⟩

content: def exists_inf (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) : { s : S K₁ K₂ V // ∀ k, IsGLB (V' b T k) (s k) } :=
  ⟨
    f' b T,
    λ k ↦
      have f'_codomain : (f' b T) k ∈ V' b T k := by admit /- proof elided -/
  ⟩

content: def infV (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) (k : Kbar K₁ K₂) :
  InfSet V where
    sInf := λ s ↦ if s = V' b T k
                  then (exists_inf b T).1 k
                  else 0

content: def f (b : S K₁ K₂ V) (T : Τ K₁ K₂ V) : S K₁ K₂ V :=
  ⟨
    λ k ↦
      have : InfSet V := infV b T k
      ⨅ x : boundedBelow b T, fc x.1 k,
    by admit /- proof elided -/
  ⟩

-- Local lemmas used:
content: @[simp]
lemma e_def : e i = λ j ↦ if i = j then 1 else 0

content: @[simp]
lemma v'_key_eq_meet {k : Key K₁ K₂} : v' v b (Kbar.key k) = ⟨v ⊓ b k, by simp⟩

content: private lemma V'_eq_range {b : S K₁ K₂ V} {T : Τ K₁ K₂ V} {k : Kbar K₁ K₂} :
  V' b T k =
  Set.range λ (x : { x : (Τc K₁ K₂ V × S K₁ K₂ V) // (T, b) ≤ (↑x.1, x.2) }) ↦ fc ↑x k

content: lemma f_eq_f' : f = f' (K₁ := K₁) (K₂ := K₂) (V := V)

content: omit [Finite K₂] [LinearOrder K₂] in
private lemma f_withdrawal_block_source_aux {l : List K₁}
  (h₀ : l.Nodup) (h : b.isWithdrawalBlock) :
  (List.foldl f' σ
    (List.map (λ s : K₁ ↦ ⟨(s, Kbar.Source, some (b.getWithdrawal h s)), by unfold Τ'.isValid; aesop⟩) l)).1
                            .Source = σ .Source + ∑ x : K₁, if x ∈ l then (↑(b.getWithdrawal h x) ⊓ σ x) else 0

-- Target theorem and proof:
lemma f_withdrawal_block_source (h : b.isWithdrawalBlock) :
  ((TransactionsInBlock π b).foldl f σ) .Source = σ .Source + ∑ k : K₁, (b.getWithdrawal h k).1 ⊓ σ k
:= by
  simp only [TransactionsInBlock]
  split <;> [simp at h; simp at h; skip]
  next B =>
  simp only [f_eq_f', TransactionsInBlock_withdrawal, List.pure_def, List.bind_eq_flatMap,
    exists_eq, Set.setOf_true, Set.toFinset_univ, Finset.mem_sort, Finset.mem_univ, forall_const,
    List.flatMap_subtype, List.unattach_attach, List.flatMap_singleton', Block.getWithdrawal]
  have : (Block.withdrawal B).getWithdrawal h = B := by ext k; simp [Block.getWithdrawal]
  simp_rw [←this]
  rw [f_withdrawal_block_source_aux (by simp)]
  simp [Finset.mem_sort]