import Mathlib

variable {m n : Type*} [Fintype m] [Fintype n] [DecidableEq m] [DecidableEq n]

/--
If the block matrix `fromBlocks A B C D : Matrix (m ⊕ n) (m ⊕ n) R` is orthogonal
and `I - A` is invertible, then the Schur-type expression
`D + C * (I - A)⁻¹ * B : Matrix n n R` is orthogonal.

Here “orthogonal” is expressed as membership in `Matrix.orthogonalGroup`.
-/
theorem orthogonal_of_block
    (A : Matrix m m ℝ) (B : Matrix m n ℝ)
    (C : Matrix n m ℝ) (D : Matrix n n ℝ)
    (hU : Matrix.fromBlocks A B C D ∈ Matrix.orthogonalGroup (m ⊕ n) ℝ)
    (hInv : IsUnit ((1 : Matrix m m ℝ) - A)) :
    D + C * ((1 : Matrix m m ℝ) - A)⁻¹ * B ∈ Matrix.orthogonalGroup n ℝ := by
  obtain ⟨ a, ha ⟩ := hInv.exists_left_inv;
  -- Using the fact that $U$ is orthogonal, we can derive the necessary equations for $A$, $B$, $C$, and $D$.
  have h_eq : A.transpose * A + C.transpose * C = 1 ∧ A.transpose * B + C.transpose * D = 0 ∧ B.transpose * A + D.transpose * C = 0 ∧ B.transpose * B + D.transpose * D = 1 := by
    have h_eq : (Matrix.fromBlocks A B C D).transpose * (Matrix.fromBlocks A B C D) = 1 := by
      convert hU.1 using 1;
    -- By definition of matrix multiplication, we can expand the product of the transpose of the block matrix and the original block matrix.
    have h_expand : (Matrix.fromBlocks A B C D).transpose * (Matrix.fromBlocks A B C D) = Matrix.fromBlocks (A.transpose * A + C.transpose * C) (A.transpose * B + C.transpose * D) (B.transpose * A + D.transpose * C) (B.transpose * B + D.transpose * D) := by
      ext i j; simp +decide [ Matrix.mul_apply, Matrix.transpose_apply, Matrix.fromBlocks ] ;
      rcases i with ( i | i ) <;> rcases j with ( j | j ) <;> simp +decide [ Matrix.mul_apply, Matrix.add_apply ];
    simp_all +decide [ ← Matrix.ext_iff, Fin.sum_univ_succ ];
    simp +decide [ Matrix.one_apply ];
  -- Now consider the expression $(D + C * (1 - A)⁻¹ * B)^T * (D + C * (1 - A)⁻¹ * B)$.
  have h_schur_simplified : (D + C * a * B).transpose * (D + C * a * B) = 1 := by
    -- Substitute the given equations into the expanded product.
    have h_subst : D.transpose * D + D.transpose * C * a * B + B.transpose * a.transpose * C.transpose * D + B.transpose * a.transpose * C.transpose * C * a * B = (1 - B.transpose * B) + (-B.transpose * A) * a * B + B.transpose * a.transpose * (-A.transpose * B) + B.transpose * a.transpose * (1 - A.transpose * A) * a * B := by
      simp_all +decide [ ← eq_sub_iff_add_eq', mul_assoc, Matrix.mul_sub, Matrix.sub_mul ];
      simp_all +decide [ Matrix.mul_assoc, ← Matrix.mul_sub ] ; abel_nf;
      simp +decide [ Matrix.add_mul, Matrix.mul_assoc ];
    -- Using the fact that $a * (1 - A) = 1$, we can simplify the expression further.
    have h_simplify : a.transpose * A.transpose = a.transpose - 1 := by
      apply_fun Matrix.transpose at ha;
      norm_num +zetaDelta at *;
      rw [ ← Matrix.mul_eq_one_comm ] at ha;
      rw [ ← ha, mul_sub, mul_one ];
      abel1;
    simp_all +decide [ Matrix.mul_assoc, sub_mul, mul_sub ];
    simp_all +decide [ Matrix.add_mul, Matrix.mul_add, ← Matrix.mul_assoc ];
    simp_all +decide [ mul_assoc, sub_mul, mul_sub ];
    -- Substitute $A * a = a - 1$ into the expression.
    have h_Aa : A * a = a - 1 := by
      rw [ ← Matrix.transpose_inj ] ; simp_all +decide [ Matrix.mul_assoc, Matrix.transpose_mul ];
    simp_all +decide [ ← Matrix.mul_assoc, ← eq_sub_iff_add_eq' ];
    rw [ show B.transpose * A * a * B = B.transpose * ( A * a ) * B by simp +decide only [ Matrix.mul_assoc ] ] ; rw [ h_Aa ] ; simp +decide [ Matrix.mul_assoc, Matrix.sub_mul, Matrix.mul_sub ] ; abel_nf;
  rw [ Matrix.inv_eq_left_inv ha ];
  exact Matrix.mem_unitaryGroup_iff'.mpr h_schur_simplified