-- CvxLean/Lib/Math/Subadditivity.lean

import Mathlib.LinearAlgebra.Matrix.PosDef

import Mathlib.LinearAlgebra.Matrix.Spectrum

import Mathlib.LinearAlgebra.Eigenspace.Basic

import Mathlib.LinearAlgebra.Matrix.LDL

import Mathlib.LinearAlgebra.Matrix.DotProduct

import CvxLean.Lib.Math.LinearAlgebra.Matrix.PosDef

import CvxLean.Lib.Math.LinearAlgebra.Matrix.Spectrum

import CvxLean.Lib.Math.LinearAlgebra.Eigenspace

namespace Finset

open BigOperators

end Finset

namespace Matrix

variable {n : Type _} [Fintype n] [DecidableEq n] [LinearOrder n] [LocallyFiniteOrderBot n]

open BigOperators Matrix

namespace IsHermitian

variable {𝕜 : Type _} [DecidableEq 𝕜] [RCLike 𝕜] {A : Matrix n n 𝕜} (hA : A.IsHermitian)

end IsHermitian

noncomputable def IsHermitian.sqrt {A : Matrix n n ℝ} (hA : A.IsHermitian) : Matrix n n ℝ :=
  hA.eigenvectorMatrix * Matrix.diagonal (fun i => (hA.eigenvalues i).sqrt) * hA.eigenvectorMatrixᵀ

-- Repository definitions used:
content: ...

-- Repository lemmas used:
content: theorem spectral_theorem (xs : OrthonormalBasis n 𝕜 (EuclideanSpace 𝕜 n)) (as : n → ℝ)
    (hxs : ∀ j, Module.End.HasEigenvector (Matrix.toLin' A) (as j) (xs j)) :
    xs.toBasis.toMatrix (Pi.basisFun 𝕜 n) * A =
    diagonal (RCLike.ofReal ∘ as) * xs.toBasis.toMatrix (Pi.basisFun 𝕜 n)

-- Local definitions used:
content: noncomputable def IsHermitian.sqrt {A : Matrix n n ℝ} (hA : A.IsHermitian) : Matrix n n ℝ :=
  hA.eigenvectorMatrix * Matrix.diagonal (fun i => (hA.eigenvalues i).sqrt) * hA.eigenvectorMatrixᵀ

-- Local lemmas used:
content: lemma eigenvectorMatrix_inv_mul : hA.eigenvectorMatrixInv * hA.eigenvectorMatrix = 1

content: theorem spectral_theorem'' :
    hA.eigenvectorMatrix * diagonal (RCLike.ofReal ∘ hA.eigenvalues) * hA.eigenvectorMatrixᴴ =
    A

content: lemma conjTranspose_eq_transpose {m n : Type _} {A : Matrix m n ℝ} : Aᴴ = Aᵀ

-- Target theorem and proof:
@[simp]
lemma PosSemidef.sqrt_mul_sqrt {A : Matrix n n ℝ} (hA : A.PosSemidef) :
    hA.1.sqrt * hA.1.sqrt = A :=
  calc
    hA.1.sqrt * hA.1.sqrt =
      hA.1.eigenvectorMatrix * (Matrix.diagonal (fun i => (hA.1.eigenvalues i).sqrt)
      * (hA.1.eigenvectorMatrixᵀ * hA.1.eigenvectorMatrix)
      * Matrix.diagonal (fun i => (hA.1.eigenvalues i).sqrt)) * hA.1.eigenvectorMatrixᵀ
:=
  calc
    hA.1.sqrt * hA.1.sqrt =
      hA.1.eigenvectorMatrix * (Matrix.diagonal (fun i => (hA.1.eigenvalues i).sqrt)
      * (hA.1.eigenvectorMatrixᵀ * hA.1.eigenvectorMatrix)
      * Matrix.diagonal (fun i => (hA.1.eigenvalues i).sqrt)) * hA.1.eigenvectorMatrixᵀ := by
      simp [IsHermitian.sqrt, Matrix.mul_assoc]
    _ = A := by
      rw [← conjTranspose_eq_transpose, hA.1.conjTranspose_eigenvectorMatrix,
        hA.1.eigenvectorMatrix_inv_mul, Matrix.mul_one, diagonal_mul_diagonal,
        ← hA.1.conjTranspose_eigenvectorMatrix]
      convert hA.1.spectral_theorem''
      rw [← Real.sqrt_mul (hA.eigenvalues_nonneg _), Real.sqrt_mul_self (hA.eigenvalues_nonneg _)]
      simp