\import Algebra.Algebra
\import Algebra.Meta
\import Algebra.Module
\import Algebra.Monoid
\import Algebra.Ordered
\import Algebra.Pointed
\import Algebra.Ring
\import Algebra.Ring.FormalSeries
\import Algebra.Ring.FormalSeries.Composition
\import Algebra.Semiring
\import Arith.Nat
\import Arith.Rat
\import Data.Or
\import Function
\import Function.Meta
\import Logic
\import Meta
\import Order.Biordered
\import Order.Lattice
\import Order.LinearOrder
\import Order.PartialOrder
\import Paths
\import Paths.Meta
\import Set.Fin
\import Set.Fin.Instances
\open Monoid (pow)

{- | Formal derivative on power series:
$$(Df)_n := (n+1) \cdot f_{n+1}.$$

If $f(x) = \sum_n f_n x^n$ formally, then $Df$ is the coefficient sequence of
$f'(x) = \sum_n (n+1) f_{n+1} x^n$.

An $A$-linear operation $A[[X]] \to A[[X]]$ satisfying:

* Linearity: $D(f + g) = Df + Dg$, $D(c \cdot f) = c \cdot Df$.

* Leibniz: $D(f \star g) = Df \star g + f \star Dg$.

* Chain rule, when $g_0 = 0$: $D(a\circ g) = ((Da)\circ g) \star Dg$. -}
\func deriv {A : Ring} (a : FSeries A) : FSeries A
  => \lam n => natCoef (suc n) * a (suc n)

\lemma deriv_ide {A : Ring} : deriv {A}  = 
  => ext \lam k => zro_*-right

{- | Leibniz rule: $D(a \star b) = Da \star b\ +\ a \star Db$.
 - At each index $n$, the identity
 - $$(n+1) \sum_{i=0}^{n+1} a_i b_{n+1-i}
 - = \sum_{i=0}^{n+1} \big(i + (n+1-i)\big)\, a_i b_{n+1-i}$$
 - splits the LHS into the two convolution terms $Da \star b$ and $a \star Db$. -}
\lemma deriv-leibniz {A : Ring} {a b : FSeries A}
  : deriv (a * b) = deriv a * b + a * deriv b
  => ext \lam n => later $ A.FinSum-ldistr
  *> A.FinSum-ext (later \lam s => pmap (* _) (pmap natCoef (inv s.) *> A.natCoef_+) *> rdistr)
  *> A.FinSum_+
  *> pmap2 (+)
      (A.FinSum-inj2
        (later \lam s => (suc s., s., pmap suc s.))
        (\lam {s} {s'} p => ext (pmap (\lam x => pred x.) p, pmap __. p))
        (later \lam s c => cases (s. arg addPath) \with {
          | , _ => pmap (* _) natCoefZero *> zro_*-left
          | suc k, q => absurd $ c (k, s., pmap (\lam x => pred (x + _)) (inv q) *> pmap pred s.) $ ext (inv q, idp)
        })
        (\lam s => inv *-assoc))
      (A.FinSum-inj2
        (later \lam s => (s., suc s., pmap suc s.))
        (\lam {s} {s'} p => ext (pmap __. p, pmap (\lam x => pred x.) p))
        (later \lam s c => cases (s. arg addPath) \with {
          | , _ => pmap (* _) natCoefZero *> zro_*-left
          | suc k, q => absurd $ c (s., k, pmap (\lam x => pred (_ + x)) (inv q) *> pmap pred s.) $ ext (idp, inv q)
        })
        (\lam s => equation.monoid {A.natCoef-comm {suc s.} {a s.}}))

-- | Power rule in a commutative ring: $D(B^{n+1}) = (n+1) \cdot B^{n} \cdot DB$.
\lemma deriv-pow-conv-suc {A : CRing} (b : FSeries A) (n : Nat)
  : deriv (pow b (suc n)) = natCoef (suc n) *c (pow b n * deriv b) \elim n
  |  => pmap deriv ide-left *> inv (pmap (*c ( * deriv b)) (natCoefSuc  *> pmap (+ ) natCoefZero *> zro-left) *> ide_*c *> ide-left)
  | suc n =>
    \let | c => pow b (suc n) * deriv b
         | d => pow b n * deriv b
    \in deriv-leibniz {_} {pow b (suc n)} *> pmap (__ * b + c) (deriv-pow-conv-suc b n) *>
      inv (pmap (*c c) (natCoefSuc _) *> *c-rdistr {_} {_} {_} {c}
        *> pmap2 (+) {natCoef _ *c c} {natCoef _ *c d * b} (pmap (natCoef _ *c) {c} {d * b}
          (pmap (* deriv b) {pow b (suc n)} {b * pow b n} (inv Ring.pow-left) *> *-assoc {_} {_} {_} {deriv b} *> *-comm {_} {_} {d})
            *> *c-comm-left {_} {_} {d}) { *c c} {c} ide_*c)

{- | Chain rule for {fcompose}. For $b$ with $b_0 = 0$:
 - $$D(a \circ b) = (Da) \circ b \star Db.$$
 - The formal-power-series counterpart of $(f \circ g)' = f'(g) \cdot g'$.
 - Commutativity of $A$ is required: the factored form
 - $D(B^n) = n \cdot B^{n-1} \cdot DB$ depends on commutativity.
 -}
\lemma deriv-compose {A : CRing} (a b : FSeries A) (b_0 : b  = A.zro) (k : Nat)
  : deriv (fcompose a b) k = (fcompose (deriv a) b * deriv b) k
  => lhs-equals-M a b k *> inv (rhs-equals-M a b b_0 k)
  \where {
    -- | The canonical "middle" form $M_k(a,b) := \sum_{m = 0}^k (D a)_m · (B^{\star m} \star D b)_k$.
    \private \func M-form {A : CRing} (a b : FSeries A) (k : Nat) : A
      => A.FinSum \lam (m : Fin (suc k)) => deriv a m * (pow b m * deriv b) k

    \private \lemma lhs-equals-M {A : CRing} (a b : FSeries A) (k : Nat)
      : deriv (fcompose a b) k = M-form a b k
      => A.BigSum-ldistr {_} {\new Array A (suc (suc k)) \lam n => a n * pow b n (suc k)}
         *> pmap (+ _) (pmap (_ *) zro_*-right *> zro_*-right)
         *> zro-left
         *> inv (A.FinSum=BigSum {_} {\lam m => natCoef (suc k) * (a (suc m) * pow b (suc m) (suc k))})
         *> A.FinSum-ext \lam m => equation.cRing {path \lam i => deriv-pow-conv-suc b m i k}

    \private \lemma rhs-equals-M {A : CRing} (a b : FSeries A) (b_0 : b  = A.zro) (k : Nat)
      : (fcompose (deriv a) b * deriv b) k = M-form a b k
      => A.FinSum-ext (later \lam s => pmap (* _) (fcompose.=FinSum {_} {deriv a}) *> A.FinSum-rdistr)
      *> A.FinSum-double-dep' (later \lam s => FinFin (suc s.))
      *> inv (A.FinSum-inj2 {_} {ProdFin (FinFin (suc k)) (FSeriesRing.PairsFinSet k)}
          (later \lam s => (toFin s. $ fin_< s. <∘l suc<=suc (<=_+ <=-refl zero<=_ <=∘ =_<= s..), s.))
          (IsInj.fromSplit
            (later \lam s => (s., toFin (s. ∧ s..) $ <=_<_suc meet-right))
            (\lam s => ext (idp, fin_nat-inj $ toFin=id *> pmap (∧ _) toFin=id *> MeetSemilattice.meet_<= (<_suc_<= $ fin_< s.))))
          (later \lam s c => \case LinearOrder.dec<_<= s.. s. \with {
            | inl q => equation.cRing {pow-conv-zero-below b_0 q}
            | inr q => absurd $ c (s., toFin s. $ <=_<_suc q) $ ext (fin_nat-inj $ toFin=id *> toFin=id, idp)
          })
          (later \lam s => path (\lam i => natCoef (suc $ toFin=id i) * a (suc $ toFin=id i) * (pow b (toFin=id i) _ * _)) *> inv *-assoc))
      *> inv A.FinSum-double
      *> A.FinSum-ext (\lam m => inv A.FinSum-ldistr)
  }