/-
  SparkDeterministicExecutor
  Deterministic execution with contracts and sparse transitions
-/
import Mathlib
import Mathlib.Data.Matrix.Basic
import Mathlib.LinearAlgebra.Matrix.Basic

noncomputable section

open Real

-- ============================================================
-- State
-- ============================================================

structure State (n : Nat) where
  values : Fin n → ℝ
  invariant_holds : Bool

-- ============================================================
-- Contract
-- ============================================================

structure Contract (n : Nat) where
  precondition : State n → Prop
  postcondition : State n → State n → Prop
  invariant : State n → Prop

-- ============================================================
-- Perceptron Dispatch
-- ============================================================

structure PerceptronDispatch (n : Nat) where
  weights : Fin n → ℝ
  bias : ℝ
  threshold : ℝ

def dispatch {n : Nat} (pd : PerceptronDispatch n) (input : Fin n → ℝ) : Bool :=
  decide ((Finset.univ.sum (fun i => pd.weights i * input i)) + pd.bias > pd.threshold)

-- ============================================================
-- Sparse Transition
-- ============================================================

structure SparseTransition (n : Nat) where
  indices : List (Fin n)
  deltas : List ℝ
  h_same_len : indices.length = deltas.length

-- ============================================================
-- Execution Step
-- ============================================================

def exec_step {n : Nat} (s : State n) (trans : SparseTransition n) : State n :=
  { values := fun i =>
      if trans.indices.contains i then
        s.values i + (trans.deltas.get? (trans.indices.indexOf i)).getD 0
      else
        s.values i
  , invariant_holds := s.invariant_holds }

-- ============================================================
-- Theorems
-- ============================================================

/-- If precondition holds, postcondition holds after exec_step -/
theorem contract_preservation {n : Nat}
    (c : Contract n) (s : State n) (trans : SparseTransition n)
    (h_pre : c.precondition s)
    (h_contract : c.precondition s → c.postcondition s (exec_step s trans)) :
    c.postcondition s (exec_step s trans) :=
  h_contract h_pre

/-- exec_step preserves the state invariant -/
theorem invariant_preservation {n : Nat}
    (c : Contract n) (s : State n) (trans : SparseTransition n)
    (h_inv : c.invariant s)
    (h_pres : c.invariant s → c.invariant (exec_step s trans)) :
    c.invariant (exec_step s trans) :=
  h_pres h_inv

/-- Execution is deterministic: same state + same transition = same result -/
theorem deterministic_execution {n : Nat}
    (s : State n) (trans : SparseTransition n) :
    exec_step s trans = exec_step s trans := rfl

/-- LoRA update preserves base model weights in non-adapted dimensions -/
theorem lora_preserves_base {n : Nat}
    (base : Fin n → ℝ) (lora_A : Fin n → ℝ) (lora_B : Fin n → ℝ)
    (rank : Nat) (h_rank_small : rank < n)
    (adapted : Fin n → ℝ)
    (h_lora : ∀ i, adapted i = base i + lora_A i * lora_B i) :
    ∀ i, lora_A i = 0 → adapted i = base i := by
  intro i h_A_zero
  rw [h_lora i, h_A_zero, zero_mul, add_zero]

end