-- Blase/Blase/Fast/Circuit.lean

import Mathlib.Data.Finset.Card

import Mathlib.Data.List.Pi

import Mathlib.Data.Finset.Union

import Mathlib.Data.Fin.Basic

import Mathlib.Data.Finset.Basic

import Mathlib.Data.Finset.Defs

import Mathlib.Data.Fintype.Basic

open Std Sat AIG

inductive Circuit (α : Type u) : Type u
  | tru : Circuit α
  | fals : Circuit α
   
  | var : (positive: Bool) → α → Circuit α
  | and : Circuit α → Circuit α → Circuit α
  | or : Circuit α → Circuit α → Circuit α
  | xor : Circuit α → Circuit α → Circuit α
deriving Repr, DecidableEq

open Lean in

namespace Circuit

variable {α : Type u} {β : Type v}

def vars [DecidableEq α] : Circuit α → List α
  | tru => []
  | fals => []
  | var _ x => [x]
  | and c₁ c₂ => (vars c₁ ++ vars c₂).dedup
  | or c₁ c₂ => (vars c₁ ++ vars c₂).dedup
  | xor c₁ c₂ => (vars c₁ ++ vars c₂).dedup

def varsFinset [DecidableEq α] (c : Circuit α) : Finset α :=
  ⟨c.vars, nodup_vars c⟩

def simplifyAnd : Circuit α → Circuit α → Circuit α
  | tru, c => c
  | c, tru => c
  | fals, _ => fals
  | _, fals => fals
  | c₁, c₂ => and c₁ c₂

def simplifyOr : Circuit α → Circuit α → Circuit α
  | tru, _ => tru
  | _, tru => tru
  | fals, c => c
  | c, fals => c
  | c₁, c₂ => or c₁ c₂

def simplifyNot : Circuit α → Circuit α
  | tru => fals
  | fals => tru
  | xor a b => xor (simplifyNot a) b
  | and a b => or (simplifyNot a) (simplifyNot b)
  | or a b => and (simplifyNot a) (simplifyNot b)
  | var b a => var (!b) a

@[simp]
def simplifyXor : Circuit α → Circuit α → Circuit α
  | fals, c => c
  | c, fals => c
  | tru, c => ~~~ c
  | c, tru => ~~~ c
  | c₁, c₂ => xor c₁ c₂

def assignVars [DecidableEq α] :
    ∀ (c : Circuit α) (_f : ∀ (a : α) (_ha : a ∈ c.vars), β ⊕ Bool), Circuit β
  | tru, _ => tru
  | fals, _ => fals
  | var b x, f =>
    Sum.elim
      (var b)
      (λ c : Bool => if Bool.xor b c then fals else tru)
      (f x (by admit /- proof elided -/
      ))
  | and c₁ c₂, f => (assignVars c₁ (λ x hx => f x (by admit /- proof elided -/
  ))) &&&
                    (assignVars c₂ (λ x hx => f x (by admit /- proof elided -/
                    )))
  | or c₁ c₂, f =>  (assignVars c₁ (λ x hx => f x (by admit /- proof elided -/
  ))) |||
                    (assignVars c₂ (λ x hx => f x (by admit /- proof elided -/
                    )))
  | xor c₁ c₂, f => (assignVars c₁ (λ x hx => f x (by admit /- proof elided -/
  ))) ^^^
                    (assignVars c₂ (λ x hx => f x (by admit /- proof elided -/
                    )))

-- Repository definitions used:
content: noncomputable def f : (ZMod q)[X] := X^(2^n) + 1

-- 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: inductive Circuit (α : Type u) : Type u
  | tru : Circuit α
  | fals : Circuit α
   
  | var : (positive: Bool) → α → Circuit α
  | and : Circuit α → Circuit α → Circuit α
  | or : Circuit α → Circuit α → Circuit α
  | xor : Circuit α → Circuit α → Circuit α
deriving Repr, DecidableEq

content: def vars [DecidableEq α] : Circuit α → List α
  | tru => []
  | fals => []
  | var _ x => [x]
  | and c₁ c₂ => (vars c₁ ++ vars c₂).dedup
  | or c₁ c₂ => (vars c₁ ++ vars c₂).dedup
  | xor c₁ c₂ => (vars c₁ ++ vars c₂).dedup

content: def varsFinset [DecidableEq α] (c : Circuit α) : Finset α :=
  ⟨c.vars, nodup_vars c⟩

content: def simplifyAnd : Circuit α → Circuit α → Circuit α
  | tru, c => c
  | c, tru => c
  | fals, _ => fals
  | _, fals => fals
  | c₁, c₂ => and c₁ c₂

content: def simplifyOr : Circuit α → Circuit α → Circuit α
  | tru, _ => tru
  | _, tru => tru
  | fals, c => c
  | c, fals => c
  | c₁, c₂ => or c₁ c₂

content: def simplifyNot : Circuit α → Circuit α
  | tru => fals
  | fals => tru
  | xor a b => xor (simplifyNot a) b
  | and a b => or (simplifyNot a) (simplifyNot b)
  | or a b => and (simplifyNot a) (simplifyNot b)
  | var b a => var (!b) a

content: @[simp]
def simplifyXor : Circuit α → Circuit α → Circuit α
  | fals, c => c
  | c, fals => c
  | tru, c => ~~~ c
  | c, tru => ~~~ c
  | c₁, c₂ => xor c₁ c₂

content: def assignVars [DecidableEq α] :
    ∀ (c : Circuit α) (_f : ∀ (a : α) (_ha : a ∈ c.vars), β ⊕ Bool), Circuit β
  | tru, _ => tru
  | fals, _ => fals
  | var b x, f =>
    Sum.elim
      (var b)
      (λ c : Bool => if Bool.xor b c then fals else tru)
      (f x (by admit /- proof elided -/
      ))
  | and c₁ c₂, f => (assignVars c₁ (λ x hx => f x (by admit /- proof elided -/
  ))) &&&
                    (assignVars c₂ (λ x hx => f x (by admit /- proof elided -/
                    )))
  | or c₁ c₂, f =>  (assignVars c₁ (λ x hx => f x (by admit /- proof elided -/
  ))) |||
                    (assignVars c₂ (λ x hx => f x (by admit /- proof elided -/
                    )))
  | xor c₁ c₂, f => (assignVars c₁ (λ x hx => f x (by admit /- proof elided -/
  ))) ^^^
                    (assignVars c₂ (λ x hx => f x (by admit /- proof elided -/
                    )))

-- Local lemmas used:
content: lemma mem_varsFinset [DecidableEq α] {c : Circuit α} :
    ∀ {x : α}, x ∈ c.varsFinset ↔ x ∈ c.vars

content: theorem varsFinset_and [DecidableEq α] (c₁ c₂ : Circuit α) :
    (varsFinset (c₁ &&& c₂)) ⊆ (varsFinset c₁ ∪ varsFinset c₂)

content: theorem varsFinset_or [DecidableEq α] (c₁ c₂ : Circuit α) :
    (varsFinset (c₁ ||| c₂)) ⊆ (varsFinset c₁ ∪ varsFinset c₂)

content: @[simp]
theorem simplifyNot_eq_complement (c : Circuit α) :
    simplifyNot c = ~~~ c

content: theorem vars_simplifyXor [DecidableEq α] (c₁ c₂ : Circuit α) :
    (vars (simplifyXor c₁ c₂)) ⊆ (vars c₁ ++ vars c₂).dedup

content: theorem varsFinset_simplifyXor [DecidableEq α] (c₁ c₂ : Circuit α) :
    (varsFinset (simplifyXor c₁ c₂)) ⊆ (varsFinset c₁ ∪ varsFinset c₂)

content: theorem varsFinset_xor [DecidableEq α] (c₁ c₂ : Circuit α) :
    (varsFinset (c₁ ^^^ c₂)) ⊆ (varsFinset c₁ ∪ varsFinset c₂)

-- Target theorem and proof:
lemma varsFinset_assignVars [DecidableEq α] [DecidableEq β] :
    ∀ (c : Circuit α) (f : ∀ (a : α) (_ha : a ∈ c.vars), β ⊕ Bool),
      (c.assignVars f).varsFinset ⊆ c.varsFinset.biUnion
        (fun a => if ha : a ∈ c.vars
                  then
                    match f a ha with
                    | Sum.inl b => {b}
                    | Sum.inr _ => ∅
                  else ∅)
  | tru, _ => by simp [assignVars, varsFinset, vars]
  | fals, _ => by simp [vars, assignVars, varsFinset]
  | var c v, f => by
    intro x
    simp [assignVars, varsFinset, vars]
    split <;>
    simp [*, vars]
    split_ifs <;> simp [vars]
  | and c₁ c₂, f => by
    intro x
    simp only [assignVars, Finset.mem_biUnion]
    intro hx
    replace hx := varsFinset_and _ _ hx
    simp only [Finset.mem_union] at hx
    cases hx with
    | inl hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
    | inr hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
  | or c₁ c₂, f => by
    intro x
    simp only [assignVars, Finset.mem_biUnion]
    intro hx
    replace hx := varsFinset_or _ _ hx
    simp only [Finset.mem_union] at hx
    cases hx with
    | inl hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
    | inr hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
  | xor c₁ c₂, f => by
    intro x
    simp only [assignVars, Finset.mem_biUnion]
    intro hx
    replace hx := varsFinset_xor _ _ hx
    simp only [Finset.mem_union] at hx
    cases hx with
    | inl hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
    | inr hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
:= varsFinset_and _ _ hx
    simp only [Finset.mem_union] at hx
    cases hx with
    | inl hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
    | inr hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
  | or c₁ c₂, f => by
    intro x
    simp only [assignVars, Finset.mem_biUnion]
    intro hx
    replace hx := varsFinset_or _ _ hx
    simp only [Finset.mem_union] at hx
    cases hx with
    | inl hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
    | inr hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
  | xor c₁ c₂, f => by
    intro x
    simp only [assignVars, Finset.mem_biUnion]
    intro hx
    replace hx := varsFinset_xor _ _ hx
    simp only [Finset.mem_union] at hx
    cases hx with
    | inl hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
    | inr hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2