import RNLAPhaseB.B3Base
import RNLAPhaseB.Complexity

open MeasureTheory
open scoped Matrix
open scoped Matrix.Norms.L2Operator
open scoped BigOperators
open scoped ENNReal

namespace RNLAPhaseB.MT2020.LaplacianSolver

open RNLAPhaseB.Complexity

/-- Weighted, loop-free, undirected graph on `V = {1,…,n}` given by a nonnegative weight
function `w : V × V → ℝ₊` with `w i i = 0` and `w i j = w j i` (book.md:1185). -/
def IsWeightedLoopFreeUndirectedGraph {n : ℕ} (w : Fin n → Fin n → ℝ) : Prop :=
  (∀ i j, 0 ≤ w i j) ∧ (∀ i, w i i = 0) ∧ (∀ i j, w i j = w j i)

/-- Elementary Laplacian `Δ_uv = (δ_u - δ_v)(δ_u - δ_v)^*` (book.md:3327-3329). -/
noncomputable def elementaryLaplacian {n : ℕ} (u v : Fin n) : Matrix (Fin n) (Fin n) ℝ :=
  Matrix.vecMulVec (Pi.single u (1 : ℝ) - Pi.single v 1) (Pi.single u (1 : ℝ) - Pi.single v 1)

/-- Graph Laplacian `L = ∑_{e = uv} w(e) Δ_e`, each unordered edge counted once. -/
noncomputable def graphLaplacian {n : ℕ} (w : Fin n → Fin n → ℝ) : Matrix (Fin n) (Fin n) ℝ :=
  ∑ u : Fin n, ∑ v : Fin n, if u < v then w u v • elementaryLaplacian u v else 0

/-- The sparsity `m`: number of strictly positive weights `w_ij` with `i ≤ j`. -/
noncomputable def graphSparsity {n : ℕ} (w : Fin n → Fin n → ℝ) : ℕ := by
  classical
  exact (Finset.univ.filter (fun p : Fin n × Fin n => p.1 ≤ p.2 ∧ 0 < w p.1 p.2)).card

/-- Semidefinite order `A ≼ B`. -/
def PsdLE {n : ℕ} (A B : Matrix (Fin n) (Fin n) ℝ) : Prop := (B - A).PosSemidef

/-- Morally lower-triangular: a (known) permutation of the rows brings `C` into
lower-triangular form. -/
def MorallyLowerTriangular {n : ℕ} (C : Matrix (Fin n) (Fin n) ℝ) : Prop :=
  ∃ σ : Equiv.Perm (Fin n), ∀ i j : Fin n, i < j → C (σ i) j = 0

/-- Moore–Penrose relations: `X` is the pseudoinverse of `M`. -/
def IsMoorePenroseInverse {n : ℕ} (M X : Matrix (Fin n) (Fin n) ℝ) : Prop :=
  M * X * M = M ∧ X * M * X = X ∧ (M * X)ᵀ = M * X ∧ (X * M)ᵀ = X * M

/-- The seminorm `‖v‖_L = (v^* L v)^{1/2}`. -/
noncomputable def lSeminorm {n : ℕ} (L : Matrix (Fin n) (Fin n) ℝ) (v : Fin n → ℝ) : ℝ :=
  Real.sqrt (v ⬝ᵥ (L *ᵥ v))

/-- State of preconditioned conjugate gradient: iterate, residual, preconditioned
residual, search direction. -/
structure PCGState (n : ℕ) where
  x : Fin n → ℝ
  r : Fin n → ℝ
  z : Fin n → ℝ
  p : Fin n → ℝ

/-- PCG initialization with zero initial guess; `X` applies the (pseudo)inverse of the
preconditioner `CC^*`. -/
noncomputable def pcgInit {n : ℕ} (X : Matrix (Fin n) (Fin n) ℝ) (f : Fin n → ℝ) : PCGState n :=
  ⟨0, f, X *ᵥ f, X *ᵥ f⟩

/-- One standard (exact-arithmetic) PCG iteration for `L x = f` with preconditioner
inverse `X`. -/
noncomputable def pcgStep {n : ℕ} (L X : Matrix (Fin n) (Fin n) ℝ) (s : PCGState n) :
    PCGState n :=
  let α := (s.r ⬝ᵥ s.z) / (s.p ⬝ᵥ (L *ᵥ s.p))
  let x' := s.x + α • s.p
  let r' := s.r - α • (L *ᵥ s.p)
  let z' := X *ᵥ r'
  let β := (r' ⬝ᵥ z') / (s.r ⬝ᵥ s.z)
  ⟨x', r', z', z' + β • s.p⟩

/-- The state after `k` PCG iterations. -/
noncomputable def pcgIterate {n : ℕ} (L X : Matrix (Fin n) (Fin n) ℝ) (f : Fin n → ℝ)
    (k : ℕ) : PCGState n :=
  (pcgStep L X)^[k] (pcgInit X f)

/-- Encoding of a PCG state as the input vector of an arithmetic program. -/
def PCGState.toInput {n : ℕ} (s : PCGState n) : Fin (n + n + n + n) → ℝ :=
  Fin.append (Fin.append (Fin.append s.x s.r) s.z) s.p

/-- Encoding of a PCG state as an output list. -/
def PCGState.toList {n : ℕ} (s : PCGState n) : List ℝ :=
  List.ofFn s.x ++ List.ofFn s.r ++ List.ofFn s.z ++ List.ofFn s.p

/-- The arithmetic program performs the PCG iteration from state `s` exactly, using at
most `bound` arithmetic operations (+, -, *, / each cost one) in one successful run. -/
def StepComputedWithin {n : ℕ} (prog : Arithmetic.Program (n + n + n + n))
    (L X : Matrix (Fin n) (Fin n) ℝ) (s : PCGState n) (bound : ℝ) : Prop :=
  ∃ output cost, prog.run s.toInput = some (output, cost) ∧
    output = (pcgStep L X s).toList ∧ (cost : ℝ) ≤ bound

/-- mt2020 cor:laplacian-solver. If `C` satisfies eqn:approx-cholesky
(`0.5 L ≼ CC^* ≼ 1.5 L`, `C` morally lower-triangular, `nnz(C) = O(m log² n)`), then
every consistent system `L x = f` is solved to relative error `ε` in `‖·‖_L` by at most
`1 + log₃(1/ε)` PCG iterations (preconditioner `CC^*`, applied via its pseudoinverse,
zero initial guess), each costing `O(m log² n)` arithmetic operations. The Big-O
hypothesis on `nnz(C)` is a family hypothesis with constant `c₀` and threshold `N₀`;
the per-iteration Big-O constant `K` and threshold `N` may depend on them only. -/
theorem laplacian_pcg_solver :
    ∀ c₀ : ℝ, ∀ N₀ : ℕ, ∃ K : ℝ, 0 < K ∧ ∃ N : ℕ,
      ∀ (n : ℕ) (w : Fin n → Fin n → ℝ) (C : Matrix (Fin n) (Fin n) ℝ),
        IsWeightedLoopFreeUndirectedGraph w →
        MorallyLowerTriangular C →
        PsdLE ((1 / 2 : ℝ) • graphLaplacian w) (C * Cᵀ) →
        PsdLE (C * Cᵀ) ((3 / 2 : ℝ) • graphLaplacian w) →
        (N₀ ≤ n → (matrixNnz C : ℝ) ≤ c₀ * ((graphSparsity w : ℝ) * Real.log n ^ 2)) →
        ∀ X : Matrix (Fin n) (Fin n) ℝ, IsMoorePenroseInverse (C * Cᵀ) X →
        ∃ prog : Arithmetic.Program (n + n + n + n),
          ∀ f : Fin n → ℝ, (∃ x, graphLaplacian w *ᵥ x = f) →
          ∀ ε : ℝ, 0 < ε → ε < 1 →
            ∃ T : ℕ, (T : ℝ) ≤ 1 + Real.logb 3 (1 / ε) ∧
              (∀ x, graphLaplacian w *ᵥ x = f →
                lSeminorm (graphLaplacian w) (x - (pcgIterate (graphLaplacian w) X f T).x)
                  ≤ ε * lSeminorm (graphLaplacian w) x) ∧
              (N ≤ n → ∀ k, k < T →
                StepComputedWithin prog (graphLaplacian w) X
                  (pcgIterate (graphLaplacian w) X f k)
                  (K * ((graphSparsity w : ℝ) * Real.log n ^ 2))) := by
  sorry

end RNLAPhaseB.MT2020.LaplacianSolver