fintype.lean

-- TODO(solson): lean/library/init/data/sum/default.lean should include this import upstream.
import init.data.sum.instances

universes u v

variables {α : Type u} {β : Type v} {n : ℕ} {a : α} {b : β} {f : α → β}

def count_cond [decidable_eq α] (a b : α) : ℕ :=
    cond (a = b) 1 0

@[simp]
def count_cond_true [decidable_eq α] {a b : α} : a = b → count_cond a b = 1 :=
    assume h, by simp [count_cond, to_bool_tt h]

@[simp]
def count_cond_false [decidable_eq α] {a b : α} : a ≠ b → count_cond a b = 0 :=
    assume h, by simp [count_cond, to_bool_ff h]

lemma count_cond_of_iff {a b : α} [decidable_eq α] [decidable_eq β] :
    (f a = f b ↔ a = b) → count_cond (f a) (f b) = count_cond a b :=
begin
    assume h,
    unfold count_cond,
    apply congr_fun _ 0,
    apply congr_fun _ 1,
    apply congr_arg cond,
    apply to_bool_congr,
    exact h,
end

def count [decidable_eq α] (a : α) : list α → ℕ
| []      := 0
| (x::xs) := count_cond a x + count xs

@[simp]
theorem count_nil [decidable_eq α] : count a list.nil = 0 := by unfold count

theorem count_append [decidable_eq α] :
    ∀ {xs ys : list α}, count a (xs ++ ys) = count a xs + count a ys
| []      ys := by simp [count]
| (x::xs) ys := by simp [count, count_append]

-- Finite types.
class fintype (α : Type u) [decidable_eq α] : Type u :=
(values : list α)
(values_pred : ∀ (a : α), count a values = 1)

instance empty.decidable_eq : decidable_eq empty.
instance empty.fintype : fintype empty := ⟨[], λ a, by cases a⟩
instance : fintype unit := ⟨[unit.star], λ _, rfl⟩
instance : fintype bool := ⟨[ff, tt], λ b, by cases b; refl⟩

lemma zero_ne_fin_succ : ∀ (x : fin n), 0 ≠ x.succ
| ⟨x, _⟩ := by {
    unfold fin.succ,
    assume h,
    injection h,
    contradiction,
}

lemma mem_map : ∀ {xs : list α}, b ∈ list.map f xs ↔ (∃ a ∈ xs, b = f a)
| [] := by {
    rw list.map,
    rw iff_false_intro (list.not_mem_nil b),
    split,
    { contradiction },
    { apply list.not_bex_nil },
}
| (x::xs) := by {
    rw list.map,
    rw list.mem_cons_iff,
    split,
    {
        assume h,
        cases h,
        case or.inl h {
            exact ⟨x, list.mem_cons_self x xs, h⟩,
        },
        case or.inr h {
            have h' : ∃ (a : α) (H : a ∈ xs), b = f a, from mem_map.mp h,
            exact let ⟨a, h1, h2⟩ := h' in ⟨a, list.mem_cons_of_mem x h1, h2⟩,
        },
    },
    {
        exact assume ⟨a, h1, h2⟩, by {
            rw list.mem_cons_iff at h1,
            cases h1,
            case or.inl h {
                rw [h2, h],
                apply list.mem_cons_self,
            },
            case or.inr h {
                apply list.mem_cons_of_mem,
                apply mem_map.mpr,
                exact ⟨a, h, h2⟩,
            },
        },
    },
}

lemma not_mem_map {xs : list α} : (∀ a ∈ xs, b ≠ f a) → b ∉ list.map f xs :=
    assume h h_mem, let ⟨a, h1, h2⟩ := mem_map.mp h_mem in h a h1 h2

lemma count_eq_zero [decidable_eq α] : ∀ {xs : list α}, a ∉ xs → count a xs = 0
| []      _ := rfl
| (x::xs) h := by {
    rw list.mem_cons_eq at h,
    have h1 : a ≠ x,  from h ∘ or.inl,
    have h2 : a ∉ xs, from h ∘ or.inr,
    rw [count, count_cond_false h1, count_eq_zero h2],
}

lemma count_map [decidable_eq α] [decidable_eq β] (h_inj : ∀ {a b}, f a = f b → a = b) :
     ∀ (xs : list α), count (f a) (list.map f xs) = count a xs
| []      := rfl
| (x::xs) := by {
    rw [list.map, count, count, count_map],
    congr,
    apply count_cond_of_iff,
    exact ⟨h_inj, λ h, by subst h⟩,
}

def fin_values : ∀ (n : ℕ), list (fin n)
| 0     := []
| (n+1) := 0 :: list.map fin.succ (fin_values n)

def fin_values_pred : ∀ (n : ℕ) (a : fin n), count a (fin_values n) = 1 
| 0     := λ f, f.elim0
| (n+1) := λ x,
    if h_eq : x = 0
    then by {
        subst h_eq,
        rw [fin_values, count, count_cond_true rfl, count_eq_zero],
        apply not_mem_map,
        intros,
        apply zero_ne_fin_succ,
    }
    else by {
        rw [fin_values, count, count_cond_false h_eq],
        -- TODO(solson): Extract as an independent theorem.
        have : ∀ x (h : x ≠ 0), x = fin.succ (fin.pred x h), assume x h, by {
            cases x,
            cases val,
            contradiction,
            unfold fin.pred,
            unfold nat.pred,
            unfold fin.succ,
        },
        rw [this x h_eq, count_map, fin_values_pred],
        -- TODO(solson): Extract as an independent theorem.
        show ∀ {a b : fin n}, fin.succ a = fin.succ b → a = b, by {
            intros a b,
            cases a,
            cases b,
            unfold fin.succ,
            assume h,
            congr,
            injection h with h',
            injection h',
        },
    }

instance fin.fintype (n : ℕ) : fintype (fin n) :=
    ⟨fin_values n, fin_values_pred n⟩

instance : fintype char := by { unfold char, apply_instance }

instance [decidable_eq α] [fintype α] : fintype (option α) := {
    values := none :: list.map some (fintype.values α),
    values_pred := λ x, match x with
    | none := by {
        rw [count, count_cond_true rfl, count_eq_zero],
        apply not_mem_map,
        intros,
        contradiction,
    }
    | some x := by {
        rw [count, count_cond_false, count_map, fintype.values_pred],
        show some x ≠ none, by contradiction,
        apply option.some.inj,
    }
    end
}

@[simp]
def cart : list α → list β → list (α × β)
| []      _  := []
| (x::xs) ys := list.map (prod.mk x) ys ++ cart xs ys

lemma count_cond_prod [decidable_eq α] [decidable_eq β] {a x : α} {b y : β} :
    count_cond (a, b) (x, y) = count_cond a x * count_cond b y :=
    have d1 : decidable (a = x), by apply_instance,
    have d2 : decidable (b = y), by apply_instance,
    match d1, d2 with
    | is_true h1, is_true h2 := by simp [h1, h2]
    | is_false h, _ := by {
        have h' : (a, b) ≠ (x, y),
        by { assume heq, injection heq, contradiction },
        simp [h, h'],
    }
    | _, is_false h := by {
        have h' : (a, b) ≠ (x, y),
        by { assume heq, injection heq, contradiction },
        simp [h, h'],
    }
    end

lemma count_cart_step [decidable_eq α] [decidable_eq β] {x : α} :
    ∀ {ys : list β},
    count (a, b) (list.map (prod.mk x) ys) = count_cond a x * count b ys
| []      := rfl
| (y::ys) := by {
    rw [list.map, count, count, count_cart_step, mul_add],
    congr,
    apply count_cond_prod,
}

theorem count_cart [decidable_eq α] [decidable_eq β] :
    ∀ {xs : list α} {ys : list β},
    count (a, b) (cart xs ys) = count a xs * count b ys
| []      ys := by simp [count]
| (x::xs) ys := by simp [count_append, count_cart, count, add_mul, count_cart_step]

instance [decidable_eq α] [decidable_eq β] [fintype α] [fintype β] :
    fintype (α × β) :=
{
    values := cart (fintype.values α) (fintype.values β),
    values_pred := λ ⟨a, b⟩, by simp [count_cart, fintype.values_pred],
}

instance {α β : Type u} [decidable_eq α] [decidable_eq β] [fintype α] [fintype β] :
    fintype (α ⊕ β) :=
{
    values :=
        list.map sum.inl (fintype.values α) ++
        list.map sum.inr (fintype.values β),
    values_pred := λ s, match s with
    | sum.inl a := by {
        rw [count_append, count_map, fintype.values_pred, count_eq_zero],
        show ∀ {a b : α}, sum.inl a = sum.inl b → a = b, by apply sum.inl.inj,
        apply not_mem_map,
        intros,
        contradiction,
    }
    | sum.inr b := by {
        rw [count_append, count_eq_zero, count_map, fintype.values_pred],
        show ∀ {a b : β}, sum.inr a = sum.inr b → a = b, by apply sum.inr.inj,
        apply not_mem_map,
        intros,
        contradiction,
    }
    end
}

def card (α : Type u) [decidable_eq α] [ft : fintype α] := ft.values.length

theorem card_fin : card (fin n) = n :=
begin
    rw card,
    unfold_projs,
    induction n,
    case nat.zero { refl },
    case nat.succ n' ih_n { simp [fin_values, ih_n] },
end

-- theorem card_prod [decidable_eq α] [decidable_eq β] [fintype α] [fintype β] :
--     card (α × β) = card α * card β :=
-- begin
--     unfold card,
--     unfold fintype.values,
--     unfold_projs,
-- end

#eval fintype.values (fin 2 × option unit ⊕ bool)
-- [(inl (0, none)), (inl (0, (some star))), (inl (1, none)), (inl (1, (some star))), (inr ff), (inr tt)]
#reduce fintype.values (bool ⊕ (bool × empty))
-- [sum.inl ff, sum.inl tt]

#reduce fintype.values (option (option empty))