\import Algebra.Field
\import Algebra.Group
\import Algebra.Meta
\import Algebra.Monoid
\import Algebra.QModule
\import Algebra.Ring.FormalSeries
\import Algebra.Ring.FormalSeries.Composition
\import Algebra.Ring.FormalSeries.Derivative
\import Algebra.StrictlyOrdered
\import Analysis.PowerSeries
\import Analysis.SeriesComposition
\import Arith.Exp
\import Arith.Real
\import Order.Biordered
\import Order.StrictOrder
\import Algebra.Pointed
\import Algebra.Semiring
\import Analysis.CauchyProduct
\import Analysis.Series
\import Arith.Int
\import Arith.Nat
\import Arith.Rat
\import Arith.Real.UpperReal
\import Function.Meta
\import Logic
\import Meta
\import Order.LinearOrder
\import Order.PartialOrder
\import Paths
\import Paths.Meta
\import Topology.BanachAlgebra
\import Topology.NormedAbGroup
\open Monoid (pow)

-- | Rational Taylor coefficients of $\log(1 + w) = \sum_{n \geq 1} \frac{(-1)^{n - 1}}{n} w^n$.
\func log1p-coef (n : Nat) : Rat \elim n
  | 0 => 0
  | suc n => pow (RatField.negative 1) n * RatField.finv (suc n)

{- | $\vert c_n \vert_{\mathbb Q} \leq 1$ for every $n$, where $c_n =$ `log1p-coef n`.
 -   For $n = 0$, $c_0 = 0$. For $n = m + 1$, $\vert c_n\vert = \vert(-1)^m\vert \cdot \vert 1/(m+1)\vert
 -   = 1 \cdot 1/(m+1) \leq 1/1 = 1.$
 -}
\lemma log1p-coef-abs-<=1 (n : Nat) : RatField.abs (log1p-coef n) <= 1 \elim n
  | 0 => rewrite RatField.abs_zro linarith
  | suc n =>
    \have | n+1>0 : 0 RatField.< suc n
              => fromInt_< (pos<pos NatOrder.zero<suc)
          | abs-eq : RatField.abs (log1p-coef (suc n)) = RatField.finv (suc n)
              => RatField.abs_*
                 *> pmap2 (*)
                      (RatField.abs_pow *> pmap (\lam x => pow x n) (RatField.abs_negative {1} *> RatField.abs_ide) *> RatField.pow_ide)
                      (RatField.abs-ofPos (RatField.finv>=0 (RatField.<=-less n+1>0)))
                 *> ide-left
    \in transportInv (<= 1) abs-eq $ RatField.finv_<= zro<ide (fromInt_<= (pos<=pos (suc<=suc zero<=_))) <=∘ =_<= RatField.finv_ide

{- | $\log(1 + w) := \sum_{n \geq 1} \frac{(-1)^{n - 1}}{n} w^n$ for
 -   $\vert w \vert_A < 1$, in any real Banach algebra $A$.
 -}
\sfunc log1+ {A : RealBanachAlgebra} (w : A) (\property |w|<1 : (A.norm w).U 1) : A
  => powerSeriesSum (\lam n => A.fromRat (log1p-coef n)) w (convertProof |w|<1)
  \where {
    \lemma log1p-conv-at-s-rat {A : RealBanachAlgebra} {s : Rat} (s>=0 : 0 <= s) (s<1 : s < 1)
      : IsConvUpperSeries (\lam n => A.norm (A.fromRat (log1p-coef n)) * pow s n)
      => transport IsConvUpperSeries (ext \lam n => pmap (_ *) (ExUpperRealSemigroup.rat-pow s>=0)) $
          log1p-conv-at-s {A} {s} (ExUpperReal.<=-rat.1 s>=0) s<1

    \lemma log1p-conv-at-s {A : RealBanachAlgebra} {s : ExUpperReal} (s>=0 : 0 <= s) (s<1 : s.U 1)
      : IsConvUpperSeries (\lam n => A.norm (A.fromRat (log1p-coef n)) * ExUpperRealSemigroup.pow s n)
      => series_<= (\lam n => log1p-coef-pow-bound n) (geometric-upper-series-conv s>=0 s<1)
      \where {
        \private \lemma log1p-coef-pow-bound {A : RealBanachAlgebra} (n : Nat) {t : ExUpperReal}
          : A.norm (A.fromRat (log1p-coef n)) * ExUpperRealSemigroup.pow t n <= ExUpperRealSemigroup.pow t n
          => ExUpperRealSemigroup.<=_* (A.norm_fromRat <=∘ ExUpperReal.<=-rat.1 (log1p-coef-abs-<=1 n)) <=-refl
            <=∘ =_<= (ExUpperRealSemigroup.ide-left ExUpperRealSemigroup.pow_>=0)
      }

    \lemma convertProof (|w|<1 : (A.norm w).U 1) : A.dist 0 w <LU convRadius (\lam n => A.fromRat (log1p-coef n))
      => \case U-rounded |w|<1 \with {
        | inP (a,x<a,a<1) => transport (<LU _) norm_dist $ convRadius.<LU-char x<a $ log1p-conv-at-s-rat (<=-less (norm>=0 x<a)) a<1
      }

    -- | $D \log(1 + w) = \tfrac{1}{1 + w}$ at the coefficient level: $(n + 1) \cdot \frac{(-1)^n}{n + 1} = (-1)^n$.
    \lemma deriv-coef-Rat (n : Nat)
      : deriv log1p-coef n = alt_pm n
      => run {
        unfold deriv,
        rewrite {2} *-comm,
        rewriteI *-assoc,
        rewrite $ finv-right (natRat/=0 suc/=0),
        ide-left
      }

    -- | $\log(1 + \cdot)$ as a formal line diffeomorphism.
    \func asFld : FormalLineDiffeomorphism RatField \cowith
      | f => log1p-coef
      | f_0 => idp
      | f1-inv => Monoid.Inv.lmake 1 $ pmap (* 1) (ide-left *> RatField.finv_ide) *> ide-left
  }

-- ============================================================================
-- The substitution-group identities: (exp − 1) ∘ log(1+·) = w  and  log(1+·) ∘ (exp − 1) = w.
-- ============================================================================

-- | The composition $(\exp - 1) \circ \log(1 + \cdot) = w$ as a formal series identity over $\mathbb Q$.
\lemma compose-coef-expm1-log1p (k : Nat) : fcompose exp-1-coef log1p-coef k = subst-ide k
  \elim k
  | 0 => pmap (+ 0) (RatField.zro_*-left {1}) *> zro-right
  | 1 => compose-coef-expm1-eq-exp-pos 1 NatOrder.zero<suc *> exp-log1p-coef-Rat_1
  | suc (suc k) => compose-coef-expm1-eq-exp-pos (suc (suc k)) NatOrder.zero<suc
                   *> exp-log1p-tail=0 k (suc (suc k)) (suc<=suc (suc<=suc zero<=_)) <=-refl
  \where {
    -- | $f_k := [w^k] (\exp \circ \log(1 + \cdot))$ over $\mathbb Q$.
    \private \func exp-log1p-coef-Rat (k : Nat) : Rat
      => fcompose exp.coef-Rat log1p-coef k

    {- | The ODE $D f = f \star \tfrac{1}{1 + w}$ (i.e. $D f_k = \sum_{j = 0}^k f_j \cdot (-1)^{k - j}$).
         From the chain rule $D(\exp \circ \log) = (D \exp \circ \log) \cdot D \log = (\exp \circ \log) \cdot \tfrac{1}{1 + w}$. -}
    \private \lemma expLog1p-coef-Rat-ODE (k : Nat)
      : deriv exp-log1p-coef-Rat k = conv exp-log1p-coef-Rat alt_pm k
      => deriv-compose exp.coef-Rat log1p-coef idp k
      *> rewrite (ext exp.deriv-coef-Rat : deriv exp.coef-Rat = exp.coef-Rat,
                  ext log1+.deriv-coef-Rat : deriv log1p-coef = alt_pm) idp

    \private \lemma exp-log1p-coef-Rat_0 : exp-log1p-coef-Rat 0 = 1
      => pmap (__ * 1 + 0) RatField.finv_ide *> ide-left *> zro-right

    \private \lemma exp-log1p-coef-Rat_1 : exp-log1p-coef-Rat 1 = 1
      => pmap2 (\lam a b => exp.coef-Rat 0 * a + (exp.coef-Rat 1 * b + 0))
          (pow-conv-zero-pos NatSemiring.zro<ide)
          (fseries-apply (ide-left {FSeriesRing RatField} {log1p-coef}))

    \private \lemma exp-log1p-coef-Rat_2 : exp-log1p-coef-Rat 2 = 0
      => RatField.nonZero-left (natRat/=0 suc/=0) $
          expLog1p-coef-Rat-ODE 1
          *> conv_BigSum exp-log1p-coef-Rat alt_pm 1
          *> pmap2 (\lam a b => a * alt_pm 1 + (b * alt_pm 0 + 0)) exp-log1p-coef-Rat_0 exp-log1p-coef-Rat_1

    {- | The composition coefficient identity at index $k$: for $k > 0$,
    the $n=0$ term of `compose-coef` vanishes because $(\log(1+w))^0_k = [k = 0]$ is zero. -}
    \private \lemma compose-coef-expm1-eq-exp-pos (k : Nat) (k>0 : 0 < k)
      : fcompose exp-1-coef log1p-coef k = fcompose exp.coef-Rat log1p-coef k
      => RatField.BigSum-ext {suc k}
          {\new Array Rat (suc k) (\lam n => exp-1-coef n * pow {FSeriesRing RatField} log1p-coef n k)}
          {\new Array Rat (suc k) (\lam n => exp.coef-Rat n * pow {FSeriesRing RatField} log1p-coef n k)}
          \case \elim __ \with {
            | 0 => zro_*-left *> inv (pmap (_ *) (pow-conv-zero-pos k>0) *> zro_*-right {_} {exp.coef-Rat 0})
            | suc n' => idp
          }

    \private \lemma exp-log1p-tail=0 (n j : Nat) (j>=2 : 2 <= j) (j<=2 : j <= suc (suc n)) : exp-log1p-coef-Rat j = 0
      \elim n
      | 0 => rewrite (<=-antisymmetric j<=2 j>=2) exp-log1p-coef-Rat_2
      | suc n => \case NatSemiring.trichotomy j (suc (suc n)) \with {
        | less j<top => exp-log1p-tail=0 n j j>=2 $ <_suc_<= j<top <=∘ id<=suc
        | equals j-eq => rewrite j-eq $ exp-log1p-tail=0 n (suc (suc n)) (suc<=suc (suc<=suc zero<=_)) <=-refl
        | greater j>top => rewrite (<=-antisymmetric j<=2 (suc_<_<= j>top)) $ exp-log1p-step n (exp-log1p-tail=0 n)
      }
      \where {
        -- | The tail $\sum_{i = 2}^{n + 2} f_i \cdot (-1)^{n + 2 - i}$ vanishes by IH: every term has factor $f_i = 0$.
        \private \lemma exp-log1p-step-tail-zero (n : Nat)
                                                 (IH : \Pi (j : Nat) -> 2 <= j -> j <= suc (suc n) -> exp-log1p-coef-Rat j = 0)
          : AddMonoid.BigSum (\new Array Rat (suc n) (exp-log1p-tail-term n)) = 0
          => AddMonoid.BigSum_zro {RatField} {\new Array Rat (suc n) (exp-log1p-tail-term n)} \lam j =>
            pmap (* _) (IH (suc (suc j)) (suc<=suc (suc<=suc zero<=_)) (suc<=suc (suc<=suc (<_suc_<= (fin_< j))))) *> zro_*-left
          \where {
            \private \func exp-log1p-tail-term (n : Nat) (i : Fin (suc n)) : Rat
              => exp-log1p-coef-Rat (suc (suc i)) * alt_pm (n -' i)
          }

      {- | Inductive step: given $f_j = 0$ for $2 \leq j \leq n + 2$, conclude $f_{n+3} = 0$.
           The ODE at $k = n + 2$ gives $(n + 3) f_{n + 3} = \sum_{j = 0}^{n + 2} f_j (-1)^{n + 2 - j}$;
           substituting $f_0 = f_1 = 1$, $f_j = 0$ for $j \geq 2$ leaves $(-1)^{n + 2} + (-1)^{n + 1} = 0$. -}
      \private \lemma exp-log1p-step (n : Nat)
                                     (IH : \Pi (j : Nat) -> 2 NatSemiring.<= j -> j NatSemiring.<= suc (suc n) -> exp-log1p-coef-Rat j = 0)
        : exp-log1p-coef-Rat (suc (suc (suc n))) = 0
        => \have
            | ode => expLog1p-coef-Rat-ODE (suc (suc n))
            | unfold-conv => conv_BigSum exp-log1p-coef-Rat (alt_pm {RatField}) (suc (suc n))
            | collapse {a b G G' T : Rat} (g0 : a = 1) (g1 : b = 1) (gp : G = negative G') (tz : T = 0)
                       : a * G + (b * G' + T) = 0
                       => equation.cRing {g0, g1, gp, tz}
           \in RatField.nonZero-left (natRat/=0 suc/=0) $
                 ode *> unfold-conv
                     *> collapse exp-log1p-coef-Rat_0 exp-log1p-coef-Rat_1
                                 (alt_pm-suc {_} {suc n})
                                 (exp-log1p-step-tail-zero n IH)
    }
}

-- ============================================================================
-- Analytic lifting to a Banach algebra, via compose-converges-{abs,radius}.
-- ============================================================================

-- | $\exp(\log(1 + w)) = 1 + w$ for $\vert w \vert_A < 1$ in any commutative real Banach algebra.
\lemma exp-log1p {A : RealBanachAlgebra} (w : A) (wc : A.IsCentral w) (|w|<1 : (A.norm w).U 1)
  : exp (log1+ w |w|<1) = 1 + w
  => \let | compose-result : IsSeriesSum
              (\lam k => fcompose (\lam n => A.fromRat (exp-1-coef n)) (\lam n => A.fromRat (log1p-coef n)) k * pow w k)
              (exp (log1+ w |w|<1) - 1)
              => compose-converges-abs wc A.fromRat_zro exp-1-coef.ps-conv
                   (powerSeries-absConv $ log1+.log1p-conv-at-s norm>=0 |w|<1)
                   (transportInv (IsSeriesSum _) (\peval log1+ w |w|<1) powerSeriesSum-isSum)
                   (exp-1-coef.seriesSum _)
          | coef-id : (\lam k => fcompose (\lam n => A.fromRat (exp-1-coef n)) (\lam n => A.fromRat (log1p-coef n)) k * pow w k)
                    = (\lam k => subst-ide k * pow w k)
              => path (\lam i k => coef-eq-helper-expm1-log1p k i * pow w k)
          | minus-eq : exp (log1+ w |w|<1) - A.ide = w
              => A.limit-unique (transport (\lam S => IsSeriesSum S _) coef-id compose-result) (subst-ide-seriesSum w)
     \in inv zro-right *> pmap (_ +) (inv negative-left) *> inv +-assoc *> pmap (+ 1) minus-eq *> +-comm
  \where {
    \lemma subst-ide_fromRat {A : QAlgebra} {k : Nat} : A.fromRat (subst-ide k) = subst-ide k \elim k
      | 0 => A.fromRat_zro
      | 1 => A.fromRat_ide
      | suc (suc k) => A.fromRat_zro

    \private \lemma coef-eq-helper-expm1-log1p {A : RealBanachAlgebra} (k : Nat)
      : fcompose (\lam n => A.fromRat (exp-1-coef n)) (\lam n => A.fromRat (log1p-coef n)) k = subst-ide k
      => run {
        *> pmap A.fromRat (compose-coef-expm1-log1p k) *> subst-ide_fromRat,
        \have A0 => A.toRatModule.*c_BigSum-rdistr {\new Array Rat (suc k) \lam n => exp-1-coef n * pow-conv log1p-coef n k},
        later $ rewrite A0,
        AddMonoid.BigSum-ext {A} {suc k}
          {\new Array A (suc k) \lam n => A.fromRat (exp-1-coef n) * pow-conv (\lam m => A.fromRat (log1p-coef m)) n k}
          {\new Array A (suc k) \lam n => A.fromRat (exp-1-coef n * pow-conv log1p-coef n k)},
          \lam i,
          rewriteI A.fromRat_*,
          pmap (_ *) fromRat-pow-conv
      }
  }