import Mathlib

set_option autoImplicit false

namespace Probe

def Arms : Fin 3 → ℕ
  | ⟨0, _⟩ => 4
  | ⟨1, _⟩ => 4
  | ⟨_, _⟩ => 3

theorem Arms_zero : Arms 0 = 4 := rfl
theorem Arms_one : Arms 1 = 4 := rfl
theorem Arms_two : Arms 2 = 3 := rfl
theorem Arms_pos (i : Fin 3) : 0 < Arms i := by fin_cases i <;> decide

inductive V : Type where
  | torso : V
  | arm : (i : Fin 3) → Fin (Arms i) → V
  deriving DecidableEq

inductive W : Type where
  | torso : W
  | arm : (i : Fin 3) → Fin (Arms i) → W
  deriving DecidableEq, Fintype

def V.toW : V → W
  | .torso => W.torso
  | .arm i v => W.arm i v

def V.ofW : W → V
  | .torso => V.torso
  | .arm i v => V.arm i v

instance : Fintype V := Fintype.ofEquiv W {
  toFun := V.ofW
  invFun := V.toW
  left_inv := by intro x; cases x <;> rfl
  right_inv := by intro x; cases x <;> rfl }

theorem V_card : Fintype.card V = 12 := by decide

inductive ED : V → V → Prop where
  | torso (i : Fin 3) : ED V.torso (V.arm i ⟨0, Arms_pos i⟩)
  | arm (i : Fin 3) (v : Fin (Arms i)) (h : v.val + 1 < Arms i) :
      ED (V.arm i v) (V.arm i ⟨v.val + 1, h⟩)

def EDB : V → V → Bool
  | .torso, .arm i v => decide (v.val = 0)
  | .arm i v, .arm j w => decide (i = j) && decide (w.val = v.val + 1)
  | _, _ => false

theorem edb_torso (y : V) : EDB V.torso y = true ↔ ED V.torso y := by
  cases y with
  | torso => exact ⟨fun h => absurd h (by decide), fun h => by cases h⟩
  | arm i v =>
      show decide (v.val = 0) = true ↔ ED V.torso (V.arm i v)
      rw [decide_eq_true_eq]
      exact ⟨fun hv => by
          have h : v = ⟨0, Arms_pos i⟩ := Fin.ext hv
          rw [h]; exact ED.torso i,
        fun h => by cases h with | torso j => rfl⟩

theorem edb_arm (i j : Fin 3) (v : Fin (Arms i)) (w : Fin (Arms j)) :
    EDB (V.arm i v) (V.arm j w) = true ↔ ED (V.arm i v) (V.arm j w) := by
  show (decide (i = j) && decide (w.val = v.val + 1)) = true ↔ _
  rw [Bool.and_eq_true, decide_eq_true_eq, decide_eq_true_eq]
  constructor
  · rintro ⟨hij, hw⟩
    subst hij
    have hlt : v.val + 1 < Arms i := by have := w.isLt; omega
    have h : w = ⟨v.val + 1, hlt⟩ := Fin.ext hw
    rw [h]; exact ED.arm i v hlt
  · intro h
    cases h with
    | arm k u hu => exact ⟨rfl, rfl⟩

theorem edb_arm_torso (i : Fin 3) (v : Fin (Arms i)) : EDB (V.arm i v) V.torso = false := rfl

theorem edb_not_torso (y : V) : EDB y V.torso = false := by
  cases y <;> rfl

theorem edb_iff (x y : V) : EDB x y = true ↔ ED x y := by
  match x, y with
  | .torso, .torso =>
      exact ⟨fun h => absurd h (by decide), fun h => by cases h⟩
  | .torso, .arm i v => exact edb_torso _
  | .arm i v, .torso =>
      exact ⟨fun h => absurd (h ▸ edb_not_torso (V.arm i v)) (by decide),
        fun h => by cases h⟩
  | .arm i v, .arm j w => exact edb_arm i j v w

instance : DecidableRel ED := fun x y => decidable_of_iff (EDB x y = true) (edb_iff x y)

example : ¬ ED V.torso V.torso := by decide

end Probe

-- ===== additions: edge closure, matrix, multiset =====
def E (x y : V) : Prop := ED x y ∨ ED y x

instance : DecidableRel E := fun x y => inferInstanceAs (Decidable (ED x y ∨ ED y x))

def entry (x y : V) : ℕ := if x = y then 1 else if E x y then 3 else 2

theorem entry_self (x : V) : entry x x = 1 := by
  unfold entry; rw [if_pos rfl]

theorem entry_of_ne {x y : V} (h : x ≠ y) : entry x y = if E x y then 3 else 2 := by
  unfold entry; rw [if_neg h]

theorem entry_of_edge {x y : V} (h : E x y) : entry x y = 3 := by
  have hne : x ≠ y := fun hc => (by cases h with
    | inl h1 => exact absurd hc (by
        intro he; subst he; cases h1 <;> decide)
    | inr h2 => exact absurd hc.symm (by
        intro he; subst he; cases h2 <;> decide)) hc
  rw [entry_of_ne hne, if_pos h]

example : E V.torso (V.arm 0 ⟨0, Arms_pos 0⟩) := Or.inl (ED.torso 0)

example (x : V) : ¬ E x x := by
  intro h
  cases h with
  | inl h1 => cases h1 <;> decide
  | inr h2 => cases h2 <;> decide

theorem E_symm {x y : V} (h : E x y) : E y x := h.symm

/-- the matrix -/
def M : Matrix V V ℕ := Matrix.of entry

theorem M_symm : M.IsSymm := Matrix.IsSymm.ext fun i j => by
  show entry j i = entry i j
  by_cases h : i = j
  · subst h; rw [entry_self, entry_self]
  · rw [entry_of_ne (Ne.symm h), entry_of_ne h]
    by_cases he : E i j
    · rw [if_pos he, if_pos (E_symm he)]
    · rw [if_neg he, if_neg (fun hc => he (E_symm hc))]

theorem M_off_diag (i j : V) (hij : i ≠ j) : entry i j ≠ 1 := by
  rw [entry_of_ne hij]
  by_cases he : E i j <;> simp [he]

@[simp] theorem M_apply (i j : V) : M i j = entry i j := Matrix.of_apply ..

example :
    (Finset.univ.val.map (fun p : V × V => M p.1 p.2)) =
      (Multiset.replicate 12 1 + Multiset.replicate 22 3 + Multiset.replicate 110 2 : Multiset ℕ) := by
  decide

end Probe
