\import Algebra.Group
\import Algebra.Meta
\import Algebra.Module
\import Algebra.Module.BilinearMap
\import Algebra.Module.LinearMap
\import Algebra.Monoid
\import Algebra.Pointed
\import Algebra.Ring
\import Function.Meta
\import Logic
\import Meta
\import Paths
\import Paths.Meta

\class SyntheticDiff \noclassifying (R : CRing) {
  | \protected isOpenProp : \Prop -> \Prop
  | \protected top-isOpenProp : isOpenProp (\Sigma)
  | \protected meet-isOpenProp {P Q : \Prop} : isOpenProp P -> isOpenProp Q -> isOpenProp (\Sigma P Q)

  \lemma meet3-isOpenProp {P Q R : \Prop} (Po : isOpenProp P) (Qo : isOpenProp Q) (Ro : isOpenProp R) : isOpenProp (\Sigma P Q R)
    => transport isOpenProp (ext (\lam s => (s., s.., s..), \lam s => (s., (s., s.)))) $ meet-isOpenProp Po (meet-isOpenProp Qo Ro)

  \func IsOpenPred {A : \Type} (U : A -> \Prop) : \Prop
    => \Pi (a : A) -> isOpenProp (U a)

  \func IsInvDense (B : \Set) : \Prop
    => \Pi (U : R -> \Prop) -> IsOpenPred U -> \Pi (f g : \Sigma (z : R) (U z) -> B)
        -> (\Pi (p : \Sigma (z : R) (u : U z)) -> Monoid.Inv p. -> f p = g p)
        -> \Pi (p : \Sigma (z : R) (u : U z)) -> f p = g p

  \lemma cancel-lem {B : LModule R} (Bd : IsInvDense B) {U : R -> \Prop} (Uo : IsOpenPred U)
                    (f g : \Sigma (z : R) (U z) -> B)
                    (s : \Pi (p : \Sigma (z : R) (u : U z)) -> p. *c f p = p. *c g p)
                    (p : \Sigma (z : R) (u : U z)) : f p = g p
    => Bd U Uo f g (\lam p p-inv => B.cancel p-inv (s p)) p

  \func IsDiff {A B : LModule R} (Bd : IsInvDense B) {U : A -> \Prop} (Uo : IsOpenPred U) (f : \Sigma (x : A) (U x) -> B) : \Set
    => \Pi (x : \Sigma (a : A) (U a)) (a : A) (t : \Sigma (r : R) (U (x. + r *c a))) -> \Sigma (y : B) (t. *c y = f (x. + t. *c a, t.) - f x)
    \where {
      \use \level levelProp (d d' : IsDiff Bd Uo f) : d = d'
        => ext \lam x a t => ext $ cancel-lem Bd (\lam r => Uo _) _ _ (\lam t => (d x a t). *> inv (d' x a t).) t

      \lemma IsDiff-eq (Bd : IsInvDense B) (Uo : IsOpenPred U) {f : \Sigma (x : A) (U x) -> B} (d : IsDiff Bd Uo f) {x x' : \Sigma (a : A) (U a)} (p : x. = x'.) {a : A} {t : \Sigma (r : R) (U (x. + r *c a))} {t' : \Sigma (r : R) (U (x'. + r *c a))} (q : t. = t'.) : (d x a t). = (d x' a t').
        => path \lam i => (d (p i, prop-dpi _ _ _ i) a (q i, prop-dpi _ _ _ i)).
    }

  \func IsDiffT {A B : LModule R} (Bd : IsInvDense B) (f : A -> B)
    => IsDiff Bd (\lam _ => top-isOpenProp) (\lam p => f p.)

  -- | The derivative is linear
  \func deriv {A B : LModule R} (Bd : IsInvDense B) {U : A -> \Prop} (Uo : IsOpenPred U)
              {f : \Sigma (x : A) (U x) -> B} (d : IsDiff Bd Uo f) (x : \Sigma (a : A) (U a)) : LinearMap A B \cowith
    | func a => (d x a (, sub-lem x)).
    | func-+ {a} {a'} =>
      \have s => cancel-lem Bd {\lam t => \Sigma (U (x. + t *c (a + a'))) (U (x. + t *c a)) (U (x. + t *c a + t *c a'))} (\lam t => meet3-isOpenProp (Uo _) (Uo _) (Uo _)) (\lam t => (d x (a + a') (t., t..)).) (\lam t => (d (x. + t. *c a, t..) a' (t., t..)). + (d x a (t., t..)).)
        (\lam t => (d _ _ (_, _)). *> pmap (- _) (pmap f (ext $ rewrite *c-ldistr $ inv +-assoc) *> inv zro-right *> pmap (_ +) (inv negative-left) *> inv +-assoc) *> +-assoc *> inv (pmap2 (+) (d _ _ (_, _)). (d _ _ (_, _)).) *> inv B.*c-ldistr)
        (, (sub-lem x, sub-lem x, rewrite (A.*c_zro-left, zro-right) (sub-lem x)))
      \in s *> +-comm *> pmap (_ +) (IsDiff.IsDiff-eq Bd Uo d (pmap (x. +) A.*c_zro-left *> zro-right) idp)
    | func-*c {r} {a} =>
      \have s => cancel-lem Bd {\lam t => \Sigma (U (x. + t *c (r *c a))) (U (x. + t * r *c a))} (\lam t => meet-isOpenProp (Uo _) (Uo _)) (\lam t => (d x (r *c a) (t., t..)).) (\lam t => r *c (d x a (t. * r, t..)).)
        (\lam t => later (rewrite ((d x (r *c a) (t., t..))., (d x a (t. * r, t..)).) $ pmap (f __ - f x) $ ext $ pmap (x. +) $ inv *c-assoc) *> *c-assoc)
        (, (sub-lem x, rewrite *c-assoc $ sub-lem x))
      \in s *> pmap (r *c) (IsDiff.IsDiff-eq Bd Uo d idp R.zro_*-left)
    \where
      \lemma sub-lem {A : LModule} {U : A -> \Prop} {a : A} (x : \Sigma (a : A) (U a)) : U (x. +  *c a)
        => transportInv U (pmap (x. +) A.*c_zro-left *> zro-right) x.

  -- | The derivative of a linear map is the map itself
  \func linear_diff {A B : LModule R} (Bd : IsInvDense B) (f : LinearMap A B) : IsDiffT Bd f
    => \lam x a t => (f a, rewrite (+-comm,f.func-+) $ inv (simplify f.func-*c))

  -- | The derivative of a constant map is zero
  \func const_diff {A B : LModule R} (Bd : IsInvDense B) (b : B) : IsDiffT Bd (\lam (_ : A) => b)
    => \lam x a t => (, B.*c_zro-right *> inv negative-right)

  -- | The derivative of a sum is the sum of derivatives
  \func +_diff {A B : LModule R} (Bd : IsInvDense B) {U : A -> \Prop} (Uo : IsOpenPred U)
               {f g : \Sigma (a : A) (U a) -> B} (df : IsDiff Bd Uo f) (dg : IsDiff Bd Uo g) : IsDiff Bd Uo (\lam x => f x + g x)
    => \lam x a t => ((df x a t). + (dg x a t)., B.*c-ldistr *> pmap2 (+) (df x a t). (dg x a t). *> equation *> inv (pmap (_ +) AddGroup.negative_+))

  {- | Differentiable maps are closed under composition.
   -
   -   The chain rule `deriv (o_diff df dg) x a = deriv dg (f x) (deriv df x a)` follows from this by `idp`.
   -}
  \func o_diff {A B C : LModule R} (Bd : IsInvDense B) (Cd : IsInvDense C) {U : A -> \Prop} (Uo : IsOpenPred U) {V : B -> \Prop} (Vo : IsOpenPred V)
               {f : \Sigma (a : A) (U a) -> \Sigma (b : B) (V b)} {g : \Sigma (b : B) (V b) -> C}
               (df : IsDiff Bd Uo (f __).) (dg : IsDiff Cd Vo g) : IsDiff Cd Uo (\lam x => g (f x))
    => \lam x a t => \have s => dg (f x) (df x a t). (t., rewrite (df x a t). $ transport V simplify (f (x. + t. *c a, t.)).)
                     \in (s., s. *> pmap (g __ - _) (ext $ rewrite (df x a t). $ +-comm *> simplify))

  -- | Bilinear maps are differentiable
  \func bilinear_diff {A B C : LModule R} (Cd : IsInvDense C) (m : BilinearMap A B C) : IsDiffT Cd (\lam (p : \Sigma A B) => m p. p.)
    => \lam (x,_) a (t,_) => (m x. a. + m a. x. + t *c m a. a., rewrite (m.linear-left.func-+, m.linear-right.func-+, m.linear-left.func-*c, m.linear-right.func-*c) $
        later (rewrite (negative-left,zro-left,*c-ldistr,*c-ldistr) $ +-assoc *> pmap (_ +) (pmap (_ + _ *c __) (inv m.linear-right.func-*c) *> inv *c-ldistr *> pmap (_ *c) (inv m.linear-right.func-+))) *> pmap (+ _) +-assoc *> +-assoc *> +-comm)
}