Mix p Ensemble pure iff pure

theorem mix_pEnsemble_pure_iff_pure {e : PEnsemble d α} :    mix (toMEnsemble e) = MState.pure ψ ↔    ∀ i : α, e.distr i ≠ 0 → MState.pure (e.states i) = MState.pure ψ := by  refine ⟨fun h i hi ↦ ?_, fun h ↦ ?_⟩  · apply MState.eq_of_sum_eq_pure ?_ ?_ ?_ e.distr.normalized i (Finset.mem_univ i)    · exact_mod_cast lt_of_le_of_ne (e.distr i).zero_le hi.symm    · exact (MState.pure_iff_purity_one _).mp ⟨ψ, rfl⟩    · exact congr_arg MState.M h.symm    · grind  · have h_sum : mix (toMEnsemble e) = ∑ i, (e.distr i).val • (MState.pure ψ).M := by      refine Finset.sum_congr rfl fun i _ => ?_      by_cases hi : e.distr i = 0      · simp [hi]      · rw [← h i hi]        rfl    simp [MState.ext_iff, h_sum, ← Finset.sum_smul]