ModelCounting.lean

import Init

/- Facts about Nat -/

theorem Nat.le_refl (a : Nat) : a ≤ a := sorry
theorem Nat.le_trans {a b c : Nat} : a ≤ b → b ≤ c → a ≤ c := sorry
theorem Nat.lt_trans {a b c : Nat} : a < b → b < c → a < c := sorry
theorem Nat.le_of_lt {a b : Nat} : a < b → a ≤ b := sorry
theorem Nat.not_le_of_lt {a b : Nat} : a < b → ¬ b ≤ a := sorry
theorem Nat.le_of_not_le {a b : Nat} : ¬ a ≤ b → b ≤ a := sorry
theorem Nat.lt_or_eq {a b : Nat} : a ≤ b → a < b ∨ a = b := sorry
theorem Nat.lt_succ_of_lt {a b : Nat} : a < b → a < b + 1 := sorry
theorem Nat.not_lt_zero (a : Nat) : ¬ a < 0 := sorry
theorem Nat.lt_of_lt_of_le {a b c : Nat} : a < b → b ≤ c → a < c := sorry
theorem Nat.lt_succ_iff_le {a b : Nat} : a < b + 1 ↔ a ≤ b:= sorry
theorem Nat.lt_succ_self (a : Nat) : a < a + 1 := Nat.lt_succ_iff_le.mpr (Nat.le_refl a)

/- Sets and finite sets -/

def Set (α : Type) := α → Prop

namespace Set

def mem (a : α) (s : Set α) := s a
def empty (α : Type) : Set α := fun a => False
def singleton (a : α) := fun b => b = a
def union (s t : Set α) := fun a => mem a s ∨ mem a t
def sep (s : Set α) (P : α → Prop) := fun a => mem a s ∧ P a

def subset (s t : Set α) : Prop := ∀ x, mem x s → mem x t

inductive Finite : Set α → Prop where
  | empty : Finite (empty α)
  | singleton : Finite (singleton a)
  | union {s t : Set α} : Finite s → Finite t → Finite (union s t)
  | subset {s t : Set α} : Finite t → subset s t → Finite s

end Set

def Finset α := { s : Set α // Set.Finite s }

namespace Finset 

def mem (a : α) (s : Finset α) := s.1 a
infix:50 " ∈ " => Finset.mem

def empty (α : Type) : Finset α := ⟨Set.empty α, Set.Finite.empty⟩
def singleton (a : α) : Finset α := ⟨Set.singleton a, Set.Finite.singleton⟩
def union (s t : Finset α) : Finset α := ⟨Set.union s.1 t.1, Set.Finite.union s.2 t.2⟩
infixl:65 " ∪ " => Finset.union

def disjoint (s t : Finset α) := ∀ a, ¬ (a ∈ s ∧ a ∈ t)

theorem not_mem_empty (a : α) : ¬ (a ∈ empty α) := fun h => h
theorem mem_singleton {a b : α} : a ∈ singleton b ↔ a = b := Iff.refl _
theorem mem_union {a : α} {s t : Finset α} : a ∈ union s t ↔ a ∈ s ∨ a ∈ t := Iff.refl _

/- axiomatize cardinality for now -/
axiom card {α : Type} : Finset α → Nat

axiom card_empty {α : Type} : card (empty α) = 0
axiom card_singleton {α : Type} (a : α) : card (singleton a) = 1
axiom card_union {α : Type} {s t : Finset α} (h : disjoint s t) : card (s ∪ t) = card s + card t

end Finset

/- Define itegs-/

inductive IteElt
  | Tr                : IteElt 
  | Fls               : IteElt
  | Var (i : Nat)     : IteElt
  | Ite (c t f : Nat) : IteElt

instance : Inhabited IteElt := ⟨IteElt.Tr⟩

open IteElt

/-- `IteEltBounded x b` says that if `x` is an if-then-else clause, then the branches are less 
than `b`-/
def IteEltBounded : IteElt → Nat → Prop
  | Tr,        _ => True
  | Fls,       _ => True
  | Var i,     _ => True
  | Ite c t f, b => t < b ∧ f < b

-- No need to bound i, since the default value of an `IteElt` is `Tr` 
def Iteg := { I : Array IteElt // ∀ i, IteEltBounded I[i] i }

instance itegCoe (n : Nat) : CoeHead Iteg (Array IteElt) where
  coe v := v.val

/- Semantics -/

def Assignment : Type := Nat → Bool

def evalUpTo (I : Iteg) (v : Assignment) : {m i : Nat} → i < m → Bool
  | 0,       i, hi => False.elim (Nat.not_lt_zero i hi)
  | (m + 1), i, hi => 
      match h:I.val[i] with
        | Tr            => true
        | Fls           => false
        | Var j         => v j
        | Ite c t f     => 
            have bdd : IteEltBounded (Ite c t f) i by { rw ←h; apply I.property }
            have ilem : i ≤ m from Nat.lt_succ_iff_le.mp hi
            cond (v c) (evalUpTo I v (Nat.lt_of_lt_of_le bdd.1 ilem))
                       (evalUpTo I v (Nat.lt_of_lt_of_le bdd.2 ilem))

/-- Evaluates line `j` of `I` wrt `v` -/
def eval (I : Iteg) (v : Assignment) (j : Nat) := evalUpTo I v (Nat.lt_succ_self j)

-- TODO: prove recursion equations for Eval.
-- Note: this just simulates well-founded recursion on nat, which is not available in Lean 4 yet.
theorem eval_rec (I : Iteg) (v : Assignment) (j : Nat) :
  eval I v j = match I.val[j] with 
                 | Tr        => true
                 | Fls       => false
                 | Var k     => v k 
                 | Ite c t f => cond (v c) (eval I v t) (eval I v f) :=
sorry

/- Variables -/

open Finset

def varsUpTo (I : Iteg) : {m i : Nat} → i < m → Finset Nat
  | 0,       i, hi => False.elim (Nat.not_lt_zero i hi)
  | (m + 1), i, hi => 
      match h:I.val[i] with
        | Tr            => empty Nat
        | Fls           => empty Nat
        | Var j         => singleton j
        | Ite c t f     =>
            have bdd : IteEltBounded (Ite c t f) i by { rw ←h; apply I.property }
            have ilem : i ≤ m from Nat.lt_succ_iff_le.mp hi
            singleton c ∪ varsUpTo I (Nat.lt_of_lt_of_le bdd.1 ilem) ∪ 
              varsUpTo I (Nat.lt_of_lt_of_le bdd.2 ilem)

/-- Holds iff line `j` of Iteg `I` depends on variable `k` -/
def vars (I : Iteg) (j : Nat) := varsUpTo I (Nat.lt_succ_self j)

-- TODO: prove recursion equation
theorem variables_rec (I : Iteg) (j : Nat) :
  vars I j = match I.val[j] with 
                   | Tr        => empty Nat
                   | Fls       => empty Nat
                   | Var k     => singleton k
                   | Ite c t f => singleton c ∪ vars I t ∪ vars I f :=
sorry

/- Free itegs -/

def free (I : Iteg) : Prop := 
∀ i, 
  match I.val[i] with
    | Tr => True
    | Fls => True
    | Var k => True
    | Ite c t f => ¬ (c ∈ vars I t) ∧ ¬ (c ∈ vars I f)

/- Counting finite variable assignments -/

def defined_on (v : Assignment) (s : Finset Nat) := ∀ n, ¬ n ∈ s → v n = false

theorem Finite_defined_on (s : Finset Nat) : Set.Finite (fun v : Assignment => defined_on v s) :=
sorry

def restrictedAssignments (s : Finset Nat) : Finset Assignment := 
⟨_, Finite_defined_on s⟩

theorem card_restrictedAssignments (s : Finset Nat) : card (restrictedAssignments s) = 2 ^ card s :=
sorry