import PDETextbook.Base

open MeasureTheory
open scoped BigOperators ENNReal

noncomputable section

namespace PDETextbook.LarssonThomee.Thm43

/-- Mesh point `x_j = j h`, `j ∈ ℤ²`. -/
def meshPt (h : ℝ) (j : ℤ × ℤ) : Euclidean 2 :=
  EuclideanSpace.single 0 ((j.1 : ℝ) * h) + EuclideanSpace.single 1 ((j.2 : ℝ) * h)

/-- Index shifts `e₁ = (1,0)`, `e₂ = (0,1)`. -/
def ev (l : Fin 2) : ℤ × ℤ := if l = 0 then (1, 0) else (0, 1)

def IsMeshNeighbor (j k : ℤ × ℤ) : Prop :=
  ∃ l : Fin 2, k = j + ev l ∨ k = j - ev l

/-- Smooth boundary: locally, after a rotation, `Ω` lies below a `C^∞` graph. -/
def HasSmoothBoundary (Ω : Set (Euclidean 2)) : Prop :=
  ∀ p ∈ frontier Ω, ∃ (R : Euclidean 2 ≃ₗᵢ[ℝ] Euclidean 2) (g : ℝ → ℝ) (r : ℝ),
    0 < r ∧ ContDiff ℝ (⊤ : ℕ∞) g ∧
    ∀ x ∈ Metric.ball p r, (x ∈ Ω ↔ (R x) 1 < g ((R x) 0))

/-- `Ω_h`: interior mesh points whose four neighbours lie in `Ω̄`. -/
def OmegaH (Ω : Set (Euclidean 2)) (h : ℝ) : Set (ℤ × ℤ) :=
  {j | meshPt h j ∈ Ω ∧ ∀ k, IsMeshNeighbor j k → meshPt h k ∈ closure Ω}

/-- `ω_h`: mesh points in `Ω` not in `Ω_h`. -/
def omegaH (Ω : Set (Euclidean 2)) (h : ℝ) : Set (ℤ × ℤ) :=
  {j | meshPt h j ∈ Ω ∧ j ∉ OmegaH Ω h}

/-- Selection preceding (4.14): neighbour `x_i = x_{nb j} ∈ Ω_h ∪ ω_h` and `x̄_j ∈ Γ`
on the axis line through `x_j, x_i`, with `x_j` between `x_i` and `x̄_j`, not a mesh point. -/
def IsAdmissibleSelection (Ω : Set (Euclidean 2)) (h : ℝ)
    (nb : ℤ × ℤ → ℤ × ℤ) (xbar : ℤ × ℤ → Euclidean 2) : Prop :=
  ∀ j ∈ omegaH Ω h,
    IsMeshNeighbor j (nb j) ∧ nb j ∈ OmegaH Ω h ∪ omegaH Ω h ∧
    xbar j ∈ frontier Ω ∧
    (∃ t : ℝ, 0 < t ∧ t < 1 ∧
      xbar j = meshPt h j + t • (meshPt h j - meshPt h (nb j))) ∧
    ∀ k, xbar j ≠ meshPt h k

/-- `α_j` of (4.14). -/
def alphaCoef (h : ℝ) (xj xb : Euclidean 2) : ℝ :=
  dist xj xb / (h + dist xj xb)

/-- `ℓ_h U_j` of (4.14). -/
def ellH (h : ℝ) (U : Euclidean 2 → ℝ) (nb : ℤ × ℤ → ℤ × ℤ)
    (xbar : ℤ × ℤ → Euclidean 2) (j : ℤ × ℤ) : ℝ :=
  U (meshPt h j) - alphaCoef h (meshPt h j) (xbar j) * U (meshPt h (nb j))
    - (1 - alphaCoef h (meshPt h j) (xbar j)) * U (xbar j)

/-- `Δ_h U_j` of (4.11). -/
def discLap (h : ℝ) (U : Euclidean 2 → ℝ) (j : ℤ × ℤ) : ℝ :=
  ∑ l : Fin 2, (U (meshPt h (j + ev l)) - 2 * U (meshPt h j) + U (meshPt h (j - ev l))) / h ^ 2

/-- `Γ_h`: points of `Γ` that neighbour a point of `Ω_h`, or are some `x̄_j`, `x_j ∈ ω_h`. -/
def GammaH (Ω : Set (Euclidean 2)) (h : ℝ) (xbar : ℤ × ℤ → Euclidean 2) : Set (Euclidean 2) :=
  {p | p ∈ frontier Ω ∧ ((∃ j ∈ OmegaH Ω h, ∃ k, IsMeshNeighbor j k ∧ p = meshPt h k) ∨
    (∃ j ∈ omegaH Ω h, p = xbar j))}

/-- Problem (4.15). -/
def SolvesScheme415 (Ω : Set (Euclidean 2)) (h : ℝ) (nb : ℤ × ℤ → ℤ × ℤ)
    (xbar : ℤ × ℤ → Euclidean 2) (f U : Euclidean 2 → ℝ) : Prop :=
  (∀ j ∈ OmegaH Ω h, -discLap h U j = f (meshPt h j)) ∧
  (∀ j ∈ omegaH Ω h, ellH h U nb xbar j = 0) ∧
  (∀ p ∈ GammaH Ω h xbar, U p = 0)

/-- Classical Laplacian `∂₁²u + ∂₂²u`. -/
def lap (u : Euclidean 2 → ℝ) (x : Euclidean 2) : ℝ :=
  ∑ l : Fin 2, iteratedFDeriv ℝ 2 u x (fun _ => EuclideanSpace.single l 1)

/-- Problem (4.9): `-Δu = f` in `Ω`, `u = 0` on `Γ`. -/
def SolvesPoisson49 (Ω : Set (Euclidean 2)) (f u : Euclidean 2 → ℝ) : Prop :=
  (∀ x ∈ Ω, -lap u x = f x) ∧ (∀ x ∈ frontier Ω, u x = 0)

/-- `D^α v` with `α = (a,b)`, derivatives taken within `S`. -/
def Dalpha (S : Set (Euclidean 2)) (v : Euclidean 2 → ℝ) (a b : ℕ) (x : Euclidean 2) : ℝ :=
  iteratedFDerivWithin ℝ (a + b) v S x
    (fun i : Fin (a + b) => if (i : ℕ) < a then EuclideanSpace.single 0 1
      else EuclideanSpace.single 1 1)

/-- `‖v‖_{C^4(Ω̄)} = max_{|α| ≤ 4} sup_{Ω̄} |D^α v|`. -/
def C4norm (Ω : Set (Euclidean 2)) (v : Euclidean 2 → ℝ) : ℝ :=
  ⨆ (p : ℕ × ℕ) (_ : p.1 + p.2 ≤ 4), ⨆ x ∈ closure Ω, |Dalpha (closure Ω) v p.1 p.2 x|

/-- Theorem 4.3: `|U - u|_{Ω_h ∪ ω_h} ≤ C h² ‖u‖_{C⁴}`, with `|·|_S` the max over `S` (4.7). -/
theorem thm_4_3 (Ω : Set (Euclidean 2)) (hΩo : IsOpen Ω) (hΩc : IsConnected Ω)
    (hΩb : Bornology.IsBounded Ω) (hΩcv : Convex ℝ Ω) (hΓ : HasSmoothBoundary Ω) :
    ∃ C : ℝ, 0 < C ∧ ∀ (M : ℕ), 0 < M →
      ∀ (nb : ℤ × ℤ → ℤ × ℤ) (xbar : ℤ × ℤ → Euclidean 2),
        IsAdmissibleSelection Ω (1 / (M : ℝ)) nb xbar →
      ∀ (f u U : Euclidean 2 → ℝ),
        ContDiffOn ℝ 4 u (closure Ω) →
        SolvesPoisson49 Ω f u →
        SolvesScheme415 Ω (1 / (M : ℝ)) nb xbar f U →
        ∀ j ∈ OmegaH Ω (1 / (M : ℝ)) ∪ omegaH Ω (1 / (M : ℝ)),
          |U (meshPt (1 / (M : ℝ)) j) - u (meshPt (1 / (M : ℝ)) j)|
            ≤ C * (1 / (M : ℝ)) ^ 2 * C4norm Ω u := by
  sorry

end PDETextbook.LarssonThomee.Thm43