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

-- 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: 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

-- Local lemmas used:
content: lemma length_TransactionsInBlock_transfer
  {b : { b : Block K₁ K₂ C Sigma V // b.isTransferBlock }} :
  ∀ (π₁ π₂ : BalanceProof K₁ K₂ C Pi V),
    (TransactionsInBlock_transfer π₁ b).length =
    (TransactionsInBlock_transfer π₂ b).length

content: lemma length_transactionsInBlock :
  (TransactionsInBlock π₁ b).length = (TransactionsInBlock π₂ b).length

-- Target theorem and proof:
lemma receiver_transactionsInBlock :
  (TransactionsInBlock π₁ b).map (λ s ↦ s.1.2.1) =
  (TransactionsInBlock π₂ b).map (λ s ↦ s.1.2.1)
:= by
  apply List.ext_get (by simp; rw [length_transactionsInBlock])
  intros n h₁ h₂
  simp; unfold TransactionsInBlock
  match b with
  | Block.deposit ..    => simp [TransactionsInBlock_deposit]
  | Block.transfer ..   => simp [TransactionsInBlock_transfer]
  | Block.withdrawal .. => simp [TransactionsInBlock_withdrawal]