import SciLean.Core.Distribution.Basic
import SciLean.Core.FunctionTransformations
import SciLean.Core.FunctionPropositions
import SciLean.Core.Notation

open MeasureTheory

namespace SciLean

open Distribution

variable
  {R} [RealScalar R]
  {W} [Vec R W]
  {X} [Vec R X] [MeasureSpace X]
  {Y} [Vec R Y] [Module ℝ Y]
  {Z} [Vec R Z] [Module ℝ Z]
  {U} [Vec R U] -- [Module ℝ U]
  {V} [Vec R V] -- [Module ℝ U]

set_default_scalar R

noncomputable
def diracDeriv (x dx : X) : 𝒟' X := fun φ ⊸ cderiv R φ x dx

@[fun_prop]
def DistribDifferentiableAt (f : X → 𝒟'(Y,Z)) (x : X) :=
  ∀ (φ : X → 𝒟 Y), CDifferentiableAt R φ x → CDifferentiableAt R (fun x => f x (φ x)) x

theorem distribDifferentiableAt_const_test_fun
    {f : X → 𝒟'(Y,Z)} {x : X}
    (hf : DistribDifferentiableAt f x)
    {φ : 𝒟 Y} :
    CDifferentiableAt R (fun x => f x φ) x := by
  apply hf
  fun_prop

@[fun_prop]
def DistribDifferentiable (f : X → 𝒟'(Y,Z)) :=
  ∀ x, DistribDifferentiableAt f x

-- TODO:
-- probably change the definition of `parDistribDeriv` to:
-- ⟨⟨fun φ =>
--    if h : DistribDifferentiableAt f x then
--      ∂ (x':=x;dx), ⟪f x', φ⟫
--    else
--      0 , sorry_proof⟩⟩
-- I believe in that case the function is indeed linear in φ

open Classical in
@[fun_trans]
noncomputable
def parDistribDeriv (f : X → 𝒟'(Y,Z)) (x dx : X) : 𝒟'(Y,Z) :=
  ⟨fun φ => ∂ (x':=x;dx), f x' φ, sorry_proof⟩

@[simp, ftrans_simp]
theorem action_parDistribDeriv (f : X → 𝒟'(Y,Z)) (x dx : X) (φ : 𝒟 Y) :
    parDistribDeriv f x dx φ = ∂ (x':=x;dx), f x' φ := rfl

----------------------------------------------------------------------------------------------------
-- Const rule --------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem DistribDiffrentiable.const_rule (T : 𝒟'(X,Y)) :
    DistribDifferentiable (fun _ : W => T) := by
  intro _ φ hφ; simp; fun_prop

@[fun_trans]
theorem parDistribDeriv.const_rule (T : 𝒟'(X,Y)) :
    parDistribDeriv (fun _ : W => T)
    =
    fun w dw =>
      0 := by
  funext w dw; ext φ
  unfold parDistribDeriv
  fun_trans

----------------------------------------------------------------------------------------------------
-- Pure --------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem dirac.arg_xy.DistribDiffrentiable_rule
    (x : W → X) (hx : CDifferentiable R x) :
    DistribDifferentiable (R:=R) (fun w => dirac (x w))  := by
  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  simp [action_dirac, dirac]
  fun_prop

@[fun_trans]
theorem dirac.arg_x.parDistribDeriv_rule
    (x : W → X) (hx : CDifferentiable R x) :
    parDistribDeriv (R:=R) (fun w => dirac (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      diracDeriv xdx.1 xdx.2 := by --= (dpure (R:=R) ydy.1 ydy.2) := by
  funext w dw; ext φ
  unfold parDistribDeriv dirac diracDeriv
  simp [pure, fwdDeriv, DistribDifferentiableAt]
  fun_trans

----------------------------------------------------------------------------------------------------
-- Composition -------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem DistribDiffrentiable.comp_rule
    (f : Y → 𝒟'(Z,U)) (g : X → Y)
    (hf : DistribDifferentiable f) (hg : CDifferentiable R g) :
    DistribDifferentiable (fun x => f (g x)) := by
  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  apply CDifferentiable.comp_rule (K:=R) (f:=fun xy : X×Y => f xy.2 (φ xy.1)) (g:=fun x => (x, g x))
    (hg:=by fun_prop)
  intro x
  sorry_proof -- is this even true ?

@[fun_trans]
theorem parDistribDeriv.comp_rule
    (f : Y → 𝒟'(Z,U)) (g : X → Y)
    (hf : DistribDifferentiable f) (hg : CDifferentiable R g) :
    parDistribDeriv (fun x => f (g x))
    =
    fun x dx =>
      let ydy := fwdDeriv R g x dx
      parDistribDeriv f ydy.1 ydy.2 := by

  funext x dx; ext φ
  unfold parDistribDeriv
  simp[hg]
  sorry_proof

----------------------------------------------------------------------------------------------------
-- Bind --------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- The assumptions here are definitely not right!!!
-- I think `f` has to be `deg`
@[fun_prop]
theorem Bind.bind.arg_fx.DistribDifferentiable_rule
    (f : X → Y → 𝒟'(Z,V)) (g : X → 𝒟'(Y,U)) (L : U ⊸ V ⊸ W)
    (hf : DistribDifferentiable (fun (x,y) => f x y)) -- `f` has to be nice enough to accomodate action of `g`
    (hg : DistribDifferentiable g) :
    DistribDifferentiable (fun x => (g x).bind (f x) L) := by

  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  simp
  sorry_proof

@[fun_trans]
theorem Bind.bind.arg_fx.parDistribDiff_rule
    (f : W → X → 𝒟'(Y,V)) (g : W → 𝒟'(X,U)) (L : U ⊸ V ⊸ W)
    (hf : DistribDifferentiable (fun (w,x) => f w x)) -- `f` has to be nice enough to accomodate action of `g`
    (hg : DistribDifferentiable g) :
    parDistribDeriv (fun w => (g w).bind (f w) L)
    =
    fun w dw =>
      ((parDistribDeriv g w dw).bind (f x · ) L)
      +
      ((g w).bind (fun x => parDistribDeriv (f · x) w dw) L) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Move these around -------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem Distribution.restrict.arg_T.IsSmoothLinearMap_rule (T : W → 𝒟'(X,Y)) (A : Set X)
    (hT : IsSmoothLinearMap R T) :
    IsSmoothLinearMap R (fun w => (T w).restrict A) := sorry_proof

@[fun_prop]
theorem Distribution.restrict.arg_T.IsSmoothLinearMap_rule_simple (A : Set X) :
    IsSmoothLinearMap R (fun (T : 𝒟'(X,Y)) => T.restrict A) := sorry_proof

@[fun_trans]
theorem Distribution.restrict.arg_T.parDistribDeriv_rule (T : W → 𝒟'(X,Y)) (A : Set X)
    (hT : DistribDifferentiable T) :
    parDistribDeriv (fun w => (T w).restrict A)
    =
    fun w dw =>
      (parDistribDeriv T w dw).restrict A := sorry_proof

@[fun_prop]
theorem Function.toDistribution.arg_f.CDifferentiable_rule (f : W → X → Y)
    (hf : ∀ x, CDifferentiable R (f · x)) :
    CDifferentiable R (fun w => (fun x => f w x).toDistribution (R:=R)) := sorry_proof

@[fun_trans]
theorem Function.toDistribution.arg_f.cderiv_rule (f : W → X → Y)
    (hf : ∀ x, CDifferentiable R (f · x)) :
    cderiv R (fun w => (fun x => f w x).toDistribution (R:=R))
    =
    fun w dw =>
      (fun x =>
        let dy := cderiv R (f · x) w dw
        dy).toDistribution := sorry_proof

@[fun_trans]
theorem toDistribution.linear_parDistribDeriv_rule (f : W → X → Y) (L : Y → Z)
    (hL : IsSmoothLinearMap R L) :
    parDistribDeriv (fun w => (fun x => L (f w x)).toDistribution)
    =
    fun w dw =>
      parDistribDeriv (fun w => (fun x => f w x).toDistribution) w dw |>.postComp (fun y ⊸ L y) := by
  funext w dw
  unfold parDistribDeriv Distribution.postComp Function.toDistribution
  ext φ
  simp [ftrans_simp] -- , Distribution.mk_extAction_simproc]
  sorry_proof

----------------------------------------------------------------------------------------------------
-- Integral ----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable [MeasureSpace X] [MeasureSpace Y] [MeasureSpace (X×Y)]

open Notation

@[fun_trans]
theorem cintegral.arg_f.cderiv_distrib_rule (f : W → X → Y) :
    cderiv R (fun w => ∫' x, f w x)
    =
    fun w dw =>
      (parDistribDeriv (fun w => (f w ·).toDistribution) w dw).extAction (fun _ => (1:R)) (fun y ⊸ fun r ⊸ r • y) := sorry_proof

@[fun_trans]
theorem cintegral.arg_f.cderiv_distrib_rule' (f : W → X → R) (A : Set X):
    cderiv R (fun w => ∫' x in A, f w x)
    =
    fun w dw =>
       (parDistribDeriv (fun w => (f w ·).toDistribution) w dw).restrict A |>.extAction (fun _ => (1:R)) (fun y ⊸ fun r ⊸ r • y) := sorry_proof

@[fun_trans]
theorem cintegral.arg_f.parDistribDeriv_rule (f : W → X → Y → Z) :
    parDistribDeriv (fun w => (fun x => ∫' y, f w x y).toDistribution (R:=R))
    =
    fun w dw =>
      let Tf := (fun w => (fun x => (fun y => f w x y).toDistribution (R:=R)).toDistribution (R:=R))
      parDistribDeriv Tf w dw |>.postComp (fun T ⊸ T.extAction (fun _ => (1:R)) (fun z ⊸ fun r ⊸ r • z)) := by
  funext w dw
  unfold parDistribDeriv postComp Function.toDistribution
  ext φ
  simp [ftrans_simp] -- , Distribution.mk_extAction_simproc]
  sorry_proof

@[fun_trans]
theorem cintegral.arg_f.parDistribDeriv_rule' (f : W → X → Y → Z) (B : X → Set Y) :
    parDistribDeriv (fun w => (fun x => ∫' y in B x, f w x y).toDistribution)
    =
    fun w dw =>
      let Tf := (fun w => (fun x => ((fun y => f w x y).toDistribution (R:=R)).restrict (B x)).toDistribution (R:=R))
      parDistribDeriv Tf w dw |>.postComp (fun T ⊸ T.extAction (fun _ => (1:R)) (fun z ⊸ fun r ⊸ r • z)) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Add ---------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem HAdd.hAdd.arg_a0a1.DistribDifferentiable_rule (f g : W → 𝒟'(X,Y))
    (hf : DistribDifferentiable f) (hg : DistribDifferentiable g) :
    DistribDifferentiable (fun w => f w + g w) := sorry_proof

@[fun_trans]
theorem HAdd.hAdd.arg_a0a1.parDistribDeriv_rule (f g : W → 𝒟'(X,Y))
    (hf : DistribDifferentiable f) (hg : DistribDifferentiable g) :
    parDistribDeriv (fun w => f w + g w)
    =
    fun w dw =>
      let dy := parDistribDeriv f w dw
      let dz := parDistribDeriv g w dw
      dy + dz := sorry_proof

----------------------------------------------------------------------------------------------------
-- Sub ---------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem HSub.hSub.arg_a0a1.DistribDifferentiable_rule (f g : W → 𝒟'(X,Y)) :
    -- (hf : DistribDifferentiable f) (hg : DistribDifferentiable g) :
    DistribDifferentiable (fun w => f w - g w)


import SciLean.Core.Distribution.Basic
import SciLean.Core.FunctionTransformations
import SciLean.Core.FunctionPropositions
import SciLean.Core.Notation

open MeasureTheory

namespace SciLean

open Distribution

variable
  {R} [RealScalar R]
  {W} [Vec R W]
  {X} [Vec R X] [MeasureSpace X]
  {Y} [Vec R Y] [Module ℝ Y]
  {Z} [Vec R Z] [Module ℝ Z]
  {U} [Vec R U] -- [Module ℝ U]
  {V} [Vec R V] -- [Module ℝ U]

set_default_scalar R

noncomputable
def diracDeriv (x dx : X) : 𝒟' X := fun φ ⊸ cderiv R φ x dx

@[fun_prop]
def DistribDifferentiableAt (f : X → 𝒟'(Y,Z)) (x : X) :=
  ∀ (φ : X → 𝒟 Y), CDifferentiableAt R φ x → CDifferentiableAt R (fun x => f x (φ x)) x

theorem distribDifferentiableAt_const_test_fun
    {f : X → 𝒟'(Y,Z)} {x : X}
    (hf : DistribDifferentiableAt f x)
    {φ : 𝒟 Y} :
    CDifferentiableAt R (fun x => f x φ) x := by
  apply hf
  fun_prop

@[fun_prop]
def DistribDifferentiable (f : X → 𝒟'(Y,Z)) :=
  ∀ x, DistribDifferentiableAt f x

-- TODO:
-- probably change the definition of `parDistribDeriv` to:
-- ⟨⟨fun φ =>
--    if h : DistribDifferentiableAt f x then
--      ∂ (x':=x;dx), ⟪f x', φ⟫
--    else
--      0 , sorry_proof⟩⟩
-- I believe in that case the function is indeed linear in φ

open Classical in
@[fun_trans]
noncomputable
def parDistribDeriv (f : X → 𝒟'(Y,Z)) (x dx : X) : 𝒟'(Y,Z) :=
  ⟨fun φ => ∂ (x':=x;dx), f x' φ, sorry_proof⟩

@[simp, ftrans_simp]
theorem action_parDistribDeriv (f : X → 𝒟'(Y,Z)) (x dx : X) (φ : 𝒟 Y) :
    parDistribDeriv f x dx φ = ∂ (x':=x;dx), f x' φ := rfl

----------------------------------------------------------------------------------------------------
-- Const rule --------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem DistribDiffrentiable.const_rule (T : 𝒟'(X,Y)) :
    DistribDifferentiable (fun _ : W => T) := by
  intro _ φ hφ; simp; fun_prop

@[fun_trans]
theorem parDistribDeriv.const_rule (T : 𝒟'(X,Y)) :
    parDistribDeriv (fun _ : W => T)
    =
    fun w dw =>
      0 := by
  funext w dw; ext φ
  unfold parDistribDeriv
  fun_trans

----------------------------------------------------------------------------------------------------
-- Pure --------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem dirac.arg_xy.DistribDiffrentiable_rule
    (x : W → X) (hx : CDifferentiable R x) :
    DistribDifferentiable (R:=R) (fun w => dirac (x w))  := by
  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  simp [action_dirac, dirac]
  fun_prop

@[fun_trans]
theorem dirac.arg_x.parDistribDeriv_rule
    (x : W → X) (hx : CDifferentiable R x) :
    parDistribDeriv (R:=R) (fun w => dirac (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      diracDeriv xdx.1 xdx.2 := by --= (dpure (R:=R) ydy.1 ydy.2) := by
  funext w dw; ext φ
  unfold parDistribDeriv dirac diracDeriv
  simp [pure, fwdDeriv, DistribDifferentiableAt]
  fun_trans

----------------------------------------------------------------------------------------------------
-- Composition -------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem DistribDiffrentiable.comp_rule
    (f : Y → 𝒟'(Z,U)) (g : X → Y)
    (hf : DistribDifferentiable f) (hg : CDifferentiable R g) :
    DistribDifferentiable (fun x => f (g x)) := by
  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  apply CDifferentiable.comp_rule (K:=R) (f:=fun xy : X×Y => f xy.2 (φ xy.1)) (g:=fun x => (x, g x))
    (hg:=by fun_prop)
  intro x
  sorry_proof -- is this even true ?

@[fun_trans]
theorem parDistribDeriv.comp_rule
    (f : Y → 𝒟'(Z,U)) (g : X → Y)
    (hf : DistribDifferentiable f) (hg : CDifferentiable R g) :
    parDistribDeriv (fun x => f (g x))
    =
    fun x dx =>
      let ydy := fwdDeriv R g x dx
      parDistribDeriv f ydy.1 ydy.2 := by

  funext x dx; ext φ
  unfold parDistribDeriv
  simp[hg]
  sorry_proof

----------------------------------------------------------------------------------------------------
-- Bind --------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- The assumptions here are definitely not right!!!
-- I think `f` has to be `deg`
@[fun_prop]
theorem Bind.bind.arg_fx.DistribDifferentiable_rule
    (f : X → Y → 𝒟'(Z,V)) (g : X → 𝒟'(Y,U)) (L : U ⊸ V ⊸ W)
    (hf : DistribDifferentiable (fun (x,y) => f x y)) -- `f` has to be nice enough to accomodate action of `g`
    (hg : DistribDifferentiable g) :
    DistribDifferentiable (fun x => (g x).bind (f x) L) := by

  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  simp
  sorry_proof

@[fun_trans]
theorem Bind.bind.arg_fx.parDistribDiff_rule
    (f : W → X → 𝒟'(Y,V)) (g : W → 𝒟'(X,U)) (L : U ⊸ V ⊸ W)
    (hf : DistribDifferentiable (fun (w,x) => f w x)) -- `f` has to be nice enough to accomodate action of `g`
    (hg : DistribDifferentiable g) :
    parDistribDeriv (fun w => (g w).bind (f w) L)
    =
    fun w dw =>
      ((parDistribDeriv g w dw).bind (f x · ) L)
      +
      ((g w).bind (fun x => parDistribDeriv (f · x) w dw) L) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Move these around -------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem Distribution.restrict.arg_T.IsSmoothLinearMap_rule (T : W → 𝒟'(X,Y)) (A : Set X)
    (hT : IsSmoothLinearMap R T) :
    IsSmoothLinearMap R (fun w => (T w).restrict A) := sorry_proof

@[fun_prop]
theorem Distribution.restrict.arg_T.IsSmoothLinearMap_rule_simple (A : Set X) :
    IsSmoothLinearMap R (fun (T : 𝒟'(X,Y)) => T.restrict A) := sorry_proof

@[fun_trans]
theorem Distribution.restrict.arg_T.parDistribDeriv_rule (T : W → 𝒟'(X,Y)) (A : Set X)
    (hT : DistribDifferentiable T) :
    parDistribDeriv (fun w => (T w).restrict A)
    =
    fun w dw =>
      (parDistribDeriv T w dw).restrict A := sorry_proof

@[fun_prop]
theorem Function.toDistribution.arg_f.CDifferentiable_rule (f : W → X → Y)
    (hf : ∀ x, CDifferentiable R (f · x)) :
    CDifferentiable R (fun w => (fun x => f w x).toDistribution (R:=R)) := sorry_proof

@[fun_trans]
theorem Function.toDistribution.arg_f.cderiv_rule (f : W → X → Y)
    (hf : ∀ x, CDifferentiable R (f · x)) :
    cderiv R (fun w => (fun x => f w x).toDistribution (R:=R))
    =
    fun w dw =>
      (fun x =>
        let dy := cderiv R (f · x) w dw
        dy).toDistribution := sorry_proof

@[fun_trans]
theorem toDistribution.linear_parDistribDeriv_rule (f : W → X → Y) (L : Y → Z)
    (hL : IsSmoothLinearMap R L) :
    parDistribDeriv (fun w => (fun x => L (f w x)).toDistribution)
    =
    fun w dw =>
      parDistribDeriv (fun w => (fun x => f w x).toDistribution) w dw |>.postComp (fun y ⊸ L y) := by
  funext w dw
  unfold parDistribDeriv Distribution.postComp Function.toDistribution
  ext φ
  simp [ftrans_simp] -- , Distribution.mk_extAction_simproc]
  sorry_proof

----------------------------------------------------------------------------------------------------
-- Integral ----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable [MeasureSpace X] [MeasureSpace Y] [MeasureSpace (X×Y)]

open Notation

@[fun_trans]
theorem cintegral.arg_f.cderiv_distrib_rule (f : W → X → Y) :
    cderiv R (fun w => ∫' x, f w x)
    =
    fun w dw =>
      (parDistribDeriv (fun w => (f w ·).toDistribution) w dw).extAction (fun _ => (1:R)) (fun y ⊸ fun r ⊸ r • y) := sorry_proof

@[fun_trans]
theorem cintegral.arg_f.cderiv_distrib_rule' (f : W → X → R) (A : Set X):
    cderiv R (fun w => ∫' x in A, f w x)
    =
    fun w dw =>
       (parDistribDeriv (fun w => (f w ·).toDistribution) w dw).restrict A |>.extAction (fun _ => (1:R)) (fun y ⊸ fun r ⊸ r • y) := sorry_proof

@[fun_trans]
theorem cintegral.arg_f.parDistribDeriv_rule (f : W → X → Y → Z) :
    parDistribDeriv (fun w => (fun x => ∫' y, f w x y).toDistribution (R:=R))
    =
    fun w dw =>
      let Tf := (fun w => (fun x => (fun y => f w x y).toDistribution (R:=R)).toDistribution (R:=R))
      parDistribDeriv Tf w dw |>.postComp (fun T ⊸ T.extAction (fun _ => (1:R)) (fun z ⊸ fun r ⊸ r • z)) := by
  funext w dw
  unfold parDistribDeriv postComp Function.toDistribution
  ext φ
  simp [ftrans_simp] -- , Distribution.mk_extAction_simproc]
  sorry_proof

@[fun_trans]
theorem cintegral.arg_f.parDistribDeriv_rule' (f : W → X → Y → Z) (B : X → Set Y) :
    parDistribDeriv (fun w => (fun x => ∫' y in B x, f w x y).toDistribution)
    =
    fun w dw =>
      let Tf := (fun w => (fun x => ((fun y => f w x y).toDistribution (R:=R)).restrict (B x)).toDistribution (R:=R))
      parDistribDeriv Tf w dw |>.postComp (fun T ⊸ T.extAction (fun _ => (1:R)) (fun z ⊸ fun r ⊸ r • z)) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Add ---------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem HAdd.hAdd.arg_a0a1.DistribDifferentiable_rule (f g : W → 𝒟'(X,Y))
    (hf : DistribDifferentiable f) (hg : DistribDifferentiable g) :
    DistribDifferentiable (fun w => f w + g w) := sorry_proof

@[fun_trans]
theorem HAdd.hAdd.arg_a0a1.parDistribDeriv_rule (f g : W → 𝒟'(X,Y))
    (hf : DistribDifferentiable f) (hg : DistribDifferentiable g) :
    parDistribDeriv (fun w => f w + g w)
    =
    fun w dw =>
      let dy := parDistribDeriv f w dw
      let dz := parDistribDeriv g w dw
      dy + dz := sorry_proof

----------------------------------------------------------------------------------------------------
-- Sub ---------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem HSub.hSub.arg_a0a1.DistribDifferentiable_rule (f g : W → 𝒟'(X,Y)) :
    -- (hf : DistribDifferentiable f) (hg : DistribDifferentiable g) :
    DistribDifferentiable (fun w => f w - g w) := sorry_proof

@[fun_trans]
theorem HSub.hSub.arg_a0a1.parDistribDeriv_rule (f g : W → 𝒟'(X,Y)) :
    -- (hf : DistribDifferentiable f) (hg : DistribDifferentiable g) :
    parDistribDeriv (fun w => f w - g w)
    =
    fun w dw =>
      let dy := parDistribDeriv f w dw
      let dz := parDistribDeriv g w dw
      dy - dz


import SciLean.Core.Rand.Rand
import SciLean.Core.Rand.Condition

import SciLean.Core.Distribution.ParametricDistribDeriv

import Mathlib.MeasureTheory.Constructions.Prod.Basic

namespace SciLean

open MeasureTheory
variable
  {R} [RealScalar R]
  {X Z} [MeasurableSpace X] [MeasurableSpace Z]

/-- Kullback–Leibler divergence of `Dₖₗ(P‖Q)` -/
noncomputable
def KLDiv (P Q : Rand X) : R := P.𝔼 (fun x => Scalar.log (P.pdf R Q.ℙ x))

/-- Evidence Lower Bound

```
𝔼_Q [log Q(Z) - log P(Z,X)]
```
reference: https://en.wikipedia.org/wiki/Variational_Bayesian_methods#Evidence_lower_bound
 -/
noncomputable
def ELBO {X Z} [MeasureSpace Z] [MeasureSpace X]
    (P : Rand (Z×X)) (Q : Rand Z) (x : X) : R :=
  - Q.𝔼 (fun z => Scalar.log (Q.pdf R volume z) - Scalar.log (P.pdf R volume (z,x)))

/-- Express `Kullback–Leibler divergence` as log evidence + ELBO -/
theorem kldiv_elbo
    {X Z} [MeasureSpace Z] [MeasureSpace X]
    (P : Rand (Z×X)) (Q : Rand Z) (x : X) :
    KLDiv Q (P.conditionSnd x)
    =
    (Scalar.log (P.snd.pdf R volume x)) - ELBO P Q x := sorry_proof

----------------------------------------------------------------------------------------------------
-- KLDiv properties --------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable
  {W} [Vec R W]
  [Vec R X]

@[fun_trans]
theorem KLDiv.arg_P.cderiv_rule (P : W → Rand X) (Q : Rand X) :
    cderiv R (fun w => KLDiv (R:=R) (P w) Q)
    =
    fun w dw =>
      let dP := parDistribDeriv (fun w => (P w).ℙ.toDistribution (R:=R)) w dw
      dP.extAction' (fun x => Scalar.log ((P w).pdf R Q.ℙ x) - 1)


import SciLean.Core.Rand.Rand
import SciLean.Core.FunctionPropositions
import SciLean.Core.FloatAsReal
-- import SciLean.Modules.Prob.DRand
-- import SciLean.Modules.Prob.FDRand

namespace SciLean.Rand

variable
  {R} [RealScalar R]
  -- {X} [NormedAddCommGroup X] [NormedSpace ℝ X] [NormedSpace R X] [CompleteSpace X] [MeasurableSpace X]
  -- {Y} [NormedAddCommGroup Y] [NormedSpace ℝ Y] [NormedSpace R Y] [CompleteSpace Y] [MeasurableSpace Y]
  -- {Z} [NormedAddCommGroup Z] [NormedSpace ℝ Z] [NormedSpace R Z] [CompleteSpace Z] [MeasurableSpace Z]
  {X : Type _} [MeasurableSpace X] [AddCommGroup X] [Module ℝ X]
  {Y : Type _} [AddCommGroup Y] [Module ℝ Y] [MeasurableSpace Y]
  {Z : Type _} [AddCommGroup Z] [Module ℝ Z]

open Rand

@[rand_pull_E]
theorem bind_pull_mean (x : Rand X) (f : X → Rand Y) :
    (x >>= (fun x' => pure (f x').mean)).mean
    =
    (x >>= f).mean := by simp[rand_push_E]

@[rand_push_E]
theorem ite_push_E {c} [Decidable c] (t e : Rand X) (φ : X → Y):
    (if c then t else e).𝔼 φ = (if c then t.𝔼 φ else e.𝔼 φ) := by
  if h : c then simp[h] else simp[h]

-- I don't think this is a desirable `rand_pull_E` theorem as it duplicates the if statement
-- @[rand_pull_E]
theorem ite_pull_E {c} [Decidable c] (t e : Rand X) (φ ψ : X → Y):
    (if c then t.𝔼 φ else e.𝔼 ψ) = (if c then t else e).𝔼 (if c then φ else ψ) := by
  if h : c then simp[h] else simp[h]

@[rand_push_E]
theorem ite_push_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t else e).mean = (if c then t.mean else e.mean) := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E]
theorem ite_pull_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t.mean else e.mean) = (if c then t else e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_t {c} [Decidable c] (t : Rand X) (e : X) :
    (if c then t.mean else e) = (if c then t else pure e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_f {c} [Decidable c] (t : X) (e : Rand X) :
    (if c then t else e.mean) = (if c then pure t else e).mean := by
  if h : c then simp[h] else simp[h]

-- this has messed up universes
@[rand_pull_E]
theorem pull_E_lambda (r : Rand Y) (f : X → Y → Z) :
    (fun x => r.𝔼 (fun y => f x y))
    =
    r.𝔼 (fun y x => f x y) := sorry_proof

-- this has messed up universes
@[rand_push_E]
theorem push_E_lambda (r : Rand Y) (f : X → Y → Z) :
    r.𝔼 (fun y x => f x y)
    =
    (fun x => r.𝔼 (fun y => f x y)) := sorry_proof

-- can't be simp as it has variable head
theorem pull_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    (f (r.𝔼 φ)) = r.𝔼 (fun x => f (φ x)) := by sorry_proof -- have := hf; sorry_proof

@[rand_push_E]
theorem push_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    r.𝔼 (fun x => f (φ x)) = (f (r.𝔼 φ)) := by rw[pull_E_affine (hf:=hf)]

@[rand_pull_E]
theorem pull_mean_add (x y : Rand X) :
    x.mean + y.mean
    =
    Rand.mean do
      let x' ← x
      let y' ← y
      return x' + y' := sorry_proof

@[rand_pull_E]
theorem pull_mean_add_1 (x : Rand X) (y : X) :
    x.mean + y
    =
    Rand.mean do
      let x' ← x
      return x' + y


import SciLean.Core.Rand.Rand
import SciLean.Core.FunctionPropositions
import SciLean.Core.FloatAsReal
-- import SciLean.Modules.Prob.DRand
-- import SciLean.Modules.Prob.FDRand

namespace SciLean.Rand

variable
  {R} [RealScalar R]
  -- {X} [NormedAddCommGroup X] [NormedSpace ℝ X] [NormedSpace R X] [CompleteSpace X] [MeasurableSpace X]
  -- {Y} [NormedAddCommGroup Y] [NormedSpace ℝ Y] [NormedSpace R Y] [CompleteSpace Y] [MeasurableSpace Y]
  -- {Z} [NormedAddCommGroup Z] [NormedSpace ℝ Z] [NormedSpace R Z] [CompleteSpace Z] [MeasurableSpace Z]
  {X : Type _} [MeasurableSpace X] [AddCommGroup X] [Module ℝ X]
  {Y : Type _} [AddCommGroup Y] [Module ℝ Y] [MeasurableSpace Y]
  {Z : Type _} [AddCommGroup Z] [Module ℝ Z]

open Rand

@[rand_pull_E]
theorem bind_pull_mean (x : Rand X) (f : X → Rand Y) :
    (x >>= (fun x' => pure (f x').mean)).mean
    =
    (x >>= f).mean := by simp[rand_push_E]

@[rand_push_E]
theorem ite_push_E {c} [Decidable c] (t e : Rand X) (φ : X → Y):
    (if c then t else e).𝔼 φ = (if c then t.𝔼 φ else e.𝔼 φ) := by
  if h : c then simp[h] else simp[h]

-- I don't think this is a desirable `rand_pull_E` theorem as it duplicates the if statement
-- @[rand_pull_E]
theorem ite_pull_E {c} [Decidable c] (t e : Rand X) (φ ψ : X → Y):
    (if c then t.𝔼 φ else e.𝔼 ψ) = (if c then t else e).𝔼 (if c then φ else ψ) := by
  if h : c then simp[h] else simp[h]

@[rand_push_E]
theorem ite_push_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t else e).mean = (if c then t.mean else e.mean) := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E]
theorem ite_pull_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t.mean else e.mean) = (if c then t else e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_t {c} [Decidable c] (t : Rand X) (e : X) :
    (if c then t.mean else e) = (if c then t else pure e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_f {c} [Decidable c] (t : X) (e : Rand X) :
    (if c then t else e.mean) = (if c then pure t else e).mean := by
  if h : c then simp[h] else simp[h]

-- this has messed up universes
@[rand_pull_E]
theorem pull_E_lambda (r : Rand Y) (f : X → Y → Z) :
    (fun x => r.𝔼 (fun y => f x y))
    =
    r.𝔼 (fun y x => f x y) := sorry_proof

-- this has messed up universes
@[rand_push_E]
theorem push_E_lambda (r : Rand Y) (f : X → Y → Z) :
    r.𝔼 (fun y x => f x y)
    =
    (fun x => r.𝔼 (fun y => f x y)) := sorry_proof

-- can't be simp as it has variable head
theorem pull_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    (f (r.𝔼 φ)) = r.𝔼 (fun x => f (φ x)) := by sorry_proof -- have := hf; sorry_proof

@[rand_push_E]
theorem push_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    r.𝔼 (fun x => f (φ x)) = (f (r.𝔼 φ)) := by rw[pull_E_affine (hf:=hf)]

@[rand_pull_E]
theorem pull_mean_add (x y : Rand X) :
    x.mean + y.mean
    =
    Rand.mean do
      let x' ← x
      let y' ← y
      return x' + y' := sorry_proof

@[rand_pull_E]
theorem pull_mean_add_1 (x : Rand X) (y : X) :
    x.mean + y
    =
    Rand.mean do
      let x' ← x
      return x' + y := sorry_proof

@[rand_pull_E]
theorem pull_mean_add_2 (x : X) (y : Rand X) :
    x + y.mean
    =
    Rand.mean do
      let y' ← y
      return x + y'


import SciLean.Core.Rand.Rand
import SciLean.Core.FunctionPropositions
import SciLean.Core.FloatAsReal
-- import SciLean.Modules.Prob.DRand
-- import SciLean.Modules.Prob.FDRand

namespace SciLean.Rand

variable
  {R} [RealScalar R]
  -- {X} [NormedAddCommGroup X] [NormedSpace ℝ X] [NormedSpace R X] [CompleteSpace X] [MeasurableSpace X]
  -- {Y} [NormedAddCommGroup Y] [NormedSpace ℝ Y] [NormedSpace R Y] [CompleteSpace Y] [MeasurableSpace Y]
  -- {Z} [NormedAddCommGroup Z] [NormedSpace ℝ Z] [NormedSpace R Z] [CompleteSpace Z] [MeasurableSpace Z]
  {X : Type _} [MeasurableSpace X] [AddCommGroup X] [Module ℝ X]
  {Y : Type _} [AddCommGroup Y] [Module ℝ Y] [MeasurableSpace Y]
  {Z : Type _} [AddCommGroup Z] [Module ℝ Z]

open Rand

@[rand_pull_E]
theorem bind_pull_mean (x : Rand X) (f : X → Rand Y) :
    (x >>= (fun x' => pure (f x').mean)).mean
    =
    (x >>= f).mean := by simp[rand_push_E]

@[rand_push_E]
theorem ite_push_E {c} [Decidable c] (t e : Rand X) (φ : X → Y):
    (if c then t else e).𝔼 φ = (if c then t.𝔼 φ else e.𝔼 φ) := by
  if h : c then simp[h] else simp[h]

-- I don't think this is a desirable `rand_pull_E` theorem as it duplicates the if statement
-- @[rand_pull_E]
theorem ite_pull_E {c} [Decidable c] (t e : Rand X) (φ ψ : X → Y):
    (if c then t.𝔼 φ else e.𝔼 ψ) = (if c then t else e).𝔼 (if c then φ else ψ) := by
  if h : c then simp[h] else simp[h]

@[rand_push_E]
theorem ite_push_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t else e).mean = (if c then t.mean else e.mean) := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E]
theorem ite_pull_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t.mean else e.mean) = (if c then t else e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_t {c} [Decidable c] (t : Rand X) (e : X) :
    (if c then t.mean else e) = (if c then t else pure e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_f {c} [Decidable c] (t : X) (e : Rand X) :
    (if c then t else e.mean) = (if c then pure t else e).mean := by
  if h : c then simp[h] else simp[h]

-- this has messed up universes
@[rand_pull_E]
theorem pull_E_lambda (r : Rand Y) (f : X → Y → Z) :
    (fun x => r.𝔼 (fun y => f x y))
    =
    r.𝔼 (fun y x => f x y) := sorry_proof

-- this has messed up universes
@[rand_push_E]
theorem push_E_lambda (r : Rand Y) (f : X → Y → Z) :
    r.𝔼 (fun y x => f x y)
    =
    (fun x => r.𝔼 (fun y => f x y)) := sorry_proof

-- can't be simp as it has variable head
theorem pull_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    (f (r.𝔼 φ)) = r.𝔼 (fun x => f (φ x)) := by sorry_proof -- have := hf; sorry_proof

@[rand_push_E]
theorem push_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    r.𝔼 (fun x => f (φ x)) = (f (r.𝔼 φ)) := by rw[pull_E_affine (hf:=hf)]

@[rand_pull_E]
theorem pull_mean_add (x y : Rand X) :
    x.mean + y.mean
    =
    Rand.mean do
      let x' ← x
      let y' ← y
      return x' + y' := sorry_proof

@[rand_pull_E]
theorem pull_mean_add_1 (x : Rand X) (y : X) :
    x.mean + y
    =
    Rand.mean do
      let x' ← x
      return x' + y := sorry_proof

@[rand_pull_E]
theorem pull_mean_add_2 (x : X) (y : Rand X) :
    x + y.mean
    =
    Rand.mean do
      let y' ← y
      return x + y' := sorry_proof

@[rand_pull_E]
theorem pull_mean_sub (x y : Rand X) :
    x.mean - y.mean
    =
    Rand.mean do
      let x' ← x
      let y' ← y
      return x' - y' := sorry_proof

@[rand_pull_E]
theorem pull_mean_smul [Module R X] (r : R) (x : Rand X) :
    r • x.mean
    =
    Rand.mean do
      let x' ← x
      return r • x' := sorry_proof

@[rand_pull_E]
theorem pull_mean_mul (r : R) (x : Rand R) :
    r * x.mean
    =
    Rand.mean do
      let x' ← x
      return r * x' := sorry_proof

@[rand_pull_E]
theorem pull_mean_div (x : Rand R) (y : R) :
    x.mean / y
    =
    Rand.mean do
      let x' ← x
      return x' / y


import Mathlib.MeasureTheory.Measure.GiryMonad
import Mathlib.MeasureTheory.Decomposition.Lebesgue
import Mathlib.MeasureTheory.Constructions.Prod.Basic

import SciLean.Core.FunctionPropositions
import SciLean.Core.FunctionSpaces
import SciLean.Core.Integral.CIntegral
import SciLean.Core.Distribution.TestFunction
import SciLean.Core.Distribution.SimpAttr
import SciLean.Util.SorryProof
import SciLean.Util.Limit

open MeasureTheory ENNReal Classical

namespace SciLean

variable
  {R} [RealScalar R]
  {W} [Vec R W] [Module ℝ W]
  {X} [Vec R X] -- [TopologicalSpace X] [space : TCOr (Vec R X) (DiscreteTopology X)]
  {Y} [Vec R Y] [Module ℝ Y]
  {Z} [Vec R Z]
  {U} [Vec R U]
  {V} [Vec R V]

set_default_scalar R

variable (R X Y)
abbrev Distribution := (𝒟 X) ⊸[R] Y
variable {R X Y}

notation "𝒟'" X => Distribution defaultScalar% X defaultScalar%
notation "𝒟'" "(" X ", " Y ")" => Distribution defaultScalar% X Y

@[app_unexpander Distribution] def unexpandDistribution : Lean.PrettyPrinter.Unexpander
  | `($(_) $_ $X $Y) => `(𝒟'($X,$Y))
  | _ => throw ()

@[ext]
theorem Distribution.ext (x y : 𝒟'(X,Y)) :
    (∀ (φ : 𝒟 X), x φ = y φ)
    →
    x = y := by

  apply SmoothLinearMap.ext

----------------------------------------------------------------------------------------------------
-- Algebra -----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- instance : Zero (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Add (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Sub (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Neg (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : SMul R (𝒟'(X,Y)) := by unfold Distribution; infer_instance
instance [Module ℝ Y] : SMul ℝ (𝒟'(X,Y)) := ⟨fun r f => ⟨fun φ => r • (f φ), sorry_proof⟩⟩

-- instance : UniformSpace (𝒟'(X,Y)) :=  by unfold Distribution; infer_instance
-- instance : Vec R (𝒟'(X,Y)) := by unfold Distribution; infer_instance
instance [Module ℝ Y] : Module ℝ (𝒟'(X,Y)) := Module.mkSorryProofs

----------------------------------------------------------------------------------------------------
-- Extended action ---------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

open BigOperators in
@[pp_dot]
noncomputable
def Distribution.extAction (T : 𝒟'(X,Y)) (φ : X → Z) (L : Y ⊸ Z ⊸ W) : W :=
  if h : ∃ (zₙ : ℕ → Z) (φₙ : ℕ → 𝒟 X), ∀ x, ∑' i, φₙ i x • zₙ i = φ x then
    let zₙ := Classical.choose h
    let φₙ := (Classical.choose_spec h).choose
    ∑' i, L (T (φₙ i)) (zₙ i)
  else
    0

namespace Distribution
scoped notation "⟪" T ", " φ "⟫[" L "]" => Distribution.extAction T φ L
end Distribution

noncomputable
abbrev Distribution.extAction' (T : 𝒟'(X,Y)) (φ : X → R) : Y := T.extAction φ (fun y ⊸ fun r ⊸ r • y)

noncomputable
abbrev Distribution.integrate (T : 𝒟'(X,Y)) : Y := T.extAction' (fun _ => 1)

@[fun_prop]
theorem TestFunction.apply_IsSmoothLinearMap : IsSmoothLinearMap R fun (φ : 𝒟 X) => (φ : X → R) := sorry_proof

theorem Distribution.mk_extAction (T : (X → R) → Y) (hT : IsSmoothLinearMap R (fun φ : 𝒟 X => T φ)) (φ : X → R) :
    Distribution.extAction (SmoothLinearMap.mk' R (fun (φ : 𝒟 X) => T φ) hT : Distribution _ _ _) φ (fun y ⊸ fun r ⊸ r • y) = T φ := sorry_proof

-- This is definitely not true as stated, what kind of condistions do we need on `φ` and `T`?
@[fun_prop]
theorem Distribution.extAction.arg_φ.IsSmoothLinearMap (T : 𝒟'(X,U)) (φ : W → X → V) (L : U ⊸ V ⊸ Z)
    (hφ : IsSmoothLinearMap R φ) :
    IsSmoothLinearMap R (fun w => T.extAction (φ w) L) := sorry_proof

@[fun_prop]
theorem Distribution.extAction.arg_T.IsSmoothLinearMap (T : W → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ Z)
    (hT : IsSmoothLinearMap R T) :
    IsSmoothLinearMap R (fun w => (T w).extAction φ L) := sorry_proof

-- open Lean Meta in
-- /-- Simproc to apply `Distribution.mk_extAction` theorem -/
-- simproc_decl Distribution.mk_extAction_simproc (Distribution.extAction (Distribution.mk (SmoothLinearMap.mk _ _)) _) := fun e => do

--   let φ := e.appArg!
--   let T := e.appFn!.appArg!

--   let .lam xName xType xBody xBi := T.appArg!.appFn!.appArg!
--     | return .continue
--   let hT := T.appArg!.appArg!

--   withLocalDecl xName xBi xType fun x => do
--   let R := xType.getArg! 0
--   let X := xType.getArg! 2
--   withLocalDecl `φ' xBi (← mkArrow X R) fun φ' => do
--     let b := xBody.instantiate1 x
--     let b := b.replace (fun e' =>
--       if e'.isAppOf ``DFunLike.coe &&
--          5 ≤ e'.getAppNumArgs &&
--          e'.getArg! 4 == x then
--           .some (mkAppN φ' e'.getAppArgs[5:])
--       else
--         .none)

--     if b.containsFVar x.fvarId! then
--       return .continue

--     let T ← mkLambdaFVars #[φ'] b
--     let prf ← mkAppM ``Distribution.mk_extAction #[T, hT, φ]
--     return .visit {expr := T.beta #[φ], proof? := prf}

----------------------------------------------------------------------------------------------------
-- Monadic structure -------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- def dirac (x : X) : Distribution X := fun φ => φ x

-- instance : Monad (Distribution R) where
--   pure := fun x => ⟨fun φ => φ x⟩
--   bind := fun x f => ⟨fun φ => ⟪x, fun x' => ⟪(f x'), φ⟫⟫⟩

-- instance : LawfulMonad (Distribution R) where
--   bind_pure_comp := by intros; rfl
--   bind_map       := by intros; rfl
--   pure_bind      := by intros; rfl
--   bind_assoc     := by intros; rfl
--   map_const      := by intros; rfl
--   id_map         := by intros; rfl
--   seqLeft_eq     := by intros; rfl
--   seqRight_eq    := by intros; rfl
--   pure_seq       := by intros; rfl

def dirac (x : X) : 𝒟' X := fun φ ⊸ φ x

open Notation
noncomputable
def Distribution.bind (x' : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) : 𝒟'(Y,W) :=
  fun φ ⊸ x'.extAction (fun x => (f x).extAction φ (fun v ⊸ fun r ⊸ r • v)) L

----------------------------------------------------------------------------------------------------
-- Basic identities --------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[simp, ftrans_simp]
theorem action_dirac (x : X) (φ : 𝒟 X) : dirac x φ = φ x := by simp[dirac]

@[simp, ftrans_simp]
theorem action_bind (x : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) (φ : 𝒟 Y) :
    x.bind f L φ = x.extAction (fun x' => (f x').extAction' φ) L := by
  simp[Distribution.bind]

-- @[simp, ftrans_simp]
-- theorem extAction_bind (x : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) (φ : Y → Z) (K : W ⊸ Z ⊸ W') :
--     (x.bind f L).extAction φ K = x.extAction (fun x' => (f x').extAction φ (sorry : V ⊸ Z ⊸ V⊗Z)) (sorry : U ⊸ (V⊗Z) ⊸ W') := by
--   simp [Distribution.bind]

----------------------------------------------------------------------------------------------------
-- Arithmetics -------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

section Arithmetics

@[simp, ftrans_simp, action_push]
theorem Distribution.zero_extAction (φ : X → V) (L : U ⊸ V ⊸ W) : (0 : 𝒟'(X,U)).extAction φ L = 0 := by
  unfold extAction; simp

-- todo: this needs some integrability condition
@[action_push]
theorem Distribution.add_extAction (T T' : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    ((T + T') : 𝒟'(X,U)).extAction φ L = T.extAction φ L + T'.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.sub_extAction (T T' : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (T - T').extAction φ L = T.extAction φ L - T'.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.smul_extAction (r : R) (T : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W)  :
    (r • T).extAction φ L = r • T.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.neg_extAction (T : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (- T).extAction φ L = - T.extAction φ L := by sorry_proof

open BigOperators in
@[action_push]
theorem Distribution.fintype_sum_extAction {I} [Fintype I] (T : I → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (∑ i, T i).extAction φ L = ∑ i, (T i).extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.indextype_sum_extAction {I} [IndexType I] (T : I → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (∑ i, T i).extAction φ L = ∑ i, (T i).extAction φ L := by sorry_proof

end Arithmetics

----------------------------------------------------------------------------------------------------
-- Distributional if statement ---------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable [MeasureSpace X]

open Classical Notation in
noncomputable
def iteD (A : Set X) (t e : 𝒟'(X,Y)) : 𝒟'(X,Y) :=
  fun φ ⊸
    t.extAction (fun x => if x ∈ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) +
    e.extAction (fun x => if x ∈ A then 0 else φ x) (fun y ⊸ fun r ⊸ r • y)

open Lean.Parser Term in
syntax withPosition("ifD " term " then "
    ppDedent(ppLine ppSpace ppSpace) term ppDedent(ppLine)
  "else"
    ppDedent(ppLine ppSpace ppSpace) term) : term

macro_rules
  | `(ifD $A then $t else $e) => `(iteD $A $t $e)

open Lean Parser in
@[app_unexpander iteD]
def unexpandIteD : Lean.PrettyPrinter.Unexpander
  | `($(_) $A $t $e) => `(ifD $A then $t else $e)
  | _ => throw ()

@[action_push]
theorem Distribution.action_iteD (A : Set X) (t e : 𝒟'(X,Y)) (φ : 𝒟 X) :
   iteD A t e φ =
        t.extAction (fun x => if x ∈ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) +
        e.extAction (fun x => if x ∉ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) := by sorry_proof

@[simp, ftrans_simp]
theorem Distribution.iteD_same (A : Set X) (u : 𝒟'(X,Y)) :
   iteD A u u = u := by sorry_proof

@[action_push]
theorem Distribution.extAction_iteD (A : Set X) (t e : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (iteD A t e).extAction φ L =
        t.extAction (fun x => if x ∈ A then φ x else 0) L +
        e.extAction (fun x => if x ∉ A then φ x else 0) L := by sorry_proof

@[fun_prop]
theorem iteD.arg_te.IsSmoothLinearMap_rule (A : Set X) (t e : W → 𝒟'(X,Y))
    (ht : IsSmoothLinearMap R t) (he : IsSmoothLinearMap R e) :
    IsSmoothLinearMap R (fun w => iteD A (t w) (e w)) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Set restriction ---------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[pp_dot]
noncomputable
def Distribution.restrict (T : 𝒟'(X,Y)) (A : Set X) : 𝒟'(X,Y) :=
  ifD A then T else 0

@[simp,ftrans_simp]
theorem restrict_univ (T : 𝒟'(X,Y))  :
    T.restrict Set.univ = T := sorry_proof

@[simp,ftrans_simp]
theorem zero_restrict (A : Set X) :
    (0 : 𝒟'(X,Y)).restrict A = 0 := sorry_proof

@[restrict_push]
theorem add_restrict (T S : 𝒟'(X,Y)) (A : Set X) :
    (T + S).restrict A = T.restrict A + S.restrict A := sorry_proof

@[restrict_pull]
theorem add_restrict' (T S : 𝒟'(X,Y)) (A : Set X) :
    T.restrict A + S.restrict A = (T + S).restrict A := sorry_proof

@[restrict_push]
theorem sub_restrict (T S : 𝒟'(X,Y)) (A : Set X) :
    (T - S).restrict A = T.restrict A - S.restrict A := sorry_proof

@[restrict_pull]
theorem sub_restrict' (T S : 𝒟'(X,Y)) (A : Set X) :
    T.restrict A - S.restrict A = (T - S).restrict A := sorry_proof

@[restrict_push]
theorem smul_restrict (r : R) (T : 𝒟'(X,Y)) (A : Set X) :
    (r • T).restrict A = r • (T.restrict A) := sorry_proof

@[restrict_pull]
theorem smul_restrict' (r : R) (T : 𝒟'(X,Y)) (A : Set X) :
    r • (T.restrict A) = (r • T).restrict A := sorry_proof

@[restrict_push]
theorem neg_restrict (T : 𝒟'(X,Y)) (A : Set X) :
    (- T).restrict A = - (T.restrict A) := sorry_proof

@[restrict_pull]
theorem neg_restrict' (T : 𝒟'(X,Y)) (A : Set X) :
    - (T.restrict A) = (- T).restrict A := sorry_proof

open BigOperators in
@[restrict_push]
theorem finset_sum_restrict {I} [Fintype I] (T : I → 𝒟'(X,Y)) (A : Set X) :
    (∑ i, T i).restrict A = ∑ i, (T i).restrict A := sorry_proof

open BigOperators in
@[restrict_pull]
theorem finset_sum_restrict' {I} [Fintype I] (T : I → 𝒟'(X,Y)) (A : Set X) :
    ∑ i, (T i).restrict A = (∑ i, T i).restrict A := sorry_proof

@[restrict_push]
theorem indextype_sum_restrict {I} [IndexType I] (T : I → 𝒟' X) (A : Set X) :
    (∑ i, T i).restrict A = ∑ i, (T i).restrict A := sorry_proof

@[restrict_pull]
theorem indextype_sum_restrict' {I} [IndexType I] (T : I → 𝒟' X) (A : Set X) :
    ∑ i, (T i).restrict A = (∑ i, T i).restrict A := sorry_proof

@[simp,ftrans_simp]
theorem iteD_restrict (T : 𝒟'(X,Y)) (A : Set X) :
    (ifD A then T else 0) = T.restrict A := by rfl

@[simp,ftrans_simp]
theorem iteD_restrict' (T : 𝒟'(X,Y)) (A : Set X) :
    (ifD A then 0 else T) = T.restrict Aᶜ := sorry_proof

@[action_push]
theorem Distribution.extAction_iteD' (A B : Set X) (t e : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    ((iteD A t e).restrict B).extAction φ L =
        (t.restrict B).extAction (fun x => if x ∈ A then φ x else 0) L +
        (e.restrict B).extAction (fun x => if x ∉ A then φ x else 0) L


import SciLean.Core.FunctionPropositions.IsSmoothLinearMap

set_option linter.unusedVariables false

namespace SciLean

variable
  (K : Type _) [RCLike K]
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {W : Type _} [Vec K W]
  {ι : Type _} [IndexType ι] [LawfulIndexType ι] [DecidableEq ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

-- Function space --------------------------------------------------------------
--------------------------------------------------------------------------------

structure SmoothLinearMap (X Y : Type _) [Vec K X] [Vec K Y] where
  toFun : X → Y
  is_smooth_linear_map : IsSmoothLinearMap K toFun

instance : FunLike (SmoothLinearMap K X Y) X Y where
  coe f := f.toFun
  coe_injective' := sorry_proof

macro X:term:25 " ⊸[" K:term "] " Y:term:26 : term =>
  `(SmoothLinearMap $K $X $Y)

macro X:term:25 " ⊸ " Y:term:26 : term =>
  `(SmoothLinearMap defaultScalar% $X $Y)

@[app_unexpander SmoothLinearMap] def unexpandSmoothLinearMap : Lean.PrettyPrinter.Unexpander
  | `($(_) $R $X $Y) => `($X ⊸[$R] $Y)
  | _ => throw ()

@[fun_prop]
theorem SmoothLinearMap_apply_right (f : X ⊸[K] Y) : IsSmoothLinearMap K (fun x => f x) := f.2

-- Lambda function notation ----------------------------------------------------
--------------------------------------------------------------------------------

variable {K}

@[simp, ftrans_simp]
theorem SmoothLinearMap.mk_eval (x : X) (f : X → Y) (hf : IsSmoothLinearMap K f) :
    mk f hf x = f x := by rfl

@[simp]
theorem SmoothLinearMap.eta_reduce (f : SmoothLinearMap K X Y) :
    (mk f.1 f.2) = f := by rfl

@[ext]
theorem SmoothLinearMap.ext (f g : X ⊸[K] Y) : (∀ x, f x = g x) → f = g := sorry_proof

variable (K)
def SmoothLinearMap.mk' (f : X → Y) (hf : IsSmoothLinearMap K f) : X ⊸[K] Y := ⟨f,hf⟩

@[simp, ftrans_simp]
theorem SmoothLinearMap.mk'_eval (x : X) (f : X → Y) (hf : IsSmoothLinearMap K f) :
    mk' K f hf x = f x := by rfl

open Lean Parser Term in
macro "fun " x:funBinder " ⊸[" K:term "] " b:term : term =>
  `(SmoothLinearMap.mk' $K (fun $x => $b) (by fun_prop))

open Lean Parser Term in
macro "fun " x:funBinder " ⊸ " b:term : term =>
  `(SmoothLinearMap.mk' defaultScalar% (fun $x => $b) (by fun_prop))

@[app_unexpander SmoothLinearMap.mk'] def unexpandSmoothLinearMapMk' : Lean.PrettyPrinter.Unexpander

  | `($(_) $R $f:term $_:term) =>
    match f with
    | `(fun $x':ident => $b:term) => `(fun $x' ⊸[$R] $b)
    | `(fun ($x':ident : $ty) => $b:term) => `(fun ($x' : $ty) ⊸[$R] $b)
    | `(fun $x':ident : $ty => $b:term) => `(fun ($x' : $ty) ⊸[$R] $b)
    | _  => throw ()
  | _  => throw ()

@[app_unexpander SmoothLinearMap.mk] def unexpandSmoothLinearMapMk : Lean.PrettyPrinter.Unexpander

  | `($(_) $f:term $_:term) =>
    match f with
    | `(fun $x':ident => $b:term) => `(fun $x' ⊸ $b)
    | `(fun ($x':ident : $ty) => $b:term) => `(fun ($x' : $ty) ⊸ $b)
    | `(fun $x':ident : $ty => $b:term) => `(fun ($x' : $ty) ⊸ $b)
    | _  => throw ()
  | _  => throw ()

-- Algebra ---------------------------------------------------------------------
--------------------------------------------------------------------------------

instance : Add (X ⊸[K] Y) := ⟨fun f g => fun x ⊸[K] f x + g x⟩
instance : Sub (X ⊸[K] Y) := ⟨fun f g => fun x ⊸[K] f x - g x⟩
instance : Neg (X ⊸[K] Y) := ⟨fun f => fun x ⊸[K] - f x⟩
instance : SMul K (X ⊸[K] Y) := ⟨fun r f => fun x ⊸[K] r • f x⟩
instance : Zero (X ⊸[K] Y) := ⟨fun x ⊸[K] 0⟩

section AlgebraSimps

variable (f g : X ⊸[K] Y) (x : X) (r : K)

@[simp, ftrans_simp]
theorem SmoothLinearMap.add_apply : (f + g) x = f x + g x := by rfl

@[simp, ftrans_simp]
theorem SmoothLinearMap.sub_apply : (f - g) x = f x - g x := by rfl

@[simp, ftrans_simp]
theorem SmoothLinearMap.neg_apply : (- f) x = - f x := by rfl

@[simp, ftrans_simp]
theorem SmoothLinearMap.smul_apply : (r • f) x = r • f x := by rfl

@[simp, ftrans_simp]
theorem SmoothLinearMap.zero_apply : (0 : X⊸[K]Y) x = 0 := by sorry_proof

@[simp,ftrans_simp]
theorem SmoothLinearMap.apply_zero (f : X ⊸[K] Y) : f 0 = 0


import Mathlib.MeasureTheory.Measure.GiryMonad
import Mathlib.MeasureTheory.Decomposition.Lebesgue
import Mathlib.MeasureTheory.Constructions.Prod.Basic

import SciLean.Core.FunctionPropositions
import SciLean.Core.FunctionSpaces
import SciLean.Core.Integral.CIntegral
import SciLean.Core.Distribution.TestFunction
import SciLean.Core.Distribution.SimpAttr
import SciLean.Util.SorryProof
import SciLean.Util.Limit

open MeasureTheory ENNReal Classical

namespace SciLean

variable
  {R} [RealScalar R]
  {W} [Vec R W] [Module ℝ W]
  {X} [Vec R X] -- [TopologicalSpace X] [space : TCOr (Vec R X) (DiscreteTopology X)]
  {Y} [Vec R Y] [Module ℝ Y]
  {Z} [Vec R Z]
  {U} [Vec R U]
  {V} [Vec R V]

set_default_scalar R

variable (R X Y)
abbrev Distribution := (𝒟 X) ⊸[R] Y
variable {R X Y}

notation "𝒟'" X => Distribution defaultScalar% X defaultScalar%
notation "𝒟'" "(" X ", " Y ")" => Distribution defaultScalar% X Y

@[app_unexpander Distribution] def unexpandDistribution : Lean.PrettyPrinter.Unexpander
  | `($(_) $_ $X $Y) => `(𝒟'($X,$Y))
  | _ => throw ()

@[ext]
theorem Distribution.ext (x y : 𝒟'(X,Y)) :
    (∀ (φ : 𝒟 X), x φ = y φ)
    →
    x = y := by

  apply SmoothLinearMap.ext

----------------------------------------------------------------------------------------------------
-- Algebra -----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- instance : Zero (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Add (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Sub (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Neg (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : SMul R (𝒟'(X,Y)) := by unfold Distribution; infer_instance
instance [Module ℝ Y] : SMul ℝ (𝒟'(X,Y)) := ⟨fun r f => ⟨fun φ => r • (f φ), sorry_proof⟩⟩

-- instance : UniformSpace (𝒟'(X,Y)) :=  by unfold Distribution; infer_instance
-- instance : Vec R (𝒟'(X,Y)) := by unfold Distribution; infer_instance
instance [Module ℝ Y] : Module ℝ (𝒟'(X,Y)) := Module.mkSorryProofs

----------------------------------------------------------------------------------------------------
-- Extended action ---------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

open BigOperators in
@[pp_dot]
noncomputable
def Distribution.extAction (T : 𝒟'(X,Y)) (φ : X → Z) (L : Y ⊸ Z ⊸ W) : W :=
  if h : ∃ (zₙ : ℕ → Z) (φₙ : ℕ → 𝒟 X), ∀ x, ∑' i, φₙ i x • zₙ i = φ x then
    let zₙ := Classical.choose h
    let φₙ := (Classical.choose_spec h).choose
    ∑' i, L (T (φₙ i)) (zₙ i)
  else
    0

namespace Distribution
scoped notation "⟪" T ", " φ "⟫[" L "]" => Distribution.extAction T φ L
end Distribution

noncomputable
abbrev Distribution.extAction' (T : 𝒟'(X,Y)) (φ : X → R) : Y := T.extAction φ (fun y ⊸ fun r ⊸ r • y)

noncomputable
abbrev Distribution.integrate (T : 𝒟'(X,Y)) : Y := T.extAction' (fun _ => 1)

@[fun_prop]
theorem TestFunction.apply_IsSmoothLinearMap : IsSmoothLinearMap R fun (φ : 𝒟 X) => (φ : X → R) := sorry_proof

theorem Distribution.mk_extAction (T : (X → R) → Y) (hT : IsSmoothLinearMap R (fun φ : 𝒟 X => T φ)) (φ : X → R) :
    Distribution.extAction (SmoothLinearMap.mk' R (fun (φ : 𝒟 X) => T φ) hT : Distribution _ _ _) φ (fun y ⊸ fun r ⊸ r • y) = T φ := sorry_proof

-- This is definitely not true as stated, what kind of condistions do we need on `φ` and `T`?
@[fun_prop]
theorem Distribution.extAction.arg_φ.IsSmoothLinearMap (T : 𝒟'(X,U)) (φ : W → X → V) (L : U ⊸ V ⊸ Z)
    (hφ : IsSmoothLinearMap R φ) :
    IsSmoothLinearMap R (fun w => T.extAction (φ w) L) := sorry_proof

@[fun_prop]
theorem Distribution.extAction.arg_T.IsSmoothLinearMap (T : W → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ Z)
    (hT : IsSmoothLinearMap R T) :
    IsSmoothLinearMap R (fun w => (T w).extAction φ L) := sorry_proof

-- open Lean Meta in
-- /-- Simproc to apply `Distribution.mk_extAction` theorem -/
-- simproc_decl Distribution.mk_extAction_simproc (Distribution.extAction (Distribution.mk (SmoothLinearMap.mk _ _)) _) := fun e => do

--   let φ := e.appArg!
--   let T := e.appFn!.appArg!

--   let .lam xName xType xBody xBi := T.appArg!.appFn!.appArg!
--     | return .continue
--   let hT := T.appArg!.appArg!

--   withLocalDecl xName xBi xType fun x => do
--   let R := xType.getArg! 0
--   let X := xType.getArg! 2
--   withLocalDecl `φ' xBi (← mkArrow X R) fun φ' => do
--     let b := xBody.instantiate1 x
--     let b := b.replace (fun e' =>
--       if e'.isAppOf ``DFunLike.coe &&
--          5 ≤ e'.getAppNumArgs &&
--          e'.getArg! 4 == x then
--           .some (mkAppN φ' e'.getAppArgs[5:])
--       else
--         .none)

--     if b.containsFVar x.fvarId! then
--       return .continue

--     let T ← mkLambdaFVars #[φ'] b
--     let prf ← mkAppM ``Distribution.mk_extAction #[T, hT, φ]
--     return .visit {expr := T.beta #[φ], proof? := prf}

----------------------------------------------------------------------------------------------------
-- Monadic structure -------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- def dirac (x : X) : Distribution X := fun φ => φ x

-- instance : Monad (Distribution R) where
--   pure := fun x => ⟨fun φ => φ x⟩
--   bind := fun x f => ⟨fun φ => ⟪x, fun x' => ⟪(f x'), φ⟫⟫⟩

-- instance : LawfulMonad (Distribution R) where
--   bind_pure_comp := by intros; rfl
--   bind_map       := by intros; rfl
--   pure_bind      := by intros; rfl
--   bind_assoc     := by intros; rfl
--   map_const      := by intros; rfl
--   id_map         := by intros; rfl
--   seqLeft_eq     := by intros; rfl
--   seqRight_eq    := by intros; rfl
--   pure_seq       := by intros; rfl

def dirac (x : X) : 𝒟' X := fun φ ⊸ φ x

open Notation
noncomputable
def Distribution.bind (x' : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) : 𝒟'(Y,W) :=
  fun φ ⊸ x'.extAction (fun x => (f x).extAction φ (fun v ⊸ fun r ⊸ r • v)) L

----------------------------------------------------------------------------------------------------
-- Basic identities --------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[simp, ftrans_simp]
theorem action_dirac (x : X) (φ : 𝒟 X) : dirac x φ = φ x := by simp[dirac]

@[simp, ftrans_simp]
theorem action_bind (x : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) (φ : 𝒟 Y) :
    x.bind f L φ = x.extAction (fun x' => (f x').extAction' φ) L := by
  simp[Distribution.bind]

-- @[simp, ftrans_simp]
-- theorem extAction_bind (x : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) (φ : Y → Z) (K : W ⊸ Z ⊸ W') :
--     (x.bind f L).extAction φ K = x.extAction (fun x' => (f x').extAction φ (sorry : V ⊸ Z ⊸ V⊗Z)) (sorry : U ⊸ (V⊗Z) ⊸ W') := by
--   simp [Distribution.bind]

----------------------------------------------------------------------------------------------------
-- Arithmetics -------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

section Arithmetics

@[simp, ftrans_simp, action_push]
theorem Distribution.zero_extAction (φ : X → V) (L : U ⊸ V ⊸ W) : (0 : 𝒟'(X,U)).extAction φ L = 0 := by
  unfold extAction; simp

-- todo: this needs some integrability condition
@[action_push]
theorem Distribution.add_extAction (T T' : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    ((T + T') : 𝒟'(X,U)).extAction φ L = T.extAction φ L + T'.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.sub_extAction (T T' : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (T - T').extAction φ L = T.extAction φ L - T'.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.smul_extAction (r : R) (T : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W)  :
    (r • T).extAction φ L = r • T.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.neg_extAction (T : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (- T).extAction φ L = - T.extAction φ L := by sorry_proof

open BigOperators in
@[action_push]
theorem Distribution.fintype_sum_extAction {I} [Fintype I] (T : I → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (∑ i, T i).extAction φ L = ∑ i, (T i).extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.indextype_sum_extAction {I} [IndexType I] (T : I → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (∑ i, T i).extAction φ L = ∑ i, (T i).extAction φ L := by sorry_proof

end Arithmetics

----------------------------------------------------------------------------------------------------
-- Distributional if statement ---------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable [MeasureSpace X]

open Classical Notation in
noncomputable
def iteD (A : Set X) (t e : 𝒟'(X,Y)) : 𝒟'(X,Y) :=
  fun φ ⊸
    t.extAction (fun x => if x ∈ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) +
    e.extAction (fun x => if x ∈ A then 0 else φ x) (fun y ⊸ fun r ⊸ r • y)

open Lean.Parser Term in
syntax withPosition("ifD " term " then "
    ppDedent(ppLine ppSpace ppSpace) term ppDedent(ppLine)
  "else"
    ppDedent(ppLine ppSpace ppSpace) term) : term

macro_rules
  | `(ifD $A then $t else $e) => `(iteD $A $t $e)

open Lean Parser in
@[app_unexpander iteD]
def unexpandIteD : Lean.PrettyPrinter.Unexpander
  | `($(_) $A $t $e) => `(ifD $A then $t else $e)
  | _ => throw ()

@[action_push]
theorem Distribution.action_iteD (A : Set X) (t e : 𝒟'(X,Y)) (φ : 𝒟 X) :
   iteD A t e φ =
        t.extAction (fun x => if x ∈ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) +
        e.extAction (fun x => if x ∉ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) := by sorry_proof

@[simp, ftrans_simp]
theorem Distribution.iteD_same (A : Set X) (u : 𝒟'(X,Y)) :
   iteD A u u = u := by sorry_proof

@[action_push]
theorem Distribution.extAction_iteD (A : Set X) (t e : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (iteD A t e).extAction φ L =
        t.extAction (fun x => if x ∈ A then φ x else 0) L +
        e.extAction (fun x => if x ∉ A then φ x else 0) L := by sorry_proof

@[fun_prop]
theorem iteD.arg_te.IsSmoothLinearMap_rule (A : Set X) (t e : W → 𝒟'(X,Y))
    (ht : IsSmoothLinearMap R t) (he : IsSmoothLinearMap R e) :
    IsSmoothLinearMap R (fun w => iteD A (t w) (e w)) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Set restriction ---------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[pp_dot]
noncomputable
def Distribution.restrict (T : 𝒟'(X,Y)) (A : Set X) : 𝒟'(X,Y) :=
  ifD A then T else 0

@[simp,ftrans_simp]
theorem restrict_univ (T : 𝒟'(X,Y))  :
    T.restrict Set.univ = T := sorry_proof

@[simp,ftrans_simp]
theorem zero_restrict (A : Set X) :
    (0 : 𝒟'(X,Y)).restrict A = 0 := sorry_proof

@[restrict_push]
theorem add_restrict (T S : 𝒟'(X,Y)) (A : Set X) :
    (T + S).restrict A = T.restrict A + S.restrict A := sorry_proof

@[restrict_pull]
theorem add_restrict' (T S : 𝒟'(X,Y)) (A : Set X) :
    T.restrict A + S.restrict A = (T + S).restrict A


import Mathlib.MeasureTheory.Measure.GiryMonad
import Mathlib.MeasureTheory.Decomposition.Lebesgue
import Mathlib.MeasureTheory.Constructions.Prod.Basic

import SciLean.Core.FunctionPropositions
import SciLean.Core.FunctionSpaces
import SciLean.Core.Integral.CIntegral
import SciLean.Core.Distribution.TestFunction
import SciLean.Core.Distribution.SimpAttr
import SciLean.Util.SorryProof
import SciLean.Util.Limit

open MeasureTheory ENNReal Classical

namespace SciLean

variable
  {R} [RealScalar R]
  {W} [Vec R W] [Module ℝ W]
  {X} [Vec R X] -- [TopologicalSpace X] [space : TCOr (Vec R X) (DiscreteTopology X)]
  {Y} [Vec R Y] [Module ℝ Y]
  {Z} [Vec R Z]
  {U} [Vec R U]
  {V} [Vec R V]

set_default_scalar R

variable (R X Y)
abbrev Distribution := (𝒟 X) ⊸[R] Y
variable {R X Y}

notation "𝒟'" X => Distribution defaultScalar% X defaultScalar%
notation "𝒟'" "(" X ", " Y ")" => Distribution defaultScalar% X Y

@[app_unexpander Distribution] def unexpandDistribution : Lean.PrettyPrinter.Unexpander
  | `($(_) $_ $X $Y) => `(𝒟'($X,$Y))
  | _ => throw ()

@[ext]
theorem Distribution.ext (x y : 𝒟'(X,Y)) :
    (∀ (φ : 𝒟 X), x φ = y φ)
    →
    x = y := by

  apply SmoothLinearMap.ext

----------------------------------------------------------------------------------------------------
-- Algebra -----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- instance : Zero (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Add (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Sub (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Neg (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : SMul R (𝒟'(X,Y)) := by unfold Distribution; infer_instance
instance [Module ℝ Y] : SMul ℝ (𝒟'(X,Y)) := ⟨fun r f => ⟨fun φ => r • (f φ), sorry_proof⟩⟩

-- instance : UniformSpace (𝒟'(X,Y)) :=  by unfold Distribution; infer_instance
-- instance : Vec R (𝒟'(X,Y)) := by unfold Distribution; infer_instance
instance [Module ℝ Y] : Module ℝ (𝒟'(X,Y)) := Module.mkSorryProofs

----------------------------------------------------------------------------------------------------
-- Extended action ---------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

open BigOperators in
@[pp_dot]
noncomputable
def Distribution.extAction (T : 𝒟'(X,Y)) (φ : X → Z) (L : Y ⊸ Z ⊸ W) : W :=
  if h : ∃ (zₙ : ℕ → Z) (φₙ : ℕ → 𝒟 X), ∀ x, ∑' i, φₙ i x • zₙ i = φ x then
    let zₙ := Classical.choose h
    let φₙ := (Classical.choose_spec h).choose
    ∑' i, L (T (φₙ i)) (zₙ i)
  else
    0

namespace Distribution
scoped notation "⟪" T ", " φ "⟫[" L "]" => Distribution.extAction T φ L
end Distribution

noncomputable
abbrev Distribution.extAction' (T : 𝒟'(X,Y)) (φ : X → R) : Y := T.extAction φ (fun y ⊸ fun r ⊸ r • y)

noncomputable
abbrev Distribution.integrate (T : 𝒟'(X,Y)) : Y := T.extAction' (fun _ => 1)

@[fun_prop]
theorem TestFunction.apply_IsSmoothLinearMap : IsSmoothLinearMap R fun (φ : 𝒟 X) => (φ : X → R) := sorry_proof

theorem Distribution.mk_extAction (T : (X → R) → Y) (hT : IsSmoothLinearMap R (fun φ : 𝒟 X => T φ)) (φ : X → R) :
    Distribution.extAction (SmoothLinearMap.mk' R (fun (φ : 𝒟 X) => T φ) hT : Distribution _ _ _) φ (fun y ⊸ fun r ⊸ r • y) = T φ := sorry_proof

-- This is definitely not true as stated, what kind of condistions do we need on `φ` and `T`?
@[fun_prop]
theorem Distribution.extAction.arg_φ.IsSmoothLinearMap (T : 𝒟'(X,U)) (φ : W → X → V) (L : U ⊸ V ⊸ Z)
    (hφ : IsSmoothLinearMap R φ) :
    IsSmoothLinearMap R (fun w => T.extAction (φ w) L) := sorry_proof

@[fun_prop]
theorem Distribution.extAction.arg_T.IsSmoothLinearMap (T : W → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ Z)
    (hT : IsSmoothLinearMap R T) :
    IsSmoothLinearMap R (fun w => (T w).extAction φ L) := sorry_proof

-- open Lean Meta in
-- /-- Simproc to apply `Distribution.mk_extAction` theorem -/
-- simproc_decl Distribution.mk_extAction_simproc (Distribution.extAction (Distribution.mk (SmoothLinearMap.mk _ _)) _) := fun e => do

--   let φ := e.appArg!
--   let T := e.appFn!.appArg!

--   let .lam xName xType xBody xBi := T.appArg!.appFn!.appArg!
--     | return .continue
--   let hT := T.appArg!.appArg!

--   withLocalDecl xName xBi xType fun x => do
--   let R := xType.getArg! 0
--   let X := xType.getArg! 2
--   withLocalDecl `φ' xBi (← mkArrow X R) fun φ' => do
--     let b := xBody.instantiate1 x
--     let b := b.replace (fun e' =>
--       if e'.isAppOf ``DFunLike.coe &&
--          5 ≤ e'.getAppNumArgs &&
--          e'.getArg! 4 == x then
--           .some (mkAppN φ' e'.getAppArgs[5:])
--       else
--         .none)

--     if b.containsFVar x.fvarId! then
--       return .continue

--     let T ← mkLambdaFVars #[φ'] b
--     let prf ← mkAppM ``Distribution.mk_extAction #[T, hT, φ]
--     return .visit {expr := T.beta #[φ], proof? := prf}

----------------------------------------------------------------------------------------------------
-- Monadic structure -------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- def dirac (x : X) : Distribution X := fun φ => φ x

-- instance : Monad (Distribution R) where
--   pure := fun x => ⟨fun φ => φ x⟩
--   bind := fun x f => ⟨fun φ => ⟪x, fun x' => ⟪(f x'), φ⟫⟫⟩

-- instance : LawfulMonad (Distribution R) where
--   bind_pure_comp := by intros; rfl
--   bind_map       := by intros; rfl
--   pure_bind      := by intros; rfl
--   bind_assoc     := by intros; rfl
--   map_const      := by intros; rfl
--   id_map         := by intros; rfl
--   seqLeft_eq     := by intros; rfl
--   seqRight_eq    := by intros; rfl
--   pure_seq       := by intros; rfl

def dirac (x : X) : 𝒟' X := fun φ ⊸ φ x

open Notation
noncomputable
def Distribution.bind (x' : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) : 𝒟'(Y,W) :=
  fun φ ⊸ x'.extAction (fun x => (f x).extAction φ (fun v ⊸ fun r ⊸ r • v)) L

----------------------------------------------------------------------------------------------------
-- Basic identities --------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[simp, ftrans_simp]
theorem action_dirac (x : X) (φ : 𝒟 X) : dirac x φ = φ x := by simp[dirac]

@[simp, ftrans_simp]
theorem action_bind (x : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) (φ : 𝒟 Y) :
    x.bind f L φ = x.extAction (fun x' => (f x').extAction' φ) L := by
  simp[Distribution.bind]

-- @[simp, ftrans_simp]
-- theorem extAction_bind (x : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) (φ : Y → Z) (K : W ⊸ Z ⊸ W') :
--     (x.bind f L).extAction φ K = x.extAction (fun x' => (f x').extAction φ (sorry : V ⊸ Z ⊸ V⊗Z)) (sorry : U ⊸ (V⊗Z) ⊸ W') := by
--   simp [Distribution.bind]

----------------------------------------------------------------------------------------------------
-- Arithmetics -------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

section Arithmetics

@[simp, ftrans_simp, action_push]
theorem Distribution.zero_extAction (φ : X → V) (L : U ⊸ V ⊸ W) : (0 : 𝒟'(X,U)).extAction φ L = 0 := by
  unfold extAction; simp

-- todo: this needs some integrability condition
@[action_push]
theorem Distribution.add_extAction (T T' : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    ((T + T') : 𝒟'(X,U)).extAction φ L = T.extAction φ L + T'.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.sub_extAction (T T' : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (T - T').extAction φ L = T.extAction φ L - T'.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.smul_extAction (r : R) (T : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W)  :
    (r • T).extAction φ L = r • T.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.neg_extAction (T : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (- T).extAction φ L = - T.extAction φ L := by sorry_proof

open BigOperators in
@[action_push]
theorem Distribution.fintype_sum_extAction {I} [Fintype I] (T : I → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (∑ i, T i).extAction φ L = ∑ i, (T i).extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.indextype_sum_extAction {I} [IndexType I] (T : I → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (∑ i, T i).extAction φ L = ∑ i, (T i).extAction φ L := by sorry_proof

end Arithmetics

----------------------------------------------------------------------------------------------------
-- Distributional if statement ---------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable [MeasureSpace X]

open Classical Notation in
noncomputable
def iteD (A : Set X) (t e : 𝒟'(X,Y)) : 𝒟'(X,Y) :=
  fun φ ⊸
    t.extAction (fun x => if x ∈ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) +
    e.extAction (fun x => if x ∈ A then 0 else φ x) (fun y ⊸ fun r ⊸ r • y)

open Lean.Parser Term in
syntax withPosition("ifD " term " then "
    ppDedent(ppLine ppSpace ppSpace) term ppDedent(ppLine)
  "else"
    ppDedent(ppLine ppSpace ppSpace) term) : term

macro_rules
  | `(ifD $A then $t else $e) => `(iteD $A $t $e)

open Lean Parser in
@[app_unexpander iteD]
def unexpandIteD : Lean.PrettyPrinter.Unexpander
  | `($(_) $A $t $e) => `(ifD $A then $t else $e)
  | _ => throw ()

@[action_push]
theorem Distribution.action_iteD (A : Set X) (t e : 𝒟'(X,Y)) (φ : 𝒟 X) :
   iteD A t e φ =
        t.extAction (fun x => if x ∈ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) +
        e.extAction (fun x => if x ∉ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) := by sorry_proof

@[simp, ftrans_simp]
theorem Distribution.iteD_same (A : Set X) (u : 𝒟'(X,Y)) :
   iteD A u u = u := by sorry_proof

@[action_push]
theorem Distribution.extAction_iteD (A : Set X) (t e : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (iteD A t e).extAction φ L =
        t.extAction (fun x => if x ∈ A then φ x else 0) L +
        e.extAction (fun x => if x ∉ A then φ x else 0) L := by sorry_proof

@[fun_prop]
theorem iteD.arg_te.IsSmoothLinearMap_rule (A : Set X) (t e : W → 𝒟'(X,Y))
    (ht : IsSmoothLinearMap R t) (he : IsSmoothLinearMap R e) :
    IsSmoothLinearMap R (fun w => iteD A (t w) (e w)) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Set restriction ---------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[pp_dot]
noncomputable
def Distribution.restrict (T : 𝒟'(X,Y)) (A : Set X) : 𝒟'(X,Y) :=
  ifD A then T else 0

@[simp,ftrans_simp]
theorem restrict_univ (T : 𝒟'(X,Y))  :
    T.restrict Set.univ = T := sorry_proof

@[simp,ftrans_simp]
theorem zero_restrict (A : Set X) :
    (0 : 𝒟'(X,Y)).restrict A = 0 := sorry_proof

@[restrict_push]
theorem add_restrict (T S : 𝒟'(X,Y)) (A : Set X) :
    (T + S).restrict A = T.restrict A + S.restrict A := sorry_proof

@[restrict_pull]
theorem add_restrict' (T S : 𝒟'(X,Y)) (A : Set X) :
    T.restrict A + S.restrict A = (T + S).restrict A := sorry_proof

@[restrict_push]
theorem sub_restrict (T S : 𝒟'(X,Y)) (A : Set X) :
    (T - S).restrict A = T.restrict A - S.restrict A := sorry_proof

@[restrict_pull]
theorem sub_restrict' (T S : 𝒟'(X,Y)) (A : Set X) :
    T.restrict A - S.restrict A = (T - S).restrict A


import Mathlib.MeasureTheory.Measure.GiryMonad
import Mathlib.MeasureTheory.Decomposition.Lebesgue
import Mathlib.MeasureTheory.Constructions.Prod.Basic

import SciLean.Core.FunctionPropositions
import SciLean.Core.FunctionSpaces
import SciLean.Core.Integral.CIntegral
import SciLean.Core.Distribution.TestFunction
import SciLean.Core.Distribution.SimpAttr
import SciLean.Util.SorryProof
import SciLean.Util.Limit

open MeasureTheory ENNReal Classical

namespace SciLean

variable
  {R} [RealScalar R]
  {W} [Vec R W] [Module ℝ W]
  {X} [Vec R X] -- [TopologicalSpace X] [space : TCOr (Vec R X) (DiscreteTopology X)]
  {Y} [Vec R Y] [Module ℝ Y]
  {Z} [Vec R Z]
  {U} [Vec R U]
  {V} [Vec R V]

set_default_scalar R

variable (R X Y)
abbrev Distribution := (𝒟 X) ⊸[R] Y
variable {R X Y}

notation "𝒟'" X => Distribution defaultScalar% X defaultScalar%
notation "𝒟'" "(" X ", " Y ")" => Distribution defaultScalar% X Y

@[app_unexpander Distribution] def unexpandDistribution : Lean.PrettyPrinter.Unexpander
  | `($(_) $_ $X $Y) => `(𝒟'($X,$Y))
  | _ => throw ()

@[ext]
theorem Distribution.ext (x y : 𝒟'(X,Y)) :
    (∀ (φ : 𝒟 X), x φ = y φ)
    →
    x = y := by

  apply SmoothLinearMap.ext

----------------------------------------------------------------------------------------------------
-- Algebra -----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- instance : Zero (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Add (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Sub (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Neg (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : SMul R (𝒟'(X,Y)) := by unfold Distribution; infer_instance
instance [Module ℝ Y] : SMul ℝ (𝒟'(X,Y)) := ⟨fun r f => ⟨fun φ => r • (f φ), sorry_proof⟩⟩

-- instance : UniformSpace (𝒟'(X,Y)) :=  by unfold Distribution; infer_instance
-- instance : Vec R (𝒟'(X,Y)) := by unfold Distribution; infer_instance
instance [Module ℝ Y] : Module ℝ (𝒟'(X,Y)) := Module.mkSorryProofs

----------------------------------------------------------------------------------------------------
-- Extended action ---------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

open BigOperators in
@[pp_dot]
noncomputable
def Distribution.extAction (T : 𝒟'(X,Y)) (φ : X → Z) (L : Y ⊸ Z ⊸ W) : W :=
  if h : ∃ (zₙ : ℕ → Z) (φₙ : ℕ → 𝒟 X), ∀ x, ∑' i, φₙ i x • zₙ i = φ x then
    let zₙ := Classical.choose h
    let φₙ := (Classical.choose_spec h).choose
    ∑' i, L (T (φₙ i)) (zₙ i)
  else
    0

namespace Distribution
scoped notation "⟪" T ", " φ "⟫[" L "]" => Distribution.extAction T φ L
end Distribution

noncomputable
abbrev Distribution.extAction' (T : 𝒟'(X,Y)) (φ : X → R) : Y := T.extAction φ (fun y ⊸ fun r ⊸ r • y)

noncomputable
abbrev Distribution.integrate (T : 𝒟'(X,Y)) : Y := T.extAction' (fun _ => 1)

@[fun_prop]
theorem TestFunction.apply_IsSmoothLinearMap : IsSmoothLinearMap R fun (φ : 𝒟 X) => (φ : X → R) := sorry_proof

theorem Distribution.mk_extAction (T : (X → R) → Y) (hT : IsSmoothLinearMap R (fun φ : 𝒟 X => T φ)) (φ : X → R) :
    Distribution.extAction (SmoothLinearMap.mk' R (fun (φ : 𝒟 X) => T φ) hT : Distribution _ _ _) φ (fun y ⊸ fun r ⊸ r • y) = T φ := sorry_proof

-- This is definitely not true as stated, what kind of condistions do we need on `φ` and `T`?
@[fun_prop]
theorem Distribution.extAction.arg_φ.IsSmoothLinearMap (T : 𝒟'(X,U)) (φ : W → X → V) (L : U ⊸ V ⊸ Z)
    (hφ : IsSmoothLinearMap R φ) :
    IsSmoothLinearMap R (fun w => T.extAction (φ w) L) := sorry_proof

@[fun_prop]
theorem Distribution.extAction.arg_T.IsSmoothLinearMap (T : W → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ Z)
    (hT : IsSmoothLinearMap R T) :
    IsSmoothLinearMap R (fun w => (T w).extAction φ L) := sorry_proof

-- open Lean Meta in
-- /-- Simproc to apply `Distribution.mk_extAction` theorem -/
-- simproc_decl Distribution.mk_extAction_simproc (Distribution.extAction (Distribution.mk (SmoothLinearMap.mk _ _)) _) := fun e => do

--   let φ := e.appArg!
--   let T := e.appFn!.appArg!

--   let .lam xName xType xBody xBi := T.appArg!.appFn!.appArg!
--     | return .continue
--   let hT := T.appArg!.appArg!

--   withLocalDecl xName xBi xType fun x => do
--   let R := xType.getArg! 0
--   let X := xType.getArg! 2
--   withLocalDecl `φ' xBi (← mkArrow X R) fun φ' => do
--     let b := xBody.instantiate1 x
--     let b := b.replace (fun e' =>
--       if e'.isAppOf ``DFunLike.coe &&
--          5 ≤ e'.getAppNumArgs &&
--          e'.getArg! 4 == x then
--           .some (mkAppN φ' e'.getAppArgs[5:])
--       else
--         .none)

--     if b.containsFVar x.fvarId! then
--       return .continue

--     let T ← mkLambdaFVars #[φ'] b
--     let prf ← mkAppM ``Distribution.mk_extAction #[T, hT, φ]
--     return .visit {expr := T.beta #[φ], proof? := prf}

----------------------------------------------------------------------------------------------------
-- Monadic structure -------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- def dirac (x : X) : Distribution X := fun φ => φ x

-- instance : Monad (Distribution R) where
--   pure := fun x => ⟨fun φ => φ x⟩
--   bind := fun x f => ⟨fun φ => ⟪x, fun x' => ⟪(f x'), φ⟫⟫⟩

-- instance : LawfulMonad (Distribution R) where
--   bind_pure_comp := by intros; rfl
--   bind_map       := by intros; rfl
--   pure_bind      := by intros; rfl
--   bind_assoc     := by intros; rfl
--   map_const      := by intros; rfl
--   id_map         := by intros; rfl
--   seqLeft_eq     := by intros; rfl
--   seqRight_eq    := by intros; rfl
--   pure_seq       := by intros; rfl

def dirac (x : X) : 𝒟' X := fun φ ⊸ φ x

open Notation
noncomputable
def Distribution.bind (x' : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) : 𝒟'(Y,W) :=
  fun φ ⊸ x'.extAction (fun x => (f x).extAction φ (fun v ⊸ fun r ⊸ r • v)) L

----------------------------------------------------------------------------------------------------
-- Basic identities --------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[simp, ftrans_simp]
theorem action_dirac (x : X) (φ : 𝒟 X) : dirac x φ = φ x := by simp[dirac]

@[simp, ftrans_simp]
theorem action_bind (x : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) (φ : 𝒟 Y) :
    x.bind f L φ = x.extAction (fun x' => (f x').extAction' φ) L := by
  simp[Distribution.bind]

-- @[simp, ftrans_simp]
-- theorem extAction_bind (x : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) (φ : Y → Z) (K : W ⊸ Z ⊸ W') :
--     (x.bind f L).extAction φ K = x.extAction (fun x' => (f x').extAction φ (sorry : V ⊸ Z ⊸ V⊗Z)) (sorry : U ⊸ (V⊗Z) ⊸ W') := by
--   simp [Distribution.bind]

----------------------------------------------------------------------------------------------------
-- Arithmetics -------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

section Arithmetics

@[simp, ftrans_simp, action_push]
theorem Distribution.zero_extAction (φ : X → V) (L : U ⊸ V ⊸ W) : (0 : 𝒟'(X,U)).extAction φ L = 0 := by
  unfold extAction; simp

-- todo: this needs some integrability condition
@[action_push]
theorem Distribution.add_extAction (T T' : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    ((T + T') : 𝒟'(X,U)).extAction φ L = T.extAction φ L + T'.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.sub_extAction (T T' : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (T - T').extAction φ L = T.extAction φ L - T'.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.smul_extAction (r : R) (T : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W)  :
    (r • T).extAction φ L = r • T.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.neg_extAction (T : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (- T).extAction φ L = - T.extAction φ L := by sorry_proof

open BigOperators in
@[action_push]
theorem Distribution.fintype_sum_extAction {I} [Fintype I] (T : I → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (∑ i, T i).extAction φ L = ∑ i, (T i).extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.indextype_sum_extAction {I} [IndexType I] (T : I → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (∑ i, T i).extAction φ L = ∑ i, (T i).extAction φ L := by sorry_proof

end Arithmetics

----------------------------------------------------------------------------------------------------
-- Distributional if statement ---------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable [MeasureSpace X]

open Classical Notation in
noncomputable
def iteD (A : Set X) (t e : 𝒟'(X,Y)) : 𝒟'(X,Y) :=
  fun φ ⊸
    t.extAction (fun x => if x ∈ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) +
    e.extAction (fun x => if x ∈ A then 0 else φ x) (fun y ⊸ fun r ⊸ r • y)

open Lean.Parser Term in
syntax withPosition("ifD " term " then "
    ppDedent(ppLine ppSpace ppSpace) term ppDedent(ppLine)
  "else"
    ppDedent(ppLine ppSpace ppSpace) term) : term

macro_rules
  | `(ifD $A then $t else $e) => `(iteD $A $t $e)

open Lean Parser in
@[app_unexpander iteD]
def unexpandIteD : Lean.PrettyPrinter.Unexpander
  | `($(_) $A $t $e) => `(ifD $A then $t else $e)
  | _ => throw ()

@[action_push]
theorem Distribution.action_iteD (A : Set X) (t e : 𝒟'(X,Y)) (φ : 𝒟 X) :
   iteD A t e φ =
        t.extAction (fun x => if x ∈ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) +
        e.extAction (fun x => if x ∉ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) := by sorry_proof

@[simp, ftrans_simp]
theorem Distribution.iteD_same (A : Set X) (u : 𝒟'(X,Y)) :
   iteD A u u = u := by sorry_proof

@[action_push]
theorem Distribution.extAction_iteD (A : Set X) (t e : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (iteD A t e).extAction φ L =
        t.extAction (fun x => if x ∈ A then φ x else 0) L +
        e.extAction (fun x => if x ∉ A then φ x else 0) L := by sorry_proof

@[fun_prop]
theorem iteD.arg_te.IsSmoothLinearMap_rule (A : Set X) (t e : W → 𝒟'(X,Y))
    (ht : IsSmoothLinearMap R t) (he : IsSmoothLinearMap R e) :
    IsSmoothLinearMap R (fun w => iteD A (t w) (e w)) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Set restriction ---------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[pp_dot]
noncomputable
def Distribution.restrict (T : 𝒟'(X,Y)) (A : Set X) : 𝒟'(X,Y) :=
  ifD A then T else 0

@[simp,ftrans_simp]
theorem restrict_univ (T : 𝒟'(X,Y))  :
    T.restrict Set.univ = T := sorry_proof

@[simp,ftrans_simp]
theorem zero_restrict (A : Set X) :
    (0 : 𝒟'(X,Y)).restrict A = 0 := sorry_proof

@[restrict_push]
theorem add_restrict (T S : 𝒟'(X,Y)) (A : Set X) :
    (T + S).restrict A = T.restrict A + S.restrict A := sorry_proof

@[restrict_pull]
theorem add_restrict' (T S : 𝒟'(X,Y)) (A : Set X) :
    T.restrict A + S.restrict A = (T + S).restrict A := sorry_proof

@[restrict_push]
theorem sub_restrict (T S : 𝒟'(X,Y)) (A : Set X) :
    (T - S).restrict A = T.restrict A - S.restrict A := sorry_proof

@[restrict_pull]
theorem sub_restrict' (T S : 𝒟'(X,Y)) (A : Set X) :
    T.restrict A - S.restrict A = (T - S).restrict A := sorry_proof

@[restrict_push]
theorem smul_restrict (r : R) (T : 𝒟'(X,Y)) (A : Set X) :
    (r • T).restrict A = r • (T.restrict A) := sorry_proof

@[restrict_pull]
theorem smul_restrict' (r : R) (T : 𝒟'(X,Y)) (A : Set X) :
    r • (T.restrict A) = (r • T).restrict A


import Mathlib.MeasureTheory.Measure.GiryMonad
import Mathlib.MeasureTheory.Decomposition.Lebesgue
import Mathlib.MeasureTheory.Constructions.Prod.Basic

import SciLean.Core.FunctionPropositions
import SciLean.Core.FunctionSpaces
import SciLean.Core.Integral.CIntegral
import SciLean.Core.Distribution.TestFunction
import SciLean.Core.Distribution.SimpAttr
import SciLean.Util.SorryProof
import SciLean.Util.Limit

open MeasureTheory ENNReal Classical

namespace SciLean

variable
  {R} [RealScalar R]
  {W} [Vec R W] [Module ℝ W]
  {X} [Vec R X] -- [TopologicalSpace X] [space : TCOr (Vec R X) (DiscreteTopology X)]
  {Y} [Vec R Y] [Module ℝ Y]
  {Z} [Vec R Z]
  {U} [Vec R U]
  {V} [Vec R V]

set_default_scalar R

variable (R X Y)
abbrev Distribution := (𝒟 X) ⊸[R] Y
variable {R X Y}

notation "𝒟'" X => Distribution defaultScalar% X defaultScalar%
notation "𝒟'" "(" X ", " Y ")" => Distribution defaultScalar% X Y

@[app_unexpander Distribution] def unexpandDistribution : Lean.PrettyPrinter.Unexpander
  | `($(_) $_ $X $Y) => `(𝒟'($X,$Y))
  | _ => throw ()

@[ext]
theorem Distribution.ext (x y : 𝒟'(X,Y)) :
    (∀ (φ : 𝒟 X), x φ = y φ)
    →
    x = y := by

  apply SmoothLinearMap.ext

----------------------------------------------------------------------------------------------------
-- Algebra -----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- instance : Zero (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Add (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Sub (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : Neg (𝒟'(X,Y)) := by unfold Distribution; infer_instance
-- instance : SMul R (𝒟'(X,Y)) := by unfold Distribution; infer_instance
instance [Module ℝ Y] : SMul ℝ (𝒟'(X,Y)) := ⟨fun r f => ⟨fun φ => r • (f φ), sorry_proof⟩⟩

-- instance : UniformSpace (𝒟'(X,Y)) :=  by unfold Distribution; infer_instance
-- instance : Vec R (𝒟'(X,Y)) := by unfold Distribution; infer_instance
instance [Module ℝ Y] : Module ℝ (𝒟'(X,Y)) := Module.mkSorryProofs

----------------------------------------------------------------------------------------------------
-- Extended action ---------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

open BigOperators in
@[pp_dot]
noncomputable
def Distribution.extAction (T : 𝒟'(X,Y)) (φ : X → Z) (L : Y ⊸ Z ⊸ W) : W :=
  if h : ∃ (zₙ : ℕ → Z) (φₙ : ℕ → 𝒟 X), ∀ x, ∑' i, φₙ i x • zₙ i = φ x then
    let zₙ := Classical.choose h
    let φₙ := (Classical.choose_spec h).choose
    ∑' i, L (T (φₙ i)) (zₙ i)
  else
    0

namespace Distribution
scoped notation "⟪" T ", " φ "⟫[" L "]" => Distribution.extAction T φ L
end Distribution

noncomputable
abbrev Distribution.extAction' (T : 𝒟'(X,Y)) (φ : X → R) : Y := T.extAction φ (fun y ⊸ fun r ⊸ r • y)

noncomputable
abbrev Distribution.integrate (T : 𝒟'(X,Y)) : Y := T.extAction' (fun _ => 1)

@[fun_prop]
theorem TestFunction.apply_IsSmoothLinearMap : IsSmoothLinearMap R fun (φ : 𝒟 X) => (φ : X → R) := sorry_proof

theorem Distribution.mk_extAction (T : (X → R) → Y) (hT : IsSmoothLinearMap R (fun φ : 𝒟 X => T φ)) (φ : X → R) :
    Distribution.extAction (SmoothLinearMap.mk' R (fun (φ : 𝒟 X) => T φ) hT : Distribution _ _ _) φ (fun y ⊸ fun r ⊸ r • y) = T φ := sorry_proof

-- This is definitely not true as stated, what kind of condistions do we need on `φ` and `T`?
@[fun_prop]
theorem Distribution.extAction.arg_φ.IsSmoothLinearMap (T : 𝒟'(X,U)) (φ : W → X → V) (L : U ⊸ V ⊸ Z)
    (hφ : IsSmoothLinearMap R φ) :
    IsSmoothLinearMap R (fun w => T.extAction (φ w) L) := sorry_proof

@[fun_prop]
theorem Distribution.extAction.arg_T.IsSmoothLinearMap (T : W → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ Z)
    (hT : IsSmoothLinearMap R T) :
    IsSmoothLinearMap R (fun w => (T w).extAction φ L) := sorry_proof

-- open Lean Meta in
-- /-- Simproc to apply `Distribution.mk_extAction` theorem -/
-- simproc_decl Distribution.mk_extAction_simproc (Distribution.extAction (Distribution.mk (SmoothLinearMap.mk _ _)) _) := fun e => do

--   let φ := e.appArg!
--   let T := e.appFn!.appArg!

--   let .lam xName xType xBody xBi := T.appArg!.appFn!.appArg!
--     | return .continue
--   let hT := T.appArg!.appArg!

--   withLocalDecl xName xBi xType fun x => do
--   let R := xType.getArg! 0
--   let X := xType.getArg! 2
--   withLocalDecl `φ' xBi (← mkArrow X R) fun φ' => do
--     let b := xBody.instantiate1 x
--     let b := b.replace (fun e' =>
--       if e'.isAppOf ``DFunLike.coe &&
--          5 ≤ e'.getAppNumArgs &&
--          e'.getArg! 4 == x then
--           .some (mkAppN φ' e'.getAppArgs[5:])
--       else
--         .none)

--     if b.containsFVar x.fvarId! then
--       return .continue

--     let T ← mkLambdaFVars #[φ'] b
--     let prf ← mkAppM ``Distribution.mk_extAction #[T, hT, φ]
--     return .visit {expr := T.beta #[φ], proof? := prf}

----------------------------------------------------------------------------------------------------
-- Monadic structure -------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- def dirac (x : X) : Distribution X := fun φ => φ x

-- instance : Monad (Distribution R) where
--   pure := fun x => ⟨fun φ => φ x⟩
--   bind := fun x f => ⟨fun φ => ⟪x, fun x' => ⟪(f x'), φ⟫⟫⟩

-- instance : LawfulMonad (Distribution R) where
--   bind_pure_comp := by intros; rfl
--   bind_map       := by intros; rfl
--   pure_bind      := by intros; rfl
--   bind_assoc     := by intros; rfl
--   map_const      := by intros; rfl
--   id_map         := by intros; rfl
--   seqLeft_eq     := by intros; rfl
--   seqRight_eq    := by intros; rfl
--   pure_seq       := by intros; rfl

def dirac (x : X) : 𝒟' X := fun φ ⊸ φ x

open Notation
noncomputable
def Distribution.bind (x' : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) : 𝒟'(Y,W) :=
  fun φ ⊸ x'.extAction (fun x => (f x).extAction φ (fun v ⊸ fun r ⊸ r • v)) L

----------------------------------------------------------------------------------------------------
-- Basic identities --------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[simp, ftrans_simp]
theorem action_dirac (x : X) (φ : 𝒟 X) : dirac x φ = φ x := by simp[dirac]

@[simp, ftrans_simp]
theorem action_bind (x : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) (φ : 𝒟 Y) :
    x.bind f L φ = x.extAction (fun x' => (f x').extAction' φ) L := by
  simp[Distribution.bind]

-- @[simp, ftrans_simp]
-- theorem extAction_bind (x : 𝒟'(X,U)) (f : X → 𝒟'(Y,V)) (L : U ⊸ V ⊸ W) (φ : Y → Z) (K : W ⊸ Z ⊸ W') :
--     (x.bind f L).extAction φ K = x.extAction (fun x' => (f x').extAction φ (sorry : V ⊸ Z ⊸ V⊗Z)) (sorry : U ⊸ (V⊗Z) ⊸ W') := by
--   simp [Distribution.bind]

----------------------------------------------------------------------------------------------------
-- Arithmetics -------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

section Arithmetics

@[simp, ftrans_simp, action_push]
theorem Distribution.zero_extAction (φ : X → V) (L : U ⊸ V ⊸ W) : (0 : 𝒟'(X,U)).extAction φ L = 0 := by
  unfold extAction; simp

-- todo: this needs some integrability condition
@[action_push]
theorem Distribution.add_extAction (T T' : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    ((T + T') : 𝒟'(X,U)).extAction φ L = T.extAction φ L + T'.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.sub_extAction (T T' : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (T - T').extAction φ L = T.extAction φ L - T'.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.smul_extAction (r : R) (T : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W)  :
    (r • T).extAction φ L = r • T.extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.neg_extAction (T : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (- T).extAction φ L = - T.extAction φ L := by sorry_proof

open BigOperators in
@[action_push]
theorem Distribution.fintype_sum_extAction {I} [Fintype I] (T : I → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (∑ i, T i).extAction φ L = ∑ i, (T i).extAction φ L := by sorry_proof

@[action_push]
theorem Distribution.indextype_sum_extAction {I} [IndexType I] (T : I → 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (∑ i, T i).extAction φ L = ∑ i, (T i).extAction φ L := by sorry_proof

end Arithmetics

----------------------------------------------------------------------------------------------------
-- Distributional if statement ---------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable [MeasureSpace X]

open Classical Notation in
noncomputable
def iteD (A : Set X) (t e : 𝒟'(X,Y)) : 𝒟'(X,Y) :=
  fun φ ⊸
    t.extAction (fun x => if x ∈ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) +
    e.extAction (fun x => if x ∈ A then 0 else φ x) (fun y ⊸ fun r ⊸ r • y)

open Lean.Parser Term in
syntax withPosition("ifD " term " then "
    ppDedent(ppLine ppSpace ppSpace) term ppDedent(ppLine)
  "else"
    ppDedent(ppLine ppSpace ppSpace) term) : term

macro_rules
  | `(ifD $A then $t else $e) => `(iteD $A $t $e)

open Lean Parser in
@[app_unexpander iteD]
def unexpandIteD : Lean.PrettyPrinter.Unexpander
  | `($(_) $A $t $e) => `(ifD $A then $t else $e)
  | _ => throw ()

@[action_push]
theorem Distribution.action_iteD (A : Set X) (t e : 𝒟'(X,Y)) (φ : 𝒟 X) :
   iteD A t e φ =
        t.extAction (fun x => if x ∈ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) +
        e.extAction (fun x => if x ∉ A then φ x else 0) (fun y ⊸ fun r ⊸ r • y) := by sorry_proof

@[simp, ftrans_simp]
theorem Distribution.iteD_same (A : Set X) (u : 𝒟'(X,Y)) :
   iteD A u u = u := by sorry_proof

@[action_push]
theorem Distribution.extAction_iteD (A : Set X) (t e : 𝒟'(X,U)) (φ : X → V) (L : U ⊸ V ⊸ W) :
    (iteD A t e).extAction φ L =
        t.extAction (fun x => if x ∈ A then φ x else 0) L +
        e.extAction (fun x => if x ∉ A then φ x else 0) L := by sorry_proof

@[fun_prop]
theorem iteD.arg_te.IsSmoothLinearMap_rule (A : Set X) (t e : W → 𝒟'(X,Y))
    (ht : IsSmoothLinearMap R t) (he : IsSmoothLinearMap R e) :
    IsSmoothLinearMap R (fun w => iteD A (t w) (e w)) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Set restriction ---------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[pp_dot]
noncomputable
def Distribution.restrict (T : 𝒟'(X,Y)) (A : Set X) : 𝒟'(X,Y) :=
  ifD A then T else 0

@[simp,ftrans_simp]
theorem restrict_univ (T : 𝒟'(X,Y))  :
    T.restrict Set.univ = T := sorry_proof

@[simp,ftrans_simp]
theorem zero_restrict (A : Set X) :
    (0 : 𝒟'(X,Y)).restrict A = 0 := sorry_proof

@[restrict_push]
theorem add_restrict (T S : 𝒟'(X,Y)) (A : Set X) :
    (T + S).restrict A = T.restrict A + S.restrict A := sorry_proof

@[restrict_pull]
theorem add_restrict' (T S : 𝒟'(X,Y)) (A : Set X) :
    T.restrict A + S.restrict A = (T + S).restrict A := sorry_proof

@[restrict_push]
theorem sub_restrict (T S : 𝒟'(X,Y)) (A : Set X) :
    (T - S).restrict A = T.restrict A - S.restrict A := sorry_proof

@[restrict_pull]
theorem sub_restrict' (T S : 𝒟'(X,Y)) (A : Set X) :
    T.restrict A - S.restrict A = (T - S).restrict A := sorry_proof

@[restrict_push]
theorem smul_restrict (r : R) (T : 𝒟'(X,Y)) (A : Set X) :
    (r • T).restrict A = r • (T.restrict A) := sorry_proof

@[restrict_pull]
theorem smul_restrict' (r : R) (T : 𝒟'(X,Y)) (A : Set X) :
    r • (T.restrict A) = (r • T).restrict A := sorry_proof

@[restrict_push]
theorem neg_restrict (T : 𝒟'(X,Y)) (A : Set X) :
    (- T).restrict A = - (T.restrict A) := sorry_proof

@[restrict_pull]
theorem neg_restrict' (T : 𝒟'(X,Y)) (A : Set X) :
    - (T.restrict A) = (- T).restrict A