-- FVIntmax/Lemma3.lean

import FVIntmax.Balance

namespace Intmax

section Lemma3

variable
  {Pi C Sigma : Type}
  {K₁ : Type} [Finite K₁]
  {K₂ : Type} [Finite K₂]
  {V : Type} [Lattice V] [AddCommGroup V]
             [CovariantClass V V (· + ·) (· ≤ ·)]
             [CovariantClass V V (Function.swap (· + ·)) (· ≤ ·)]

variable {Tc : Τc K₁ K₂ V} {T : Τ K₁ K₂ V} {b : S K₁ K₂ V} {Tstar : List (Τ K₁ K₂ V)}

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

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

content: def fStar (Ts : List (Τ K₁ K₂ V)) (s₀ : S K₁ K₂ V) : S K₁ K₂ V :=
  Ts.foldl f s₀

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 -/
  ⟩

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 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 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: def e (i : Kbar K₁ K₂) : Kbar K₁ K₂ → ℤ := λ j ↦ if i = j then 1 else 0

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

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

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: class abbrev PreWithZero (α : Type) := Preorder α, Zero α

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

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: def NonNeg (α : Type) [PreWithZero α] := { a : α // 0 ≤ a }

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 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: class Injective {α ω : Type} (f : α → ω) where
  h : ComputationallyInfeasible (¬ Function.Injective f)

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 isComplete (τ : Τ K₁ K₂ V) :=
  match τ with | ⟨(_, _, v), _⟩ => v.isSome

content: def TransactionsInBlocks
  (π : BalanceProof K₁ K₂ C Pi V) (bs : List (Block K₁ K₂ C Sigma V)) : List (Τ K₁ K₂ V) :=
  (bs.map (TransactionsInBlock π)).flatten

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 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: 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: 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: 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: 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: 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: abbrev isDepositBlock (b : Block K₁ K₂ C Sigma V) := b matches (Block.deposit _ _) 

content: def Bal (π : BalanceProof K₁ K₂ C Pi V) (bs : List (Block K₁ K₂ C Sigma V)) : S K₁ K₂ V :=
  fStar (TransactionsInBlocks π bs) (.initial K₁ K₂ V)

content: def initial : Scontract K₁ K₂ V C Sigma := []

content: abbrev Scontract (K₁ K₂ V : Type) [PreWithZero V] (C Sigma : Type) :=
  List (Block K₁ K₂ C Sigma V)

content: def initial (K₁ K₂ V : Type) [PreWithZero V] : S K₁ K₂ V :=
  ⟨S'.initial K₁ K₂ V, S'.isValid_initial⟩

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

content: prefix:max "!" => UniqueTokenT

content: postfix:max "₊" => NonNeg

-- Repository lemmas used:
content: lemma f_IsGLB_of_V' {b : S K₁ K₂ V} {T : Τ K₁ K₂ V} {k : Kbar K₁ K₂} :
  IsGLB (V' b T k) (f b T k)

-- Local definitions used:

-- Local lemmas used:
content: @[simp]
lemma sum_fc_eq_sum : ∑ (k : Kbar K₁ K₂), fc (Tc, b) k = ∑ (k : Kbar K₁ K₂), b k

content: lemma sum_f_le_sum : ∑ (k : Kbar K₁ K₂), f b T k ≤ ∑ (k : Kbar K₁ K₂), b k

content: private lemma sum_fStar_le_zero_aux (h : ∑ (k : Kbar K₁ K₂), b k ≤ 0) :
  ∑ (k : Kbar K₁ K₂), fStar Tstar b k ≤ 0

-- Target theorem and proof:
lemma lemma3 : Bal π bs .Source ≤ 0
:= by
  dsimp [Bal]
  generalize eq : TransactionsInBlocks π _ = blocks
  generalize eq₁ : S.initial K₁ K₂ V = s₀
  generalize eq₂ : fStar blocks s₀ = f
  have : ∑ x ∈ {Kbar.Source}, f x = 
         ∑ x : Kbar K₁ K₂, f x - ∑ x ∈ Finset.univ \ {Kbar.Source}, f x := by simp
  rw [←Finset.sum_singleton (a := Kbar.Source) (f := f), this, sub_nonpos]
  have := sum_fStar_le_zero_aux (Tstar := blocks) (b := s₀)
  have eq₃ : ∑ x : Kbar K₁ K₂, f x ≤ 0 := by aesop
  have eq₄ : 0 ≤ ∑ x ∈ Finset.univ \ {Kbar.Source}, f x := Finset.sum_nonneg λ i ↦ by rcases i <;> aesop
  exact le_trans eq₃ eq₄