import PDETextbook.Base

open MeasureTheory
open scoped BigOperators ENNReal

noncomputable section

namespace PDETextbook.LarssonThomee.Thm13_2

/-- A triangle given by its three vertices in R^2. -/
abbrev Tri := Fin 3 → Euclidean 2

/-- The closed triangle K (convex hull of its vertices). -/
def triSet (K : Tri) : Set (Euclidean 2) := convexHull ℝ (Set.range K)

/-- Omega subset R^2 is a bounded convex domain whose boundary is a polygon
(the closure is the convex hull of finitely many points). -/
def IsConvexPolygonalDomain (Ω : Set (Euclidean 2)) : Prop :=
  IsOpen Ω ∧ Ω.Nonempty ∧ Bornology.IsBounded Ω ∧ Convex ℝ Ω ∧
  ∃ P : Finset (Euclidean 2), closure Ω = convexHull ℝ (P : Set (Euclidean 2))

/-- A triangulation of Omega (Sect. 5.2): closed nondegenerate triangles covering the
closure of Omega; two distinct triangles meet in the empty set, a common node or a common
edge; no node lies in the interior of an edge. -/
structure Triangulation (Ω : Set (Euclidean 2)) where
  tris : Finset Tri
  nondeg : ∀ K ∈ tris, AffineIndependent ℝ K
  cover : closure Ω = ⋃ K ∈ tris, triSet K
  conforming : ∀ K ∈ tris, ∀ K' ∈ tris, triSet K ≠ triSet K' →
    triSet K ∩ triSet K' = ∅ ∨
    (∃ i j, K i = K' j ∧ triSet K ∩ triSet K' = {K i}) ∨
    (∃ i i' j j', i ≠ i' ∧ K i = K' j ∧ K i' = K' j' ∧
      triSet K ∩ triSet K' = segment ℝ (K i) (K i'))
  noHangingNode : ∀ K ∈ tris, ∀ K' ∈ tris, ∀ (i a b : Fin 3), a ≠ b →
    K i ∉ openSegment ℝ (K' a) (K' b)

/-- h = max_K diam(K). -/
def meshSize {Ω : Set (Euclidean 2)} (T : Triangulation Ω) : ℝ :=
  ⨆ K ∈ T.tris, Metric.diam (triSet K)

/-- All angles of all triangles are bounded below by alpha (uniform in h). -/
def MinAngleBound {Ω : Set (Euclidean 2)} (T : Triangulation Ω) (α : ℝ) : Prop :=
  ∀ K ∈ T.tris, ∀ i j l : Fin 3, i ≠ j → j ≠ l → i ≠ l →
    α ≤ EuclideanGeometry.angle (K i) (K j) (K l)

/-- S_h subset H_0^1: continuous piecewise linear functions on T_h vanishing on Gamma. -/
def Sh (Ω : Set (Euclidean 2)) (T : Triangulation Ω) : Set (Euclidean 2 → ℝ) :=
  { χ | ContinuousOn χ (closure Ω) ∧
        (∀ K ∈ T.tris, ∃ g : Euclidean 2 →ᵃ[ℝ] ℝ, ∀ x ∈ triSet K, χ x = g x) ∧
        (∀ x ∈ frontier Ω, χ x = 0) ∧ InH10Closure Ω χ }

open Classical in
/-- A weak gradient (unique a.e. on the open set Omega); zero if v is not in H^1. -/
def weakGrad (Ω : Set (Euclidean 2)) (v : Euclidean 2 → ℝ) : Fin 2 → Euclidean 2 → ℝ :=
  if h : InH1 Ω v then h.choose else 0

/-- Energy form a(v,w) = ∫_Ω ∇v·∇w. -/
def energyForm (Ω : Set (Euclidean 2)) (v w : Euclidean 2 → ℝ) : ℝ :=
  ∫ x in Ω, ∑ i, weakGrad Ω v i x * weakGrad Ω w i x

/-- L2 scalar product (v,w) = ∫_Ω v w (real-valued). -/
def ip (Ω : Set (Euclidean 2)) (v w : Euclidean 2 → ℝ) : ℝ := ∫ x in Ω, v x * w x

/-- L2(Omega) norm. -/
def l2norm (Ω : Set (Euclidean 2)) (v : Euclidean 2 → ℝ) : ℝ :=
  (eLpNorm v 2 (domainMeasure Ω)).toReal

/-- H^1 seminorm |v|_1 = (Σ_i ‖∂_i v‖^2)^{1/2}. -/
def seminorm1 (Ω : Set (Euclidean 2)) (v : Euclidean 2 → ℝ) : ℝ :=
  Real.sqrt (∑ i, (eLpNorm (weakGrad Ω v i) 2 (domainMeasure Ω)).toReal ^ 2)

/-- Ritz projection (5.49): R ∈ S_h and a(R - v, χ) = 0 for all χ ∈ S_h. -/
def IsRitzProjection (Ω : Set (Euclidean 2)) (T : Triangulation Ω)
    (v R : Euclidean 2 → ℝ) : Prop :=
  R ∈ Sh Ω T ∧ ∀ χ ∈ Sh Ω T, energyForm Ω (fun x => R x - v x) χ = 0

/-- Time levels t_n = n k. -/
def tn (k : ℝ) (n : ℕ) : ℝ := (n : ℝ) * k

/-- Scheme (13.5) with t_n = n k, for n ≥ 1, U^n ∈ S_h. -/
def SolvesScheme (Ω : Set (Euclidean 2)) (T : Triangulation Ω) (k : ℝ)
    (f : ℝ → Euclidean 2 → ℝ) (U : ℕ → Euclidean 2 → ℝ) : Prop :=
  (∀ n, U n ∈ Sh Ω T) ∧
  ∀ n : ℕ, 1 ≤ n → ∀ χ ∈ Sh Ω T,
    ip Ω (fun x => (U (n + 1) x - 2 * U n x + U (n - 1) x) / k ^ 2) χ
      + energyForm Ω (fun x => (1 / 4 : ℝ) * (U (n + 1) x + 2 * U n x + U (n - 1) x)) χ
      = ip Ω (f (tn k n)) χ

def timeDeriv (u : ℝ → Euclidean 2 → ℝ) (t : ℝ) (x : Euclidean 2) : ℝ :=
  deriv (fun s => u s x) t

def spaceLaplacian (u : ℝ → Euclidean 2 → ℝ) (t : ℝ) (x : Euclidean 2) : ℝ :=
  ∑ i, partialDeriv (fun y => partialDeriv (u t) i y) i x

/-- u solves (13.1): u_tt - Δu = f in Ω×R_+, u = 0 on Γ×R_+, u(0)=v, u_t(0)=w. -/
def IsWaveSolution (Ω : Set (Euclidean 2)) (f : ℝ → Euclidean 2 → ℝ)
    (v w : Euclidean 2 → ℝ) (u : ℝ → Euclidean 2 → ℝ) : Prop :=
  (∀ t : ℝ, 0 < t → ∀ x ∈ Ω,
      timeDeriv (timeDeriv u) t x - spaceLaplacian u t x = f t x) ∧
  (∀ t : ℝ, 0 < t → ∀ x ∈ frontier Ω, u t x = 0) ∧
  (∀ x ∈ Ω, u 0 x = v x ∧ timeDeriv u 0 x = w x)

/-- Theorem 13.2. "Appropriate regularity conditions for u" is realized as u smooth
(C^m for every m) jointly in (t,x). θ^n = U^n - R_h u(t_n), ∂_t θ^0 = (θ^1-θ^0)/k,
U^{n+1/2} = (U^n+U^{n+1})/2, ∂_t U^n = (U^{n+1}-U^n)/k, and errors are measured at the
time t_n + (1/2) k. -/
theorem thm_13_2 (Ω : Set (Euclidean 2)) (hΩ : IsConvexPolygonalDomain Ω)
    (α : ℝ) (hα : 0 < α)
    (f : ℝ → Euclidean 2 → ℝ) (v w : Euclidean 2 → ℝ) (u : ℝ → Euclidean 2 → ℝ)
    (hsmooth : ∀ m : ℕ, ContDiff ℝ m (fun p : ℝ × Euclidean 2 => u p.1 p.2))
    (hu : IsWaveSolution Ω f v w u) (C₀ : ℝ) :
    ∃ Cu : ℝ → ℝ, Monotone Cu ∧
      ∀ T : Triangulation Ω, MinAngleBound T α → meshSize T < 1 →
      ∀ k : ℝ, 0 < k →
      ∀ U : ℕ → Euclidean 2 → ℝ, SolvesScheme Ω T k f U →
      ∀ R₀ R₁ : Euclidean 2 → ℝ,
        IsRitzProjection Ω T (u (tn k 0)) R₀ → IsRitzProjection Ω T (u (tn k 1)) R₁ →
        l2norm Ω (fun x => ((U 1 x - R₁ x) - (U 0 x - R₀ x)) / k)
          + seminorm1 Ω (fun x => U 0 x - R₀ x) + seminorm1 Ω (fun x => U 1 x - R₁ x)
          ≤ C₀ * (meshSize T ^ 2 + k ^ 2) →
        ∀ n : ℕ,
          (l2norm Ω (fun x => (U n x + U (n + 1) x) / 2 - u (tn k n + (1 / 2 : ℝ) * k) x)
            + l2norm Ω (fun x => (U (n + 1) x - U n x) / k
                - timeDeriv u (tn k n + (1 / 2 : ℝ) * k) x)
            ≤ Cu (tn k n) * (meshSize T ^ 2 + k ^ 2)) ∧
          seminorm1 Ω (fun x => (U n x + U (n + 1) x) / 2 - u (tn k n + (1 / 2 : ℝ) * k) x)
            ≤ Cu (tn k n) * (meshSize T + k ^ 2) := by
  sorry

end PDETextbook.LarssonThomee.Thm13_2