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 := by sorry_proof

----------------------------------------------------------------------------------------------------
-- Distributiona product  --------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable {X₁} [Vec R X₁] {X₂} [Vec R X₂] {Y₁} [Vec R Y₁] {Y₂} [Vec R Y₂]

-- can we extended to vector valued distributions?
noncomputable
def Distribution.prod (p : X₁ → X₂ → X) (T : 𝒟' (X₁,Y₁)) (S : X₁ → 𝒟'(X₂,Y₂)) (L : Y₁ ⊸ Y₂ ⊸ Z) : 𝒟'(X,Z) :=
 ⟨fun φ => T.extAction (fun x₁ => S x₁ ⟨fun x₂ => φ (p x₁ x₂), sorry_proof⟩) L, sorry_proof⟩

@[simp, ftrans_simp]
theorem Distribution.prod_restrict (p : X₁ → X₂ → X) (T : 𝒟'(X₁,Y₁)) (S : X₁ → 𝒟'(X₂,Y₂)) (L : Y₁ ⊸ Y₂ ⊸ Z) (A : Set X) :
    (prod p T S L).restrict A = prod p (T.restrict (A.preimage1 p)) (fun x₁ => (S x₁).restrict (p x₁ ⁻¹' A)) L := sorry_proof

@[action_push]
theorem Distribution.prod'_extAction (p : X₁ → X₂ → X) (T : 𝒟'(X₁,Y₁)) (S : X₁ → 𝒟'(X₂,Y₂)) (L : Y₁ ⊸ Y₂ ⊸ Z) (K : Z ⊸ R ⊸ Z) (φ : X → R) :
    (prod p T S L).extAction φ K
    =
    T.extAction (fun x₁ => (S x₁).extAction (fun x₂ => φ (p x₁ x₂)) (fun y₂ ⊸ fun r ⊸ r • y₂)) (fun y₁ ⊸ fun y₂ ⊸ K (L y₁ y₂) 1) := sorry_proof

-- @[action_push]
-- theorem Distribution.prod'_extAction' (p : X₁ → X₂ → X) (T : 𝒟'(X₁,Y₁)) (S : X₁ → 𝒟'(X₂,Y₂)) (L : Y₁ ⊸ Y₂ ⊸ U) (φ : X → V) (K : U ⊸ V ⊸ W) :
--     (prod p T S L).extAction φ K
--     =
--     T.extAction (fun x₁ => (S x₁).extAction (fun x₂ => φ (p x₁ x₂)) (sorry : Y₂ ⊸ V ⊸ Y₂⊗V)) (fun y₁ ⊸ fun yv ⊸ ) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Post Composition --------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

noncomputable
def Distribution.postComp (T : 𝒟'(X,Y)) (f : Y ⊸ Z) : 𝒟'(X,Z) := fun φ ⊸ f (T φ)

-- @[pp_dot]
-- noncomputable
-- abbrev Distribution.postExtAction (T : 𝒟'(X,𝒟'(Y,U))) (φ : Y → V) (L : U ⊸ V ⊸ W) : 𝒟'(X,W) :=
--   T.postComp (fun u ⊸ u.extAction φ L)

@[fun_prop]
theorem Distribution.postComp.arg_T.IsSmoothLinarMap_rule (T : W → 𝒟'(X,Y)) (f : Y ⊸ Z)
    (hT : IsSmoothLinearMap R T) :
    IsSmoothLinearMap R (fun w => (T w).postComp f) := by unfold postComp; sorry_proof

@[simp, ftrans_simp]
theorem postComp_id (u : 𝒟'(X,Y)) :
    (u.postComp (fun y ⊸ y)) = u := sorry_proof

@[simp, ftrans_simp]
theorem postComp_comp (x : 𝒟'(X,U)) (g : U ⊸ V) (f : V ⊸ W) :
    (x.postComp g).postComp f
    =
    x.postComp (fun u ⊸ f (g u)) := sorry_proof

@[simp, ftrans_simp]
theorem postComp_assoc (x : 𝒟'(X,U)) (y : U ⊸ 𝒟'(Y,V)) (f : V ⊸ W) (φ : Y → R) :
    (x.postComp y).postComp (fun T ⊸ T.postComp f)
    =
    (x.postComp (fun u ⊸ (y u).postComp f)) := sorry_proof

@[action_push]
theorem postComp_extAction (x : 𝒟'(X,U)) (f : U ⊸ V) (φ : X → W) (L : V ⊸ W ⊸ Z) :
    (x.postComp y).extAction φ L
    =
    (x.extAction φ (fun u ⊸ fun w ⊸ L (f u) w)) := sorry_proof

@[action_push]
theorem postComp_restrict_extAction (x : 𝒟'(X,U)) (f : U ⊸ V) (A : Set X) (φ : X → W) (L : V ⊸ W ⊸ Z) :
    ((x.postComp f).restrict A).extAction φ L
    =
    ((x.restrict A).extAction φ (fun u ⊸ fun w ⊸ (L (f u) w))) := sorry_proof

-- @[simp, ftrans_simp, action_push]
-- theorem Distribution.zero_postExtAction (φ : Y → R) : (0 : 𝒟'(X,𝒟'(Y,Z))).postExtAction φ = 0 := by sorry_proof

-- -- todo: this needs some integrability condition
-- @[action_push]
-- theorem Distribution.add_postExtAction (T T' : 𝒟'(X,𝒟'(Y,Z))) (φ : Y → R) :
--     (T + T').postExtAction φ = T.postExtAction φ + T'.postExtAction φ := by sorry_proof

-- @[action_push]
-- theorem Distribution.sub_postExtAction (T T' : 𝒟'(X,𝒟'(Y,Z))) (φ : Y → R) :
--     (T - T').postExtAction φ = T.postExtAction φ - T'.postExtAction φ := by sorry_proof

-- @[action_push]
-- theorem Distribution.smul_postExtAction (r : R) (T : 𝒟'(X,𝒟'(Y,Z))) (φ : Y → R) :
--     (r • T).postExtAction φ = r • T.postExtAction φ := by sorry_proof

-- @[action_push]
-- theorem Distribution.neg_postExtAction (T : 𝒟'(X,𝒟'(Y,Z))) (φ : Y → R) :
--     (- T).postExtAction φ = - T.postExtAction φ := by sorry_proof

-- open BigOperators in
-- @[action_push]
-- theorem Distribution.fintype_sum_postExtAction {I} [Fintype I] (T : I → 𝒟'(X,𝒟'(Y,Z))) (φ : Y → R) :
--     (∑ i, T i).postExtAction φ = ∑ i, (T i).postExtAction φ := by sorry_proof

-- @[action_push]
-- theorem Distribution.ifD_postExtAction (T T' : 𝒟'(X,𝒟'(Y,Z))) (A : Set X) (φ : Y → R) :
--     (ifD A then T else T').postExtAction φ = ifD A then T.postExtAction φ else T'.postExtAction φ := by sorry_proof

-- -- @[action_push]
-- -- theorem Distribution.indextype_sum_postExtAction {I} [IndexType I] (T : I → 𝒟'(X,𝒟'(Y,Z))) (φ : Y → R) :
-- --     (∑ i, T i).postExtAction φ = ∑ i, (T i).postExtAction φ := by sorry_proof

----------------------------------------------------------------------------------------------------
-- Functions as distributions ----------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[coe, fun_trans]
noncomputable
def _root_.Function.toDistribution (f : X → Y) : 𝒟'(X,Y) :=
  fun φ ⊸ ∫' x, φ x • f x

def Distribution.IsFunction (T : 𝒟'(X,Y)) : Prop :=
  ∃ (f : X → Y), ∀ (φ : 𝒟 X),
      T φ = ∫' x, φ x • f x

noncomputable
def Distribution.toFunction (T : 𝒟'(X,Y)) : X → Y :=
  if h : T.IsFunction then
    choose h
  else
    0

@[action_push]
theorem Function.toDistribution_action (f : X → Y) (φ : 𝒟 X) :
    f.toDistribution φ = ∫' x, φ x • f x


import SciLean.Core.Rand.Rand
import SciLean.Core.Rand.Tactic
import SciLean.Core.Rand.Distributions.Uniform
import SciLean.Core.Rand.Distributions.Sphere

import SciLean.Core.FunctionTransformations
import SciLean.Core.Notation.CDeriv
import SciLean.Core.Notation.FwdDeriv

import SciLean.Core.FloatAsReal
import SciLean.Tactic.LetUtils
import SciLean.Tactic.LetEnter
import SciLean.Util.RecVal

import SciLean.Tactic.ConvInduction

import SciLean.Core.FunctionSpaces

open MeasureTheory

namespace SciLean.Rand

variable {Y : Type} [SemiHilbert Float Y] [Module ℝ Y] [IsScalarTower ℝ Float Y] [MeasurableSpace Y]
set_default_scalar Float

def pi' := 3.14159265359

open RealScalar in
noncomputable
def harmonicRec (n : ℕ) (φ : Vec3 → Float) (g : Vec3 → Y) (x : Vec3) : Y :=
  match n with
  | 0 => g x
  | m+1 => (4*(pi':Float))⁻¹ • ∫' (x' : sphere (0:Vec3) (1:Float)), harmonicRec m φ g (x + φ x • x'.1)

def walkOnSpheres (φ : Vec3 → Float) (g : Vec3 → Y) (n : ℕ) (x : Vec3) : Rand Y := do
  let f ←
    derive_random_approx
      (fun x => harmonicRec n φ g x)
    by
      induction n n' prev h
        . simp[harmonicRec]
        . simp[harmonicRec,h]
          simp only [smul_push,cintegral.arg_f.push_lambda]
          rw[integral_as_uniform_E Float]
      rw[pull_E_nat_recOn (x₀:=_) (r:=_) (hf:=by fun_prop)]
      simp (config:={zeta:=false})
  return f x

@[fun_prop]
theorem harmonicRec.arg_x.CDifferentiable_rule (n : ℕ)
    (φ : Vec3 → Float) (g : Vec3 → Y)
    (hφ : CDifferentiable Float φ) (hg : CDifferentiable Float g) :
    CDifferentiable Float (fun x : Vec3 => harmonicRec n φ g x) := by
  induction n <;> (simp[harmonicRec]; fun_prop)

noncomputable
def harmonicRec_fwdDeriv (n : ℕ)
    (φ : Vec3 → Float) (φ' : Vec3 → Vec3 → Float×Float)
    (g : Vec3 → Y) (g' : Vec3 → Vec3 → Y×Y) : Vec3 → Vec3 → Y×Y :=
    (∂> x, harmonicRec n φ g x)
  rewrite_by
    assuming (hφ' : (∂> φ) = φ') (hφ : CDifferentiable Float φ)
             (hg' : (∂> g) = g') (hg : CDifferentiable Float g)
    induction n n' du h
      . simp[harmonicRec]; autodiff
      . simp[harmonicRec];
        simp only [smul_push]; autodiff

def harmonicRec.arg_x.fwdDeriv_randApprox (n : ℕ)
    (φ : Vec3 → Float) (φ' : Vec3 → Vec3 → Float×Float)
    (g : Vec3 → Y) (g' : Vec3 → Vec3 → Y×Y)
    (x dx : Vec3) : Rand (Y×Y) := do
  derive_random_approx
    (harmonicRec_fwdDeriv n φ φ' g g' x dx)
  by
    unfold harmonicRec_fwdDeriv
    conv =>
      pattern (fun _ _ _ _ => cintegral _ _)
      enter [n,du]
      conv =>
        conv => enter [x];rw[cintegral.arg_f.push_lambda]
        rw[cintegral.arg_f.push_lambda]
      rw[integral_as_uniform_E Float]
    rw[pull_E_nat_recOn (x₀:=_) (r:=_) (hf:=by fun_prop)]
    simp (config:={zeta:=false})

set_default_scalar Float

noncomputable
def harmonicRec' (n : ℕ) (φ : Vec3 ⟿FD Float) (g : Vec3 ⟿FD Y) (x : Vec3) : Y :=
  match n with
  | 0 => g x
  | m+1 => (4*(pi':Float))⁻¹ • ∫' (x' : sphere (0:Vec3) (1:Float)), harmonicRec' m φ g (x + φ x • x'.1)

@[fun_prop]
theorem harmonicRec'_CDifferentiable (n : ℕ) :
    CDifferentiable Float (fun (w : (Vec3 ⟿FD Float)×(Vec3 ⟿FD Y)×Vec3) => harmonicRec' n w.1 w.2.1 w.2.2)

:= by
  induction n <;> (simp[harmonicRec']; fun_prop)


import SciLean.Core.Objects.Vec
import SciLean.Core.Objects.SemiInnerProductSpace
import SciLean.Core.Objects.FinVec
import SciLean.Core.Objects.Scalar
import SciLean.Core.FunctionPropositions.CDifferentiable

import SciLean.Core.Meta.ToAnyPoint

set_option linter.unusedVariables false

open LeanColls

namespace SciLean

local notation "∞" => (⊤ : ℕ∞)

variable
  (K : Type _) [RCLike K] (n : ℕ∞)
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {ι : Type _} [IndexType ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

/-- `CDifferentiableAt f x` - conveniently differentiable function `f` at point `x`. -/
@[irreducible, fun_prop]
def ContCDiffAt (f : X → Y) (x : X) : Prop :=
  ∀ (c : K → X),
      c 0 = x
      →
      Curve.ContDiffAt c 0 n
      →
      Curve.ContDiffAt (f∘c) 0 n

/-- `ContCDiff f` - conveniently differentiable function `f`.  -/
@[fun_prop]
def ContCDiff (f : X → Y) : Prop := ∀ x, ContCDiffAt K n f x

@[fun_prop, to_any_point]
theorem ContCDiffAt.id_rule (x : X) :
    ContCDiffAt K n (fun x : X => x) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.const_rule (y : Y) (x : X) :
    ContCDiffAt K n (fun _ : X => y) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.comp_rule (x : X)
    (f : Y → Z) (g : X → Y)
    (hf : ContCDiffAt K n f (g x)) (hg : ContCDiffAt K n g x) :
    ContCDiffAt K n (fun x => f (g x)) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.apply_rule
    (i : ι) (x) : ContCDiffAt K n (fun x : (i : ι) → E i => x i) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.pi_rule (x : X)
    (f : X → (i : ι) → E i)
    (hf : ∀ i, ContCDiffAt K n (f · i) x) :
    ContCDiffAt K n (fun x i => f x i) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.le_rule (x : X) (f : X → Y) (hf : ContCDiffAt K m f x) (h : n ≤ m) :
    ContCDiffAt K n f x := sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiff.le_rule (f : X → Y) (hf : ContCDiff K m f) (h : n ≤ m) :
    ContCDiff K n f := sorry_proof

@[fun_prop]
theorem ContCDiff.ContCDiffAt_rule (x : X) (f : X → Y) (hf : ContCDiff K n f) :
    ContCDiffAt K n f x := hf x

@[fun_prop]
theorem ContCDiff.id_rule :
    ContCDiff K n (fun x : X => x) := by intro x; fun_prop

@[fun_prop]
theorem ContCDiff.const_rule (y : Y) :
    ContCDiff K n (fun _ : X => y) := by intro x; fun_prop

@[fun_prop]
theorem ContCDiff.comp_rule
    (f : Y → Z) (g : X → Y)
    (hf : ContCDiff K n f) (hg : ContCDiff K n g) :
    ContCDiff K n (fun x => f (g x)) := by intro x; fun_prop

@[fun_prop]
theorem ContCDiff.apply_rule
    (i : ι) : ContCDiff K n (fun x : (i : ι) → E i => x i) := by intro x; fun_prop

@[fun_prop]
theorem ContCDiff.pi_rule (x : X)
    (f : X → (i : ι) → E i)
    (hf : ∀ i, ContCDiff K n (f · i)) :
    ContCDiff K n (fun x i => f x i) := by intro x; apply ContCDiffAt.pi_rule; fun_prop

----------------------------------------------------------------------------------------------------

-- transition rules to CDifferentiable

@[fun_prop]
theorem CDifferentaibleAt.ContCDiffAt_rule (x : X) (f : X → Y) (hf : ContCDiffAt K n f x) (h : 0 < n) :
    CDifferentiableAt K f x := sorry_proof

@[fun_prop]
theorem CDifferentaibleAt.ContCDiffAt_rule' (x : X) (f : X → Y) (hf : ContCDiffAt K ∞ f x) :
    CDifferentiableAt K f x := by fun_prop (disch:=aesop)

@[fun_prop]
theorem CDifferentaible.ContCDiff_rule (f : X → Y) (hf : ContCDiff K n f) (h : 0 < n) :
    CDifferentiable K f := by intro x; fun_prop (disch:=assumption)

@[fun_prop]
theorem CDifferentaible.ContCDiff_rule' (f : X → Y) (hf : ContCDiff K ∞ f) :
    CDifferentiable K f


import Lean.Meta.Tactic.Simp

import Mathlib.Tactic.FunProp.RefinedDiscrTree

namespace SciLean

namespace Tactic.RefinedSimp

open Lean Meta Mathlib.Meta.FunProp

initialize registerTraceClass `Meta.Tactic.simp.guard

/-- Optional guard on simp theorems. You can require that an argument is not for example identity
function or constant function after unification. -/
inductive ArgGuard where
  /-- Argument can't be identity function -/
  | notId
  /-- Argument can't be constant function -/
  | notConst
  deriving Inhabited, BEq, Repr

/-- Same as SimpTheorem but works with RefinedDiscrTree rather than with normal DescrTree.

It has one additional feature and that is argument guard. For example, you can say that do not apply
this theorem if theorem argument `f` unifies to identity function.
-/
structure RefinedSimpTheorem where
  keys        : List RefinedDiscrTree.DTExpr := []
  /--
    It stores universe parameter names for universe polymorphic proofs.
    Recall that it is non-empty only when we elaborate an expression provided by the user.
    When `proof` is just a constant, we can use the universe parameter names stored in the declaration.
   -/
  levelParams : Array Name := #[]
  proof       : Expr
  priority    : Nat  := eval_prio default
  post        : Bool := true
  /-- `perm` is true if lhs and rhs are identical modulo permutation of variables. -/
  perm        : Bool := false
  /--
    `origin` is mainly relevant for producing trace messages.
    It is also viewed an `id` used to "erase" `simp` theorems from `SimpTheorems`.
  -/
  origin      : Meta.Origin
  /-- `rfl` is true if `proof` is by `Eq.refl` or `rfl`. -/
  rfl         : Bool
  /-- Array of `(theorem argument id, argument guard)` specifying additional constraints on when
  to apply this theorem. For example, if the theorem has arugument `(f : X → X)` with index `3` then
  `guards := #[(3,.notId)]` will stop applying this theorem if `f` unifies to identity function. -/
  guards      : Array (ℕ × ArgGuard) := #[]
  deriving Inhabited, BEq

def RefinedSimpTheorem.getValue (simpThm : RefinedSimpTheorem) : MetaM Expr := do
  if simpThm.proof.isConst && simpThm.levelParams.isEmpty then
    let info ← getConstInfo simpThm.proof.constName!
    if info.levelParams.isEmpty then
      return simpThm.proof
    else
      return simpThm.proof.updateConst! (← info.levelParams.mapM (fun _ => mkFreshLevelMVar))
  else
    let us ← simpThm.levelParams.mapM fun _ => mkFreshLevelMVar
    return simpThm.proof.instantiateLevelParamsArray simpThm.levelParams us

def RefinedSimpTheorem.toSimpTheorem (simpThm : RefinedSimpTheorem) : SimpTheorem := {
  keys        := #[]
  levelParams := simpThm.levelParams
  proof       := simpThm.proof
  priority    := simpThm.priority
  post        := simpThm.post
  perm        := simpThm.perm
  origin      := simpThm.origin
  rfl         := simpThm.rfl
}

abbrev RefinedSimpTheoremTree := RefinedDiscrTree RefinedSimpTheorem

structure RefinedSimpTheorems where
  pre          : RefinedSimpTheoremTree := {}
  post         : RefinedSimpTheoremTree := {}
  deriving Inhabited

/-- -/
initialize refinedSimpTheoremsExt : SimpleScopedEnvExtension RefinedSimpTheorem RefinedSimpTheorems ←
  registerSimpleScopedEnvExtension {
    name     := by exact decl_name%
    initial  := {}
    addEntry := fun t thm =>
      if thm.post then
        {t with post := thm.keys.foldl (init:=t.post) (fun tree key => tree.insertDTExpr key thm) }
      else
        {t with pre := thm.keys.foldl (init:=t.pre) (fun tree key => tree.insertDTExpr key thm) }
  }

def getTheoremFromConst (declName : Name) (prio : Nat := eval_prio default) (post := true)
    (guards : Array (ℕ × ArgGuard)) : MetaM RefinedSimpTheorem := do
  let info ← getConstInfo declName
  let (_,_,b') ← forallMetaTelescope info.type
  let keys := ← RefinedDiscrTree.mkDTExprs b'.appFn!.appArg! {} false
  let thm : RefinedSimpTheorem := {
    keys        := keys
    levelParams := info.levelParams.toArray
    proof       := info.value!
    priority    := prio
    post        := post
    perm        := false
    origin      := .decl declName
    rfl         := false
    guards      := guards
  }

  return thm

def addTheorem (declName : Name) (attrKind : AttributeKind := .global)
    (prio : Nat := eval_prio default) (post := true) (guards : Array (ℕ × ArgGuard) := #[]) :
    MetaM Unit := do
  let thm ← getTheoremFromConst declName prio post guards
  refinedSimpTheoremsExt.add thm attrKind

/-- Check if `thm` can be applied to `e` and if the theorem argument `A : W → Set X` is not
a constant function. -/
def theoremGuard (e : Expr) (thm : RefinedSimpTheorem) : MetaM Bool

:= do
  if thm.guards.size = 0 then return true

  let val  ← thm.getValue
  let type ← inferType val

  let (xs, _, type) ← forallMetaTelescopeReducing type
  let type ← whnf (← instantiateMVars type)
  let lhs := type.appFn!.appArg!

  if (← isDefEq e lhs) then
    for (i,guardType) in thm.guards do
      let x ← instantiateMVars xs[i]!

      match guardType with
      | .notId =>
        let .lam _ _ body _ := x
          | continue
        if body == .bvar 0 then
          trace[Meta.Tactic.simp.guard] "not applying {← ppOrigin thm.origin} to \
                          {← ppExpr e} bacause {← ppExpr x} is identity function"
          return false
      | .notConst =>
        let .lam _ _ body _ := x
          | continue
        if ¬body.hasLooseBVars then
          trace[Meta.Tactic.simp.guard] "not applying {← ppOrigin thm.origin} to \
                          {← ppExpr e} bacause {← ppExpr x} is constant function"
          return false

    return true
  else
    return false


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


import SciLean.Core.FunctionTransformations
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R C} [Scalar R C] [RealScalar R]
  {W} [Vec C W]
  {U} [SemiInnerProductSpace C U]

--------------------------------------------------------------------------------
-- Exp -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem exp.arg_x.CDifferentiable_rule
  (x : W → C) (hx : CDifferentiable C x)
  : CDifferentiable C fun w => exp (x w) := sorry_proof

@[fun_trans]
theorem exp.arg_x.ceriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : cderiv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      xdx.2 * e := sorry_proof

@[fun_trans]
theorem exp.arg_x.fwdDeriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    fwdDeriv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      (e, xdx.2 * e) := by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem exp.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => exp (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate
@[fun_trans]
theorem exp.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => exp (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (exp xdx.1, fun dy => xdx.2 (conj (exp xdx.1) * dy)) := by
  unfold revDeriv
  fun_trans only [fwdDeriv, smul_push, ftrans_simp]

--------------------------------------------------------------------------------
-- Log -------------------------------------------------------------------------
--------------------------------------------------------------------------------

section Log

variable
  {R : Type _} [RealScalar R]
  {W : Type _} [Vec R W]
  {U : Type _} [SemiInnerProductSpace R U]

@[fun_prop]
theorem log.arg_x.CDifferentiableAt_rule
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    CDifferentiableAt R (fun w => log (x w)) w := sorry_proof

@[fun_prop]
theorem log.arg_x.CDifferentiable_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    CDifferentiable R (fun w => log (x w)) := by
  intro x; fun_prop (disch:=aesop)

@[fun_trans]
theorem log.arg_x.cderiv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    cderiv R (fun w => log (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / abs xdx.1 := sorry_proof

@[fun_trans]
theorem log.arg_x.cderiv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    cderiv R (fun w => log (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / abs xdx.1 := by
  funext x
  fun_trans (disch:=aesop)

@[fun_trans]
theorem log.arg_x.fwdDeriv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    fwdDeriv R (fun w => log (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      let l := log xdx.1
      (l, xdx.2 / abs xdx.1) :=
by
  unfold fwdDeriv; fun_trans (disch:=assumption); simp[fwdDeriv]

@[fun_trans]
theorem log.arg_x.fwdDeriv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    fwdDeriv R (fun w => log (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      let l := log xdx.1
      (l, xdx.2 / abs xdx.1) :=
by
  unfold fwdDeriv; fun_trans (disch:=assumption); simp[fwdDeriv]

@[fun_prop]
theorem log.arg_x.HasAdjDiffAt_rule
    (u : U) (x : U → R) (hx : HasAdjDiffAt R x u) (hu : x u ≠ 0) :
    HasAdjDiffAt R (fun u => log (x u)) u := by
  constructor
  fun_prop (disch:=aesop)
  fun_trans (disch:=aesop) [fwdDeriv]; fun_prop

@[fun_prop]
theorem log.arg_x.HasAdjDiff_rule
    (x : U → R) (hx : HasAdjDiff R x) (hu : ∀ u, x u ≠ 0) :
    HasAdjDiff R (fun u => log (x u)) := by
  intro u;
  fun_prop (disch:=aesop)

open ComplexConjugate
@[fun_trans]
theorem log.arg_x.revDeriv_rule_at
    (u : U) (x : U → R) (hx : HasAdjDiffAt R x u) (hu : x u ≠ 0) :
    revDeriv R (fun u => log (x u)) u
    =
    let xdx := revDeriv R x u
    (log xdx.1, fun dy => xdx.2 ((abs (x u))⁻¹ * dy)) := by
  unfold revDeriv
  fun_trans (disch:=aesop) only [fwdDeriv, smul_push, ftrans_simp]

open ComplexConjugate

@[fun_trans]
theorem log.arg_x.revDeriv_rule
    (x : U → R) (hx : HasAdjDiff R x) (hu : ∀ u, x u ≠ 0) :
    revDeriv R (fun u => log (x u))
    =
    fun u =>
      let xdx := revDeriv R x u
      (log xdx.1, fun dy => xdx.2 ((abs (x u))⁻¹ * dy))


import SciLean.Core.FunctionPropositions.CDifferentiable
import SciLean.Core.FunctionPropositions.HasSemiAdjoint

import SciLean.Core.FunctionTransformations.CDeriv

import Lean.Meta.Tactic.Assumption

set_option linter.unusedVariables false

open LeanColls

namespace SciLean

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

@[fun_prop]
def HasAdjDiffAt (f : X → Y) (x : X) : Prop :=
    CDifferentiableAt K f x ∧ HasSemiAdjoint K (cderiv K f x)

@[fun_prop]
def HasAdjDiff (f : X → Y) : Prop := ∀ x, HasAdjDiffAt K f x

@[fun_prop]
theorem HasAdjDiff.hasAdjDiffAt (f : X → Y) (x) (hf : HasAdjDiff K f) :
    HasAdjDiffAt K f x


import SciLean.Core.Objects.SemiInnerProductSpace
import SciLean.Core.FunctionPropositions.IsLinearMap
import SciLean.Core.FunctionPropositions.CDifferentiable
import SciLean.Core.FunctionPropositions.IsSmoothLinearMap
import SciLean.Core.Simp
import SciLean.Core.Meta.GenerateLinearMapSimp

import Mathlib.Tactic.FunTrans.Attr
import Mathlib.Tactic.FunTrans.Elab

set_option linter.unusedVariables false

open LeanColls

namespace SciLean

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

@[fun_prop]
structure HasSemiAdjoint (f : X → Y) : Prop where
  semiAdjoint_exists : ∃ f' : Y → X,
     ∀ x y, TestFunction x → ⟪y, f x⟫[K] = ⟪f' y, x⟫[K]
   -- at some point I was convinced that these conditions are important
   -- maybe add: ∀ x, TestFunction x → TestFunction (f x) - I think this is important for proving linearity of f'
   -- maybe add: ∀ y, TestFunction y → TestFunction (f' y)
  -- Right now we have no use for functions that have semiAdjoint and are not differentiable
  -- so we just assume that all are differentiable
  is_differentiable : CDifferentiable K f

/-- Generalization of adjoint of linear map `f : X → Y`.

If `f : X → Y` is linear map between Hilbert spaces then `semiAdjoint K f = adjoint K f`.

`semiAdjoint` is a generalization of adjoint to spaces that are not necessarily complete
and might have inner product defined only on a dense subset, see `SemiInnerProductSpace`
for more information.
 -/
@[fun_trans]
noncomputable
def semiAdjoint (f : X → Y) (y : Y) : X :=
  match Classical.dec (HasSemiAdjoint K f) with
  | isTrue h => Classical.choose h.semiAdjoint_exists y
  | isFalse _ => 0

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

/-- `semiAdjoint K f ·` is always linear because either `f` has adjoint and is linear or
`semiAdjoint K f ·` is zero function and thus linear too. -/
@[fun_prop]
theorem semiAdjoint.arg_y.IsLinearMap_rule (f : X → Y) :
    IsLinearMap K (fun y => semiAdjoint K f y) := sorry_proof

#generate_linear_map_simps SciLean.semiAdjoint.arg_y.IsLinearMap_rule

@[fun_prop]
theorem semiAdjoint.arg_y.CDifferentiable_rule (f : X → Y) :
    CDifferentiable K (fun y => semiAdjoint K f y) := sorry_proof

@[fun_prop]
theorem semiAdjoint.arg_y.IsSmoothLinearMap_rule (f : X → Y) :
    IsSmoothLinearMap K (fun y => semiAdjoint K f y) := by constructor; fun_prop; fun_prop

-- Do we need joint smoothness in `w` and `x` for `f` ???
@[fun_prop]
theorem semiAdjoint.arg_f.IsSmoothLinearMap_rule (f : W → X → Y)
    (hf₁ : ∀ x, IsSmoothLinearMap K (f · x)) (hf₂ : ∀ w, HasSemiAdjoint K (f w ·)) :
    IsSmoothLinearMap K (fun w => semiAdjoint K (f w)) := sorry_proof

@[fun_prop]
theorem semiAdjoint.arg_fy.CDifferentiable_rule (f : W → X → Y) (y : W → Y)
    (hf₁ : CDifferentiable K (fun (w,x) => f w x)) (hf₂ : ∀ w, HasSemiAdjoint K (f w ·))
    (hy : CDifferentiable K y) :
    CDifferentiable K (fun w => semiAdjoint K (f w) (y w)) := sorry_proof

@[fun_prop]
theorem HasSemiAdjoint.CDifferentiable (f : X → Y) (hf : HasSemiAdjoint K f) :
    CDifferentiable K f


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


import SciLean.Core.FunctionPropositions.ContCDiff
import SciLean.Core.FunctionTransformations

import SciLean.Core.Notation.CDeriv

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)]

local notation "∞" => (⊤ : ℕ∞)

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

structure ContCDiffMap (n : ℕ∞) (X Y : Type _)  [Vec K X] [Vec K Y] where
  toFun : X → Y
  is_cont_cdiff_map : ContCDiff K n toFun

variable (n : ℕ∞)

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

macro X:term:25 " ⟿[" K:term "," n:term "] " Y:term:26 : term =>
  `(ContCDiffMap $K $n $X $Y)

macro X:term:25 " ⟿[" n:term "] " Y:term:26 : term =>
  `(ContCDiffMap defaultScalar% $n $X $Y)

macro X:term:25 " ⟿ " Y:term:26 : term =>
  `(ContCDiffMap defaultScalar% ∞ $X $Y)

@[app_unexpander ContCDiffMap] def unexpandContCDiffMap : Lean.PrettyPrinter.Unexpander
  | `($(_) $R $n $X $Y) => `($X ⟿[$R,$n] $Y)
  | _  => throw ()

@[fun_prop]
theorem ContCDiffMap_apply_right (f : X ⟿[K,n] Y) : ContCDiff K n (fun x => f x) := f.2

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

variable {K}

-- def ContCDiffMap.mk' (f : X → Y) (hf : ContCDiff K n f) :
--     ContCDiffMap K n X Y := ⟨f, hf⟩

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

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

variable (K)
def ContCDiffMap.mk' (f : X → Y) (hf : ContCDiff K n f) : X ⟿[K,n] Y := ⟨f,hf⟩

open Lean Parser Term
macro "fun " x:funBinder " ⟿[" K:term "," n:term "] " b:term : term =>
  `(ContCDiffMap.mk' $K $n (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

open Lean Parser Term
macro "fun " x:funBinder " ⟿[" n:term "] " b:term : term =>
  `(ContCDiffMap.mk' defaultScalar% $n (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

open Lean Parser Term
macro "fun " x:funBinder " ⟿ " b:term : term =>
  `(ContCDiffMap.mk' defaultScalar% ∞ (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

@[app_unexpander ContCDiffMap.mk'] def unexpandContCDiffMapMk : Lean.PrettyPrinter.Unexpander

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

@[simp, ftrans_simp]
theorem ContCDiffMap_eta (f : X ⟿[K,n] Y) : (fun x ⟿[K,n] f x) = f := by rfl

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

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

section AlgebraSimps

variable (f g : X ⟿[K,n] Y) (x : X) (r : K)

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

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

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


import SciLean.Core.FunctionPropositions
import SciLean.Core.FunctionTransformations
import SciLean.Data.ArrayType.Algebra

open SciLean

set_option linter.unusedVariables false

open LeanColls

section GenericArrayType

variable
  {K : Type} [RCLike K]
  {Cont : Type} {Idx : Type |> outParam} {Elem : Type |> outParam}
  [ArrayType Cont Idx Elem] [IndexType Idx] [LawfulIndexType Idx] [DecidableEq Idx]

-- Indexed.get -----------------------------------------------------------------
--------------------------------------------------------------------------------

section OnVectorSpaces

variable {R : Type} [CommSemiring R] [AddCommGroup Elem] [Module R Elem]

@[fun_prop]
theorem LeanColls.Indexed.get.arg_cont.IsLinearMap_rule_simple (idx : Idx) :
    IsLinearMap R (fun xs : Cont => xs[idx]) := by
  constructor
  . intros; simp
  . intros; simp

#generate_linear_map_simps LeanColls.Indexed.get.arg_cont.IsLinearMap_rule_simple

end OnVectorSpaces

section OnNormedSpaces

variable
  {X : Type _} [NormedAddCommGroup X] [NormedSpace K X]
  [NormedAddCommGroup Elem] [NormedSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.get.arg_cont.IsContinuousLinearMap_rule_simple (idx : Idx) :
    IsContinuousLinearMap K (λ cont : Cont => cont[idx]) := sorry_proof

-- automatically infer: Differentiable, DifferentiableAt, ContDiff, fderiv, fwdFDeriv

end OnNormedSpaces

section OnVec

variable
  {X : Type} [Vec K X]
  [Vec K Elem]

@[fun_prop]
theorem GetElem.getElem.arg_cont.IsSmoothLinearMap_rule_simple
  (idx : Idx)
  : IsSmoothLinearMap K (fun xs : Cont => xs[idx]) := sorry_proof

-- automatically infer: CDifferentiable, CDifferentiableAt, CContDiff, cderiv, fwdDeriv

end OnVec

section OnSemiInnerProductSpace

variable
  {X : Type} [SemiInnerProductSpace K X]
  [SemiInnerProductSpace K Elem]

@[fun_prop]
theorem GetElem.getElem.arg_cont.HasSemiAdjoint_rule_simple (idx : Idx) :
    HasSemiAdjoint K (fun cont : Cont => cont[idx]) := sorry_proof

@[fun_trans]
theorem GetElem.getElem.arg_cont.semiAdjoint_rule_simple (idx : Idx) :
    semiAdjoint K (fun cont : Cont => cont[idx])
    =
    fun elem => oneHot (X:=Cont) idx elem :=
by
  sorry_proof

@[fun_trans]
theorem GetElem.getElem.arg_cont.revCDeriv_rule
  (f : X → Cont) (idx : Idx)
  (hf : HasAdjDiff K f)
  : revDeriv K (fun x => (f x)[idx])
    =
    fun x =>
      let ydf := revDerivProj K Idx f x
      (ydf.1[idx],
       fun delem => ydf.2 idx delem) :=
by
  unfold revDeriv;
  conv =>
    lhs; fun_trans; fun_trans; simp

@[fun_trans]
theorem GetElem.getElem.arg_cont.revCDerivUpdate_rule
  (f : X → Cont) (idx : Idx)
  (hf : HasAdjDiff K f)
  : revDerivUpdate K (fun x => (f x)[idx])
    =
    fun x =>
      let ydf := revDerivProjUpdate K Idx f x
      (ydf.1[idx],
       fun delem dx => ydf.2 idx delem dx) :=
by
  unfold revDerivUpdate;
  conv =>
    lhs; fun_trans; fun_trans; simp

@[fun_trans]
theorem GetElem.getElem.arg_cont.revCDerivProj_rule
  {J ElemJ} [StructType Elem J ElemJ] [IndexType J] [LawfulIndexType J] [DecidableEq J]
  [∀ j, SemiInnerProductSpace K (ElemJ j)] [SemiInnerProductSpaceStruct K Elem J ElemJ]
  (f : X → Cont) (idx : Idx)
  (hf : HasAdjDiff K f)
  : revDerivProj K J (fun x => (f x)[idx])
    =
    fun x =>
      let ydf := revDerivProj K (Idx×J) f x
      (ydf.1[idx],
       fun j delem => ydf.2 (idx,j) delem) :=
by
  unfold revDerivProj
  conv =>
    lhs; fun_trans; fun_trans; simp
  funext x; simp[revDerivProj]
  funext j de
  simp [revDeriv]
  congr
  sorry_proof

end OnSemiInnerProductSpace

-- Indexed.set ----------------------------------------------------------------
--------------------------------------------------------------------------------

section OnVectorSpaces

variable {R : Type} [CommSemiring R] [AddCommGroup Elem] [Module R Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_contelem.IsLinearMap_rule_simple (idx : Idx) :
    IsLinearMap R (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) := by
  constructor
  . intros; simp; ext; sorry_proof
  . intros; simp; ext; sorry_proof

#generate_linear_map_simps LeanColls.Indexed.set.arg_contelem.IsLinearMap_rule_simple

end OnVectorSpaces

section OnNormedSpaces

variable
  {X : Type _} [NormedAddCommGroup X] [NormedSpace K X]
  [NormedAddCommGroup Elem] [NormedSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_contelem.IsContinuousLinearMap_rule_simple :
    IsContinuousLinearMap K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) :=
  sorry_proof

-- automatically infer: Differentiable, DifferentiableAt, ContDiff, fderiv, fwdFDeriv

end OnNormedSpaces

section OnVec

variable
  {X : Type} [Vec K X]
  [Vec K Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_contelem.IsSmoothLinearMap_rule_simple (idx : Idx) :
    IsSmoothLinearMap K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) := sorry_proof

-- automatically infer: CDifferentiable, CDifferentiableAt, CContDiff, cderiv, fwdDeriv

end OnVec

section OnSemiInnerProductSpace

variable
  {X : Type} [SemiInnerProductSpace K X]
  [SemiInnerProductSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_cont.HasSemiAdjoint_rule_simple (idx : Idx) :
    HasSemiAdjoint K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) := sorry_proof

@[fun_trans]
theorem LeanColls.Indexed.set.arg_cont.semiAdjoint_rule_simple (idx : Idx) :
    semiAdjoint K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem)
    =
    fun cont' : Cont => (Indexed.set cont' idx 0, cont'[idx]) := sorry_proof

@[fun_trans]
theorem LeanColls.Indexed.set.arg_contelem.revCDeriv_rule
  (cont : X → Cont) (idx : Idx) (elem : X → Elem)
  (hcont : HasAdjDiff K cont) (helem : HasAdjDiff K elem)
  : revDeriv K (fun x => Indexed.set (cont x) idx (elem x))
    =
    fun x =>
      let cdc := revDeriv K cont x
      let ede := revDerivUpdate K elem x
      (Indexed.set cdc.1 idx ede.1,
       fun dc : Cont =>
         let dci := dc[idx]
         let dc := Indexed.set dc idx 0
         ede.2 dci (cdc.2 dc)) := by
  unfold revDeriv;
  conv =>
    lhs; fun_trans; fun_trans; simp

@[fun_trans]
theorem LeanColls.Indexed.set.arg_contelem.revCDerivUpdate_rule
  (cont : X → Cont) (idx : Idx) (elem : X → Elem)
  (hcont : HasAdjDiff K cont) (helem : HasAdjDiff K elem)
  : revDerivUpdate K (fun x => Indexed.set (cont x) idx (elem x))
    =
    fun x =>
      let cdc := revDerivUpdate K cont x
      let ede := revDerivUpdate K elem x
      (Indexed.set cdc.1 idx ede.1,
       fun (dc : Cont) dx =>
         let dci := dc[idx]
         let dc := Indexed.set dc idx 0
         ede.2 dci (cdc.2 dc dx))

:= by
  unfold revDerivUpdate
  conv =>
    lhs; fun_trans; fun_trans; simp
  simp[revDerivUpdate,add_assoc]


import Mathlib.Data.Erased
import Mathlib.Control.Random
import Mathlib.MeasureTheory.Integral.Bochner
import Mathlib.MeasureTheory.Decomposition.Lebesgue

import SciLean.Core.FunctionPropositions.Bijective
import SciLean.Core.Objects.Scalar
import SciLean.Core.Integral.CIntegral
import SciLean.Core.Rand.SimpAttr

import Mathlib.MeasureTheory.Measure.GiryMonad

open MeasureTheory ENNReal BigOperators Finset

namespace SciLean

abbrev erase (a : α) : Erased α := .mk a

@[simp,ftrans_simp]
theorem erase_out {α} (a : α) : (erase a).out = a := by simp[erase]

/-- `x : Rand X` is a random variable of type `X`

You can:
  - generate sample with `x.get : IO X`
  - get probability measure with `x.ℙ : Measure X`

The internal fields `spec` and `rand` are just an internal implementation of `Rand` and should not
be accessed by normal users.

TODO: Hide implementation using quotients or something like that
-/
structure Rand (X : Type _)  where
  /-- `spec` defines a probability measure by computing an expectation. This means if `x : Rand X`
  corresponds to a probability measure `μ` then for `φ : X → ℝ`
  ```
  x.spec.out φ = ∫ x, φ x ∂μ
  ```

  Using `(X→ℝ)→ℝ` instead of `Measure X` for the specification of random variables has the
  advantage that we can reuse Lean's `do` notation.
  -/
  spec : Erased ((X→ℝ)→ℝ)
  /-- `rand` is a pseudo randon number generator implemented using the "Standard" number generator
  -/
  rand : StateM StdGen X

namespace Rand

def _root_.Function.IsMeasure {X} [MeasurableSpace X] (F : (X → ℝ) → ℝ) : Prop :=
  ∃ μ : Measure X, ∀ (f : X → ℝ), F f = ∫ x, f x ∂μ

open Classical in
/-- Probability measure of a random variable -/
@[pp_dot]
noncomputable
def ℙ {X} [MeasurableSpace X] (r : Rand X) : Measure X :=
  if h : r.spec.out.IsMeasure then
    choose h
  else
    0

/-- Specification of `x : Rand X` is really saying that it is a probability measure. -/
class LawfulRand (x : Rand X) [MeasurableSpace X] where
  is_measure : x.spec.out.IsMeasure
  is_prob : IsProbabilityMeasure x.ℙ

variable {X Y Z : Type _}
  [MeasurableSpace X]
  [MeasurableSpace Y]

instance instIsProbabilityMeasureℙ (x : Rand X) [inst : LawfulRand x] : IsProbabilityMeasure (x.ℙ) := inst.is_prob

/-- Extensionality of random variable.

WARNING: This theorem is inconsistent!!! The random generators `x.rand` and `y.rand` might differ.
         We are not trying to model pseudo-random numbers. We assume that every random number
         generator is a true random number generator. Thus the result of any probabilistic program
         should be independent on the exact generator up to some randomness.

TODO: We might quotient all the random number generators corresponding to the measure `x.ℙ`  under
      the assumption that they are all true random generators. I believe that such type would be
      a singleton i.e. all the random number generators are all the same.
-/
@[ext]
axiom ext (x y : Rand X) : x.spec.out = y.spec.out → x = y

/-- Generate rundom number using IO randomness -/
def get (x : Rand X) : IO X := do
  let stdGen ← IO.stdGenRef.get
  let (res, new) := x.rand stdGen
  let _ ← IO.stdGenRef.set new
  pure res

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

instance : Monad Rand where
  pure x := {
    spec := erase (fun φ => φ x),
    rand := pure x
  }
  bind x f := {
    spec := erase (fun φ => x.spec.out (fun x => (f x).spec.out φ)),
    rand := bind x.rand (fun x => (f x).rand)
  }

instance : LawfulMonad Rand where
  bind_pure_comp := by intros; rfl
  bind_map       := by intros; rfl
  pure_bind      := by intros; ext; simp[Bind.bind,Pure.pure]
  bind_assoc     := by intros; ext; simp[Bind.bind,Pure.pure]
  map_const      := by intros; ext; rfl
  id_map         := by intros; ext; simp[Bind.bind,Pure.pure,id,Functor.map]
  seqLeft_eq     := by intros; ext; simp[Bind.bind,Pure.pure,Seq.seq,Function.const,Functor.map,SeqLeft.seqLeft]
  seqRight_eq    := by intros; ext; simp[Bind.bind,Pure.pure,Seq.seq,Function.const,Functor.map,SeqRight.seqRight]
  pure_seq       := by intros; ext; simp[Bind.bind,Pure.pure,Seq.seq,Functor.map]

-- this needs some integrability and lawfulness of Rand
theorem swap_bind (f : X → Y → Z) (x : Rand X) (y : Rand Y) :
    (do let x' ← x; let y' ← y; pure (f x' y'))
    =
    (do let y' ← y; let x' ← x; pure (f x' y')) := by
  sorry_proof

@[simp, ftrans_simp]
theorem pure_ℙ (x : X) : (pure x : Rand X).ℙ = Measure.dirac x := sorry_proof

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

instance [Add X] : HAdd X (Rand X) (Rand X) := ⟨fun x' x => do
  let x'' ← x
  pure (x' + x'')⟩

instance [Add X] : HAdd (Rand X) X (Rand X) := ⟨fun x x' => do
  let x'' ← x
  pure (x'' + x')⟩

-- instance [Add X] : HAdd (Rand X) (Rand X) (Rand X) := ⟨fun x y =>
--   let x' ~ x
--   let y' ~ y
--   pure (x' + y')⟩

-- todo: add simp theorems that inline these operations

----------------------------------------------------------------------------------------------------
-- Simple Random Variable functions ----------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[pp_dot]
abbrev map (r : Rand X) (f : X → Y) : Rand Y := do
  let x' ← r
  return f x'

/-- Marginal distribution for the first component of a pair. -/
@[pp_dot]
abbrev fst (r : Rand (X×Y)) : Rand X := do
  let (x,_) ← r
  return x

/-- Marginal distribution for the second component of a pair. -/
@[pp_dot]
abbrev snd (r : Rand (X×Y)) : Rand Y := do
  let (_,y) ← r
  return y

@[simp, ftrans_simp]
theorem map_ℙ  (r : Rand X) (f : X → Y) :
  (r.map f).ℙ = r.ℙ.map f := sorry_proof

----------------------------------------------------------------------------------------------------
-- Expected Value ----------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

section ExpectedValue

variable
  {R} [RealScalar R]
  [AddCommGroup Y] [Module ℝ Y]
  [AddCommGroup Z] [Module ℝ Z]
  -- [AddCommGroup U] [TopologicalSpace U] [TopologicalAddGroup U] [Module ℝ U] [LocallyConveUSpace ℝ U]
  {U} [AddCommGroup U] [Module ℝ U]

@[pp_dot]
noncomputable
def 𝔼 (r : Rand X) (φ : X → Y) : Y := ∫' x, φ x ∂r.ℙ

@[simp, ftrans_simp, rand_push_E]
theorem pure_𝔼 (x : X) (φ : X → Y) :
    (pure (f:=Rand) x).𝔼 φ = φ x := by simp [𝔼]

-- What are the right assumptions here? Lambda lawfulness of `x` and `f x'` and integrability of `φ`
@[rand_push_E]
theorem bind_E (r : Rand X) (f : X → Rand Y) (φ : Y → Z) :
    (r >>= f).𝔼 φ = r.𝔼 (fun x' => (f x').𝔼 φ) := by simp[𝔼]; sorry_proof

-- todo: We might want this to hold without lawfulness
-- consider adding as a property inside of `Distribution` or `Rand`

@[simp, ftrans_simp, rand_push_E]
theorem E_zero (r : Rand X) :
    r.𝔼 (fun _ => (0 : U)) = 0


import SciLean.Core.FunctionPropositions.ContCDiff
import SciLean.Core.FunctionTransformations

import SciLean.Core.Notation.CDeriv

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)]

local notation "∞" => (⊤ : ℕ∞)

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

structure ContCDiffMap (n : ℕ∞) (X Y : Type _)  [Vec K X] [Vec K Y] where
  toFun : X → Y
  is_cont_cdiff_map : ContCDiff K n toFun

variable (n : ℕ∞)

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

macro X:term:25 " ⟿[" K:term "," n:term "] " Y:term:26 : term =>
  `(ContCDiffMap $K $n $X $Y)

macro X:term:25 " ⟿[" n:term "] " Y:term:26 : term =>
  `(ContCDiffMap defaultScalar% $n $X $Y)

macro X:term:25 " ⟿ " Y:term:26 : term =>
  `(ContCDiffMap defaultScalar% ∞ $X $Y)

@[app_unexpander ContCDiffMap] def unexpandContCDiffMap : Lean.PrettyPrinter.Unexpander
  | `($(_) $R $n $X $Y) => `($X ⟿[$R,$n] $Y)
  | _  => throw ()

@[fun_prop]
theorem ContCDiffMap_apply_right (f : X ⟿[K,n] Y) : ContCDiff K n (fun x => f x) := f.2

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

variable {K}

-- def ContCDiffMap.mk' (f : X → Y) (hf : ContCDiff K n f) :
--     ContCDiffMap K n X Y := ⟨f, hf⟩

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

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

variable (K)
def ContCDiffMap.mk' (f : X → Y) (hf : ContCDiff K n f) : X ⟿[K,n] Y := ⟨f,hf⟩

open Lean Parser Term
macro "fun " x:funBinder " ⟿[" K:term "," n:term "] " b:term : term =>
  `(ContCDiffMap.mk' $K $n (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

open Lean Parser Term
macro "fun " x:funBinder " ⟿[" n:term "] " b:term : term =>
  `(ContCDiffMap.mk' defaultScalar% $n (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

open Lean Parser Term
macro "fun " x:funBinder " ⟿ " b:term : term =>
  `(ContCDiffMap.mk' defaultScalar% ∞ (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

@[app_unexpander ContCDiffMap.mk'] def unexpandContCDiffMapMk : Lean.PrettyPrinter.Unexpander

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

@[simp, ftrans_simp]
theorem ContCDiffMap_eta (f : X ⟿[K,n] Y) : (fun x ⟿[K,n] f x) = f := by rfl

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

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

section AlgebraSimps

variable (f g : X ⟿[K,n] Y) (x : X) (r : K)

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


import SciLean.Core.FunctionTransformations
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R C} [Scalar R C]
  {W} [Vec C W]
  {U} [SemiInnerProductSpace C U]

--------------------------------------------------------------------------------
-- Sin -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem sin.arg_x.CDifferentiable_rule
    (x : W → C) (hx : CDifferentiable C x) :
    CDifferentiable C fun w => sin (x w) := sorry_proof

@[fun_trans]
theorem sin.arg_x.ceriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    cderiv C (fun w => sin (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let c := cos xdx.1
      xdx.2 * c := sorry_proof

@[fun_trans]
theorem sin.arg_x.fwdDeriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    fwdDeriv C (fun w => sin (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let s := sin xdx.1
      let c := cos xdx.1
      (s, xdx.2 * c) := by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem sin.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => sin (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate
@[fun_trans]
theorem sin.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => sin (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (sin xdx.1, fun dy => xdx.2 (conj (cos xdx.1) * dy)) := by
  unfold revDeriv
  fun_trans only [fwdDeriv, smul_push, ftrans_simp]

--------------------------------------------------------------------------------
-- Cos -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem cos.arg_x.CDifferentiable_rule
  (x : W → C) (hx : CDifferentiable C x)
  : CDifferentiable C fun w => cos (x w) := sorry_proof

@[fun_trans]
theorem cos.arg_x.ceriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : cderiv C (fun w => cos (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let s := -sin xdx.1
      xdx.2 * s := sorry_proof

@[fun_trans]
theorem cos.arg_x.fwdDeriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : fwdDeriv C (fun w => cos (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let s := sin xdx.1
      let c := cos xdx.1
      (c, xdx.2 * -s) :=
by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem cos.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => cos (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate

@[fun_trans]
theorem cos.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => cos (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (cos xdx.1, fun dy => xdx.2 (- conj (sin xdx.1) * dy))

:= by
  unfold revDeriv
  fun_trans only [fwdDeriv, smul_push, neg_push, ftrans_simp]


import SciLean.Core.FunctionPropositions.ContCDiff
import SciLean.Core.FunctionTransformations

import SciLean.Core.Notation.CDeriv

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)]

local notation "∞" => (⊤ : ℕ∞)

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

structure ContCDiffMap (n : ℕ∞) (X Y : Type _)  [Vec K X] [Vec K Y] where
  toFun : X → Y
  is_cont_cdiff_map : ContCDiff K n toFun

variable (n : ℕ∞)

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

macro X:term:25 " ⟿[" K:term "," n:term "] " Y:term:26 : term =>
  `(ContCDiffMap $K $n $X $Y)

macro X:term:25 " ⟿[" n:term "] " Y:term:26 : term =>
  `(ContCDiffMap defaultScalar% $n $X $Y)

macro X:term:25 " ⟿ " Y:term:26 : term =>
  `(ContCDiffMap defaultScalar% ∞ $X $Y)

@[app_unexpander ContCDiffMap] def unexpandContCDiffMap : Lean.PrettyPrinter.Unexpander
  | `($(_) $R $n $X $Y) => `($X ⟿[$R,$n] $Y)
  | _  => throw ()

@[fun_prop]
theorem ContCDiffMap_apply_right (f : X ⟿[K,n] Y) : ContCDiff K n (fun x => f x) := f.2

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

variable {K}

-- def ContCDiffMap.mk' (f : X → Y) (hf : ContCDiff K n f) :
--     ContCDiffMap K n X Y := ⟨f, hf⟩

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

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

variable (K)
def ContCDiffMap.mk' (f : X → Y) (hf : ContCDiff K n f) : X ⟿[K,n] Y := ⟨f,hf⟩

open Lean Parser Term
macro "fun " x:funBinder " ⟿[" K:term "," n:term "] " b:term : term =>
  `(ContCDiffMap.mk' $K $n (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

open Lean Parser Term
macro "fun " x:funBinder " ⟿[" n:term "] " b:term : term =>
  `(ContCDiffMap.mk' defaultScalar% $n (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

open Lean Parser Term
macro "fun " x:funBinder " ⟿ " b:term : term =>
  `(ContCDiffMap.mk' defaultScalar% ∞ (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

@[app_unexpander ContCDiffMap.mk'] def unexpandContCDiffMapMk : Lean.PrettyPrinter.Unexpander

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

@[simp, ftrans_simp]
theorem ContCDiffMap_eta (f : X ⟿[K,n] Y) : (fun x ⟿[K,n] f x) = f := by rfl

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

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

section AlgebraSimps

variable (f g : X ⟿[K,n] Y) (x : X) (r : K)

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

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


import SciLean.Core.FunctionTransformations
import SciLean.Core.Functions.Exp
import SciLean.Tactic.Autodiff
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean

variable
  {R C} [Scalar R C] [RealScalar R]
  {W} [Vec R W]
  {U} [SemiHilbert R U]

set_default_scalar R

----------------------------------------------------------------------------------------------------
-- Gaussian ----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

open Scalar RealScalar in
def gaussian (μ : U) (σ : R) (x : U) : R :=
  let x' := σ⁻¹ • (x - μ)
  1/(σ*sqrt (2*(pi : R))) * exp (- ‖x'‖₂²/2)

open Scalar RealScalar in
@[simp, ftrans_simp]
theorem log_gaussian (μ : U) (σ : R) (x : U) :
    log (gaussian μ σ x)
    =
    let x' := σ⁻¹ • (x - μ)
    (- ‖x'‖₂²/2 - log σ - log (sqrt (2*(pi :R)))) := by

  unfold gaussian
  simp [log_inv,log_mul,log_div,log_exp,log_one]
  ring

-- TODO: add deriving statements!!!
-- derive_cdifferentiable_at gaussian μ x by unfold gaussian; fun_prop
-- derive_cdifferentiable gaussian μ x

-- derive_cdifferentiable_at gaussian μ σ x
--   assuming σ ≠ 0

-- -- this can't have non-compositional version
-- derive_cdifferentiable gaussian μ σ x
--   assuming σ ≠ 0

@[fun_prop]
theorem gaussian.arg_μx.CDifferentiableAt_rule (w : W)
    (μ : W → U) (σ : R) (x : W → U)
    (hμ : CDifferentiableAt R μ w) (hx : CDifferentiableAt R x w) :
    CDifferentiableAt R (fun w => gaussian (μ w) σ (x w)) w := by

  unfold gaussian
  fun_prop

@[fun_prop]
theorem gaussian.arg_μx.CDifferentiable_rule
    (μ : W → U) (σ : R) (x : W → U)
    (hμ : CDifferentiable R μ) (hx : CDifferentiable R x) :
    CDifferentiable R (fun w => gaussian (μ w) σ (x w)) := by

  intro w
  fun_prop

@[fun_trans]
theorem gaussian.arg_μx.cderiv_rule
    (μ : W → U) (σ : R) (x : W → U)
    (hμ : CDifferentiable R μ) (hx : CDifferentiable R x) :
    fwdDeriv R (fun w => gaussian (μ w) σ (x w))
    =
    fun w dw =>
      let μdμ := fwdDeriv R μ w dw
      let xdx := fwdDeriv R x w dw
      let xdx' := σ⁻¹ • (xdx - μdμ)
      let g := gaussian μdμ.1 σ xdx.1
      (g, - ⟪xdx'.1, xdx'.2⟫ * g) := by

  sorry_proof

@[fun_prop]
theorem gaussian.arg_μσx.CDifferentiableAt_rule (w : W)
    (μ : W → U) (σ : W → R) (x : W → U)
    (hμ : CDifferentiableAt R μ w) (hσ : CDifferentiableAt R σ w) (hx : CDifferentiableAt R x w)
    (hσ' : σ w ≠ 0) :
    CDifferentiableAt R (fun w => gaussian (μ w) (σ w) (x w)) w := by

  unfold gaussian
  fun_prop (disch:=first | assumption | sorry_proof)

@[fun_prop]
theorem gaussian.arg_μσx.CDifferentiable_rule
    (μ : W → R) (σ : W → R) (x : W → R)
    (hμ : CDifferentiable R μ) (hσ : CDifferentiable R σ) (hx : CDifferentiable R x)
    (hσ' : ∀ w, σ w ≠ 0) :
    CDifferentiable R (fun w => gaussian (μ w) (σ w) (x w))

:= by

  intro x
  fun_prop (disch:=aesop)


import Mathlib.MeasureTheory.Measure.GiryMonad
import Mathlib.MeasureTheory.Decomposition.Lebesgue

import SciLean.Core.FunctionPropositions
import SciLean.Util.SorryProof

open MeasureTheory ENNReal

namespace SciLean

local notation "∞" => (⊤ : ℕ∞)

-- variable
--   {R} [RealScalar R]
--   {W} [Vec R W] [Module ℝ W]
--   {X}

-- TODO: Move somewhere
class TCOr (A B : Sort _) where
  val : PSum A B

set_option checkBinderAnnotations false in
instance {A B} [inst : A] : TCOr A B where
  val := .inl inst

set_option checkBinderAnnotations false in
instance {A B} [inst : B] : TCOr A B where
  val := .inr inst

structure IsTestFunction
    {R : Type u} [RealScalar R]
    {X : Type v} [TopologicalSpace X] [space : TCOr (Vec R X) (DiscreteTopology X)]
    (f : X → R) : Prop where
  is_smooth :
    match space.val with
    | .inl _ => ContCDiff R ∞ f
    | .inr _ => True
  has_compact_support : HasCompactSupport f

structure TestFunctionSpace (R : Type u) [RealScalar R] (X : Type v)
    [TopologicalSpace X] [space : TCOr (Vec R X) (DiscreteTopology X)] where
  toFun : X → R
  is_test_fun : IsTestFunction toFun

notation "𝒟" X => TestFunctionSpace defaultScalar% X

@[app_unexpander TestFunctionSpace] def unexpandTestFunctionSpace : Lean.PrettyPrinter.Unexpander
  | `($(_) $_ $X) => `(𝒟 $X)
  | _ => throw ()

variable
  {R} [RealScalar R]
  [TopologicalSpace X]
  [space : TCOr (Vec R X) (DiscreteTopology X)] -- here we are getting topology in `Vec R X` for the second time :( this will cause issues ...

set_default_scalar R

instance : FunLike (𝒟 X) X R where
  coe f x := f.toFun x
  coe_injective' := sorry_proof

instance : Add (𝒟 X) := ⟨fun f g : 𝒟 X => ⟨fun x => f x + g x, sorry_proof⟩⟩
instance : Sub (𝒟 X) := ⟨fun f g : 𝒟 X => ⟨fun x => f x - g x, sorry_proof⟩⟩
instance : SMul R (𝒟 X) := ⟨fun r f => ⟨fun x => r * f x, sorry_proof⟩⟩
instance : Neg (𝒟 X) := ⟨fun f => ⟨fun x => - f x, sorry_proof⟩⟩
instance : Zero (𝒟 X) := ⟨⟨fun x => - 0, sorry_proof⟩⟩

instance : UniformSpace (𝒟 X) := sorry
instance : Vec R (𝒟 X) := Vec.mkSorryProofs

----------------------------------------------------------------------------------------------------
-- Test function approximation ---------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

/-- Restrict function `f` by a ball of radius `n` and do smooth by convolution with radius `1/n`.

If `X` is discrete space then we do not need to do smoothing. -/
opaque testFunApprox (n : ℕ) (f : X → R) : 𝒟 X

----------------------------------------------------------------------------------------------------
-- Properties --------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem TestFunctionSpace.eval_CDifferentiable (x : X) :
    CDifferentiable R (fun φ : 𝒟 X => φ x) := sorry_proof

@[fun_prop]
theorem TestFunctionSpace.eval_IsLinearMap (x : X) :
    IsLinearMap R (fun φ : 𝒟 X => φ x) := sorry_proof

@[fun_prop]
theorem TestFunctionSpace.eval_IsSmoothLinearMap (x : X) :
    IsSmoothLinearMap R (fun φ : 𝒟 X => φ x) := by constructor <;> fun_prop

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

-- @[fun_prop]
-- theorem TestFunctionSpace.eval_CDifferentiable_rule :
--     CDifferentiable R (fun (φx : (𝒟 X)×X) => φx.1 φx.2) := sorry_proof

@[fun_prop]
theorem TestFunctionSpace.eval_CDifferentiableAt_rule (w : W)
    (φ : W → 𝒟 X) (x : W → X)
    (hφ : CDifferentiableAt R φ w) (hx : CDifferentiableAt R x w) :
    CDifferentiableAt R (fun w : W => (φ w) (x w)) w := sorry_proof

@[fun_prop]
theorem TestFunctionSpace.eval_CDifferentiable_rule'
    (φ : W → 𝒟 X) (x : W → X)
    (hφ : CDifferentiable R φ) (hx : CDifferentiable R x) :
    CDifferentiable R (fun w : W => (φ w) (x 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' φ


import Mathlib.MeasureTheory.Measure.MeasureSpaceDef
import Mathlib.MeasureTheory.Measure.Restrict
import Mathlib.MeasureTheory.Measure.Dirac
-- import Mathlib.Data.RCLike.Basic
-- import Mathlib.Analysis.LocallyConvex.Basic
-- import Mathlib.Topology.Algebra.Module.LocallyConvex

import SciLean.Core.FunctionTransformations
import SciLean.Util.SorryProof

open MeasureTheory

namespace SciLean

/-- Convenient integral - the integral I need :)
It should be Bochner integral but it should integrate function valued function point wise i.e.
```
∫ x, fun y => f x y = fun y => ∫ x, f x y
```
where rhs can be understoods as Bochenr integral and lhs defined thie `cintegral`. -/
noncomputable
opaque cintegral {α} [MeasurableSpace α] {X} [AddCommGroup X] [Module ℝ X]
  -- dragging along all of these typeclasses is really really annoying
  -- [AddCommGroup X] [TopologicalSpace X] [TopologicalAddGroup X] [Module ℝ X] [LocallyConvexSpace ℝ X]
    (f : α → X) (μ : Measure α) : X := 0

opaque CIntegrable {α} [MeasurableSpace α] {X} [AddCommGroup X] [Module ℝ X]
    (f : α → X) (μ : Measure α) : Prop

open Lean Parser  Term
syntax "∫' " funBinder (" in " term)? ", " term:60 (" ∂" term:70)? : term

macro_rules
| `(∫' $x:funBinder, $b) => `(cintegral (fun $x => $b) (by volume_tac))
| `(∫' $x:funBinder, $b ∂$μ) => `(cintegral (fun $x => $b) $μ)
| `(∫' $x:funBinder in $set, $b) => `(cintegral (fun $x => $b) (Measure.restrict (by volume_tac) $set))
| `(∫' $x:funBinder in $set, $b ∂$μ) => `(cintegral (fun $x => $b) (Measure.restrict $μ $set))

@[app_unexpander cintegral] def unexpandCIntegral : Lean.PrettyPrinter.Unexpander

  | `($(_) $f:term volume) =>
    match f with
    | `(fun $x => $b) => `(∫' $x, $b)
    | `(fun $x $xs* => $b) => `(∫' $x, fun $xs* => $b)
    | _ => throw ()

  | `($(_) $f:term $μ) =>
    match f, μ with
    -- |      `(fun $x => $b), `(Measure.restrict volume $A) => `(∫' $x in $A, $b)
    -- | `(fun $x $xs* => $b), `(Measure.restrict volume $A) => `(∫' $x in $A, (fun $xs* => $b))
    |      `(fun $x => $b), `(Measure.restrict $μ $A) => `(∫' $x in $A, $b ∂$μ)
    | `(fun $x $xs* => $b), `(Measure.restrict $μ $A) => `(∫' $x in $A, (fun $xs* => $b) ∂$μ)
    -- |      `(fun $x => $b), `(volume) => `(∫' $x, $b)
    -- | `(fun $x $xs* => $b), `(volume) => `(∫' $x, (fun $xs* => $b))
    |      `(fun $x => $b), _ => `(∫' $x, $b ∂$μ)
    | `(fun $x $xs* => $b), _ => `(∫' $x, (fun $xs* => $b) ∂$μ)
    | _, _ => throw ()

  | _ => throw ()

----------------------------------------------------------------------------------------------------
-- Algebra -----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
section Algebra

variable
  {α} [MeasurableSpace α]
  {β} [MeasurableSpace β]
  {X} [AddCommGroup X] [Module ℝ X]
  {Y} [AddCommGroup Y] [Module ℝ Y]

@[simp, ftrans_simp]
theorem cintegral_zero (μ : Measure α): ∫' _, (0 : X) ∂μ = 0 := sorry_proof

@[add_pull] -- @[integral_push]
theorem cintegral_add {f g : α → X} {μ : Measure α}
    (hf : CIntegrable f μ) (hg : CIntegrable g μ) :
    ∫' x, f x + g x ∂μ = ∫' x, f x ∂μ + ∫' x, g x ∂μ := sorry_proof

@[add_push] -- @[integral_push]
theorem cintegral_add' {f g : α → X} {μ : Measure α}
    (hf : CIntegrable f μ) (hg : CIntegrable g μ) :
    ∫' x, f x ∂μ + ∫' x, g x ∂μ = ∫' x, f x + g x ∂μ  := sorry_proof

@[smul_pull]-- @[integral_push]
theorem cintegral_smul {R} [Semiring R] [Module R X] {f : α → X} (r : R) :
    ∫' x, r • f x ∂μ = r • ∫' x, f x ∂μ := sorry_proof

@[smul_push] -- @[integral_pull]
theorem cintegral_smul' {R} [Semiring R] [Module R X] {f : α → X} (r : R) :
    r • ∫' x, f x ∂μ = ∫' x, r • f x ∂μ := sorry_proof

@[add_pull] -- @[integral_push]
theorem cintegral_add_measures {f : α → X} {μ ν : Measure α}
    (hμ : CIntegrable f μ) (hν : CIntegrable f ν) :
    ∫' x, f x ∂(μ+ν) = ∫' x, f x ∂μ + ∫' x, f x ∂ν := sorry_proof

@[simp, ftrans_simp]
theorem cintegral_dirac {f : α → X} (p : α) :
    ∫' x, f x ∂(Measure.dirac p) = f p := sorry_proof

@[simp,ftrans_simp]
theorem cintegral_unit {f : Unit → X} :
    ∫' x, f x  = f () := sorry_proof

-- @[integral_push]
theorem cintegral_prod_mk {f : α → X} {g : α → Y} {μ : Measure α} :
    ∫' x, (f x, g x) ∂μ = (∫' x, f x ∂μ, ∫' x, g x ∂μ) := sorry_proof

-- @[integral_pull]
theorem cintegral_prod_mk' {f : α → X} {g : α → Y} {μ : Measure α} :
    (∫' x, f x ∂μ, ∫' x, g x ∂μ) = ∫' x, (f x, g x) ∂μ := sorry_proof

theorem cintegral_lambda {α} (f : α → β → X) (μ : Measure β) :
    (fun x => ∫' y, f x y ∂μ)
    =
    ∫' y, (fun x => f x y) ∂μ := sorry_proof

theorem cintegral.arg_f.push_lambda {α} (f : α → β → X) (μ : Measure β) :
    (fun x => ∫' y, f x y ∂μ)
    =
    ∫' y, (fun x => f x y) ∂μ := sorry_proof

@[simp, ftrans_simp]
theorem cintegral_measure_map (g : β → X) (f : α → β) (μ : Measure α) :
    (∫' y, g y ∂(μ.map f))
    =
    ∫' x, g (f x) ∂μ := sorry_proof

end Algebra

----------------------------------------------------------------------------------------------------
-- Integral simproc --------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

section Simplifier

variable {α} [MeasurableSpace α] {X} [AddCommGroup X] [Module ℝ X]

theorem integral_in_set_simp
    (f g : α → X) (A : Set α) (h : ∀ x, x ∈ A → f x = g x) (μ : Measure α) :
    ∫' x in A, f x ∂μ
    =
    ∫' x in A, g x ∂μ := sorry_proof

open Lean Meta in
simproc_decl integral_simproc (cintegral _ (Measure.restrict _ _)) := fun e => do
  let .app e  μ' := e  | return .continue
  let .app _  f  := e  | return .continue
  let .app μ' A  := μ' | return .continue
  let .app _  μ  := μ' | return .continue

  let .forallE xName xType _ _ ← inferType f | return .continue

  withLocalDeclD xName xType fun x => do
  withLocalDeclD `h (← mkAppM ``Membership.mem #[x,A]) fun h => do
    let fx := f.beta #[x]
    let r ← Simp.simp fx

    let h' ← mkLambdaFVars #[x,h] (← r.getProof)
    let g ← mkLambdaFVars #[x] r.expr

    let prf ← mkAppM ``integral_in_set_simp #[f,g,A,h',μ]
    let rhs := (← inferType prf).appArg!

    return .visit { expr := rhs, proof? := .some prf }

end Simplifier

----------------------------------------------------------------------------------------------------
-- Domain simps  -----------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

section DomainSimps

variable
  {R} [RealScalar R] [@DecidableRel R (· ≤ ·)]
  {X} [MeasureSpace X]
  {Y} [AddCommGroup Y] [Module ℝ Y]
  {Z} [AddCommGroup Z] [Module ℝ Z]

-- we prefer specifying sets with `0 ≤ ...`
-- TODO: Add seme regularity conditions on φ and ψ
-- @[simp,ftrans_simp]
theorem split_integral_of_ite (φ ψ : X → R) (f g : X → Y) :
    (∫' x, if ψ x ≤ φ x then f x else g x)
    =
    (∫' x in {x' | ψ x' ≤ φ x'}, f x)
    +
    (∫' x in {x' | φ x' < ψ x'}, g x) := sorry_proof

-- @[simp,ftrans_simp]
-- theorem split_integral_of_ite' (φ ψ : X → R) (f g : X → Y) (h : X → Y → Z):
--     (∫' x, h x (if ψ x ≤ φ x then f x else g x))
--     =
--     (∫' x in {x' | 0 ≤ φ x' - ψ x'}, h x (f x))
--     +
--     (∫' x in {x' | 0 ≤ ψ x' - φ x'}, h x (g x)) := sorry_proof

-- @[simp,ftrans_simp]
theorem split_integral_over_set_of_ite (φ ψ : X → R) (f g : X → Y) (A : Set X) :
    (∫' x in A, if ψ x ≤ φ x then f x else g x)
    =
    (∫' x in {x' | 0 ≤ φ x' - ψ x'} ∩ A, f x)
    +
    (∫' x in {x' | 0 ≤ ψ x' - φ x'} ∩ A, g x) := by sorry_proof

-- attribute [ftrans_simp] MeasureTheory.integral_zero sub_nonneg

end DomainSimps

----------------------------------------------------------------------------------------------------
-- Differentiation ---------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
section Differentiation

variable
  {α} [MeasurableSpace α]
  {β} [MeasurableSpace β]
  {R} [RCLike R]
  {X} [Vec R X] [Module ℝ X]
  {Y} [Vec R Y] [Module ℝ Y]
  {Z} [Vec R Z] [Module ℝ Z]

@[fun_prop]
theorem cintegral.arg_f.CDifferentiable_rule
    (f : X → β → Z) (μ : Measure β) (hf : ∀ x, CDifferentiable R (f · x)) :
    CDifferentiable R (fun x => ∫' y, f x y ∂μ) := by have := hf; sorry_proof

@[fun_trans]
theorem cintegral.arg_f.cderiv_rule
    (f : X → β → Z) (μ : Measure β) (hf : ∀ y, CDifferentiable R (f · y)) :
    (cderiv R  fun x => ∫' y, f x y ∂μ)
    =
    fun x dx => ∫' y, cderiv R (f · y) x dx ∂μ := by have := hf; sorry_proof

@[fun_trans]
theorem cintegral.arg_f.fwdDeriv_rule
    (f : X → β → Z) (μ : Measure β) (hf : ∀ x, CDifferentiable R (f · x)) :
    (fwdDeriv R fun x => ∫' y, f x y ∂μ)
    =
    fun x dx => ∫' y, fwdDeriv R (f · y) x dx ∂μ := by
  unfold fwdDeriv; fun_trans [cintegral_prod_mk]

@[fun_prop]
theorem cintegral.arg_f.IsLinearMap_rule
    (f : X → β → Z) (μ : Measure β) (hf : ∀ y, IsLinearMap R (f · y)) :
    IsLinearMap R (fun x => ∫' y, f x y ∂μ) := by have := hf; sorry_proof

@[fun_prop]
theorem cintegral.arg_f.IsSmoothLinearMap_rule
    (f : X → β → Z) (μ : Measure β) (hf : ∀ y, IsSmoothLinearMap R (f · y)) :
    IsSmoothLinearMap R (fun x => ∫' y, f x y ∂μ)