import data.fintype.basic
import data.fintype.card
import data.finset.basic

import tactic.ring

open equiv fintype

section definitions

def is_derangement  {α : Type*} (f : perm α) : Prop := (∀ x : α, f x ≠ x)

def derangements (α : Type*) := {f : perm α // is_derangement f}

-- TODO fintype seems too restrictive; how do i correctly loosen this?
variables {α : Type*} [decidable_eq α] [fintype α]

instance derangement_decidable : decidable_pred (λ f : perm α, is_derangement f) :=
begin
  intro f,
  apply fintype.decidable_forall_fintype,
end

instance derangements_fintype : fintype (derangements α) := begin
  exact subtype.fintype _,
end

end definitions

section nice_lemmas
-- here's some functions that i may or may not use
-- TODO delete the ones i don't, and rename the rest to follow lean conventions

universe u
variables {α β : Type u}

def derangement_congr (e : α ≃ β) : (derangements α) ≃ (derangements β) :=
begin
  refine subtype_equiv (perm_congr e) _,
  intro σ,
  unfold is_derangement,
  rw ← not_iff_not,
  push_neg,
  split,
  { rintro ⟨a, ha⟩,
    use (e a),
    simp,
    assumption },
  { rintro ⟨b, hb⟩,
    use (e.symm b),
    conv_rhs { rw ← hb },
    simp },
end

def perm_subtype_preimage {p : α → Prop} [decidable_pred p] (f₀ : perm {x // p x}) :
{f : perm α // f ∘ coe = coe ∘ f₀} ≃ perm {x // ¬ p x} :=
begin
  -- we want to use equiv.set.compl:
  -- {e : α ≃ β // ∀ x : s, e x = e₀ x} ≃ ((sᶜ : set α) ≃ (tᶜ : set β))
  let s : set α := {a : α | p a},

  transitivity {f : perm α // ∀ x : s, f x = f₀ x},
  {
    apply subtype_equiv_right _,
    intro f,
    split,
    {
      rintros hf x,
      calc f x = (f ∘ coe) x  : rfl
      ...      = (coe ∘ f₀) x : by rw hf
      ...      = f₀ x         : rfl,
    },
    {
      intro hf,
      funext,
      simp,
      rw hf _,
    }
  },

  transitivity perm (sᶜ : set α),
  {
    exact equiv.set.compl f₀,
  },

  apply perm_congr,
  refl,
end

end nice_lemmas

section counting

-- TODO can i hide this in some way? it's computationally expensive...
def num_derangements (n : ℕ) : ℕ := fintype.card (derangements (fin n))

lemma num_derangements_invariant {α : Type} [fintype α] [decidable_eq α] :
fintype.card (derangements α) = num_derangements (fintype.card α) :=
begin
  unfold num_derangements,
  rw fintype.card_eq,
  obtain ⟨x⟩ := equiv_fin α, -- TODO what's going on here? it's related to non-computability i think
  use derangement_congr x,
end

lemma num_derangements_invariant_apply {α : Type} [fintype α] [decidable_eq α]
{n : ℕ} (hα : n = fintype.card α) :
fintype.card (derangements α) = num_derangements n :=
begin
  rw hα,
  exact num_derangements_invariant,
end

end counting

section type1

variables {α : Type*} [decidable_eq α]

def perm_ignore_two_cycle (a b : α) :
{ f : perm α // f a = b ∧ f b = a } ≃ perm {k : α // k ≠ a ∧ k ≠ b} :=
begin
  let pred : α → Prop := λ k, k = a ∨ k = b,
  let pred_a : pred a := by {left, refl},
  let pred_b : pred b := by {right, refl},

  let a_and_b := {k : α // pred k},
  let swap : perm a_and_b := equiv.swap ⟨a, pred_a⟩ ⟨b, pred_b⟩,

  transitivity {f : perm α // f ∘ coe = coe ∘ swap},
  {
    refine subtype_equiv_right _,
    intro f,
    split,
    {
      rintro ⟨hf_a, hf_b⟩,
      funext,
      cases x with x x_pred,
      simp,
      dsimp [swap],
      cases x_pred with x_eq_a x_eq_b,
      {
        -- why does subst work when rw doesn't??
        subst x_eq_a,
        simp,
        assumption,
      },
      {
        subst x_eq_b,
        simp,
        assumption,
      },
    },
    {
      intro hf,
      split,
      {
        let a' : a_and_b := ⟨a, pred_a⟩,
        calc f a = (f ∘ coe) a' : by simp
        ...      = coe (swap a') : by rw hf
        ...      = b            : by simp
      },
      {
        let b' : a_and_b := ⟨b, pred_b⟩,
        calc f b = (f ∘ coe) b' : by simp
        ...      = coe (swap b') : by rw hf
        ...      = a            : by simp
      }
    }
  },

  transitivity perm {k : α // ¬ pred k},
  {
    exact perm_subtype_preimage swap,
  },

  apply perm_congr,
  refine subtype_equiv_right _,
  intro x,
  dsimp [pred],
  push_neg,
end

-- TODO show that the perm_ignore_two_cycle preserves derangements
def derangement_ignore_two_cycle_helper (a b : α) (hab : a ≠ b) :
∀ f : { f : perm α // f a = b ∧ f b = a },
is_derangement f.val ↔ is_derangement (perm_ignore_two_cycle a b f) :=
begin
  intro f,

  have key : ∀ x : {x // x ≠ a ∧ x ≠ b}, f x = (perm_ignore_two_cycle a b f) x,
  {
    dsimp [perm_ignore_two_cycle, subtype_equiv_right, perm_subtype_preimage, equiv.set.compl],
    simp,
  },

  split,
  {
    intros h_derangement x,
    apply_fun subtype.val,
    simp,
    rw ← key x,
    exact h_derangement x,
  },
  {
    intros h_derangement x,
    by_cases H : x ≠ a ∧ x ≠ b,
    {
      -- TODO this is a silly juggle to avoid 'motive is not type correct'
      intro contra,
      specialize h_derangement ⟨x, H⟩,
      revert h_derangement,
      contrapose!,

      intro _,
      apply subtype.coe_injective,
      rw ← key _,
      exact contra,
    },
    {
      rw not_and_distrib at H,
      push_neg at H,
      cases H with Ha Hb,
      {
        rw [Ha, f.property.left],
        exact hab.symm,
      },
      {
        rw [Hb, f.property.right],
        exact hab,
      }
    }
  },
end

def derangement_ignore_two_cycle (a b : α) (hab : a ≠ b) :
{ f : derangements α // f.val a = b ∧ f.val b = a } ≃ derangements {k : α // k ≠ a ∧ k ≠ b} :=
begin
  -- TODO this is a mess. i gotta specify these predicates explicitly, which
  -- seems suboptimal... why can't it infer them?
  let foo : perm α → Prop := λ f, f a = b ∧ f b = a,
  let bar : {f : perm α // foo f} → Prop := λ f, is_derangement (perm_ignore_two_cycle a b f),

  transitivity {x // is_derangement x ∧ foo x},
  {
    exact subtype_subtype_equiv_subtype_inter _ foo,
  },

  -- TODO this seems... verbose
  -- but it's less verbose than the equivalent calc block would be, i think?
  apply equiv.trans _ (subtype_equiv_of_subtype (perm_ignore_two_cycle a b)),
  apply equiv.trans _ (subtype_equiv_right (derangement_ignore_two_cycle_helper a b hab)),
  apply equiv.trans _ (subtype_subtype_equiv_subtype_inter _ _).symm,

  refine subtype_equiv_right _,
  intro x,
  tauto,
end

end type1

section recursion

@[derive decidable_pred]
def type1 {n : ℕ} : derangements (fin(n+1)) → Prop := λ σ, σ.val (σ.val 0) = 0

@[derive decidable_pred]
def type2 {n : ℕ} : derangements (fin(n+1)) → Prop := λ σ, ¬type1 σ

def value_at_zero {n : ℕ} (σ : derangements (fin(n+1))) : fin(n+1) := σ.val 0


lemma size_of_type1' {n : ℕ} (m : fin(n+1)) :
fintype.card {σ : derangements (fin(n+2)) // type1 σ ∧ value_at_zero σ = m.succ} = num_derangements n :=
begin
  have h_0_ne_m := fin.succ_ne_zero m,

  have : fintype.card {k : fin(n+2) // k ≠ 0 ∧ k ≠ m.succ } = n,
  {
    sorry, -- TODO this is obvious to the human, but i don't know how to
    -- 'get into' the start of this proof. must be doable! but how? :\
  },

  -- TODO this is way too verbose. are there tricks to make this easier?
  rw ← num_derangements_invariant_apply this.symm,
  apply card_congr,
  apply equiv.trans _ (derangement_ignore_two_cycle 0 m.succ h_0_ne_m.symm),
  apply subtype_equiv_right,

  intro σ,
  split,
  {
    rintro ⟨h_split, h_m⟩,
    split,
    { exact h_m },
    { rw ← h_m, exact h_split },
  },
  {
    rintro ⟨h_0, h_m⟩,
    split,
    { unfold1 type1, rw h_0, exact h_m },
    { exact h_0 }
  }
end

lemma size_of_type1 {n : ℕ} :
fintype.card {σ : derangements (fin(n+2)) // type1 σ} = (n + 1) * num_derangements n :=
begin
  let lhs := {σ : derangements (fin(n+2)) // type1 σ},
  have : lhs ≃ Σ m : fin(n+2), {σ : lhs // value_at_zero σ.val = m}, from
  (sigma_preimage_equiv _).symm,
  rw card_congr this,
  rw fintype.card_sigma,
  rw fin.sum_univ_succ,

  have : {σ : lhs // value_at_zero σ.val = 0} ≃ empty,
  {
    apply equiv_empty,
    intro σ,
    have h_eq_0 : σ.val.val.val 0 = 0 := σ.property,
    have h_ne_0 : σ.val.val.val 0 ≠ 0 := σ.val.val.property 0,
    exact h_ne_0 h_eq_0,
  },
  rw card_congr this,

  have : ∀ m : fin(n+1), card {σ : lhs // value_at_zero σ.val = m.succ} = num_derangements n,
  {
    intro m,
    rw ← size_of_type1' m,
    apply card_congr,
    exact subtype_subtype_equiv_subtype_inter type1 (λ σ, value_at_zero σ = m.succ),
  },
  simp_rw this,
  simp,
end

lemma size_of_type2 {n : ℕ} :
fintype.card {σ : derangements (fin(n+2)) // type2 σ} = (n + 1) * num_derangements (n+1) :=
begin
  sorry,
end


lemma num_derangements_recursive (n : ℕ) :
num_derangements (n+2) = (n + 1) * (num_derangements n + num_derangements (n+1)) :=
begin
  -- TODO would it be better to get subsets of perm (fin(_)) instead?
  let splitting_pred : derangements (fin(n+2)) → Prop := λ σ, σ.val (σ.val 0) = 0,
  let type1_derangements := {σ : derangements (fin(n+2)) // type1 σ},
  let type2_derangements := {σ : derangements (fin(n+2)) // type2 σ},

  calc num_derangements (n+2) = fintype.card ( derangements (fin(n+2)) ) : rfl
  ...                         = fintype.card ( type1_derangements ⊕ type2_derangements )         : card_congr (sum_compl splitting_pred).symm
  ...                         = fintype.card type1_derangements + fintype.card type2_derangements : fintype.card_sum _ _
  ...                         = (n+1) * num_derangements n + (n+1) * num_derangements (n+1)     : by rw [size_of_type1, size_of_type2]
  ...                         = (n+1) * (num_derangements n + num_derangements (n+1))           : by ring,
  -- done!
end

end recursion

section explicit

-- TODO theorem about alternating sum of stuff

end explicit

#eval num_derangements 0
#eval num_derangements 1
#eval num_derangements 2
#eval num_derangements 3
#eval num_derangements 4
#eval num_derangements 5