-- Möbius inversion: if g(n) = ∑_{d ∣ n} f(d) for every n ≥ 1 then
-- f(n) = ∑_{d ∣ n} μ(d) · g(n / d) for every n ≥ 1, and conversely, proved
-- by interchanging the double sum over divisors and the key identity
-- ∑_{d ∣ m} μ(d) = [m = 1]
-- Part of Cumulus, the Grace17 mathematics library.
-- Copyright (c) 2026 Grant Passmore.
-- Released under the Apache License 2.0. See LICENSE in the repository root.
-- Authors: Grant Passmore

module Grace.Cumulus.Combinatorics.MoebiusInversion
import Moebius
import DivisorInterchange

-- The weight (d, e) ↦ a(d) · b(e) on pairs of naturals.
defgraph productWeights (a b : Set) (q ∈ ℕ × ℕ) : ℤ :=
  integerMulValue(a(firstProjection(ℕ, ℕ)(q)), b(secondProjection(ℕ, ℕ)(q)))
requires weightsA : a : ℕ → ℤ
requires weightsB : b : ℕ → ℤ
proof
  have firstFunction := firstProjection.function at ℕ, ℕ
  have secondFunction := secondProjection.function at ℕ, ℕ
  have firstMember := FunctionBetween.apply_closed with firstFunction, member
  have secondMember := FunctionBetween.apply_closed with secondFunction, member
  have left := FunctionBetween.apply_closed with weightsA, firstMember
  have right := FunctionBetween.apply_closed with weightsB, secondMember
  exact integerMulValue.integer at a(firstProjection(ℕ, ℕ)(q)),
    b(secondProjection(ℕ, ℕ)(q)) with *

theorem productWeights.pair (a b d e : Set)
  requires weightsA : a : ℕ → ℤ
  requires weightsB : b : ℕ → ℤ
  requires dMember : d ∈ ℕ
  requires eMember : e ∈ ℕ
  shows productWeights(a, b)((d, e)) = integerMulValue(a(d), b(e))
proof
  have pairMember := product.pair with dMember, eMember
  have value := productWeights.value at a, b, (d, e) with weightsA, weightsB,
    pairMember
  have firstValue := firstProjection.pair at ℕ, ℕ, d, e with dMember, eMember
  have secondValue := secondProjection.pair at ℕ, ℕ, d, e with dMember, eMember
  rw [value, firstValue, secondValue]
  rfl

-- Row d of the pairs μ(d) · f(e) sums to μ(d) · ∑_{e ∣ n/d} f(e).
theorem moebiusRow.value (f : Set) (n : Nat) (d : Set)
  requires weightsF : f : ℕ → ℤ
  requires nonzero : n ≠ 0
  requires member : d ∈ divisors(n)
  shows rowDivisorSums(productWeights(moebius, f), n)(d) =
    integerMulValue(moebius(d), divisorSum(quotientNat(n, d), f))
proof
  have moebiusFunction := moebius.function
  have pairWeights := productWeights.function at moebius, f with
    moebiusFunction, weightsF
  have naturals := divisors.subset_naturals at n
  have dNat := divisors.natural with member
  have dMember := naturals with d, member
  have quotientNatural := quotientNat.natural at d, n
  have innerNaturals := divisors.subset_naturals at quotientNat(n, d)
  have rowValue := rowDivisorSums.value at productWeights(moebius, f), n, d
    with pairWeights
  have rows := rowTerms.function at productWeights(moebius, f), d with
    pairWeights, dMember
  have scaledValues : ∀ e ∈ divisors(quotientNat(n, d)),
      rowTerms(productWeights(moebius, f), d)(e) =
        integerMulValue(moebius(d), f(e))
    intros e he
    have eMember := innerNaturals with e, he
    have termValue := rowTerms.value at productWeights(moebius, f), d, e with
      pairWeights, dMember
    have pairValue := productWeights.pair at moebius, f, d, e with
      moebiusFunction, weightsF, dMember, eMember
    cc
  have dInteger := moebius.integer at d
  have scaled := divisorSum.scale at quotientNat(n, d), f,
    rowTerms(productWeights(moebius, f), d), moebius(d) with weightsF, rows,
      scaledValues
  rw [rowValue, scaled]
  rfl

-- Column e of the pairs μ(d) · f(e) sums to f(e) · ∑_{d ∣ n/e} μ(d).
theorem moebiusColumn.value (f : Set) (n : Nat) (e : Set)
  requires weightsF : f : ℕ → ℤ
  requires nonzero : n ≠ 0
  requires member : e ∈ divisors(n)
  shows columnDivisorSums(productWeights(moebius, f), n)(e) =
    integerMulValue(f(e), divisorSum(quotientNat(n, e), moebius))
proof
  have moebiusFunction := moebius.function
  have pairWeights := productWeights.function at moebius, f with
    moebiusFunction, weightsF
  have naturals := divisors.subset_naturals at n
  have eNat := divisors.natural with member
  have eMember := naturals with e, member
  have eInteger := FunctionBetween.apply_closed with weightsF, eMember
  have quotientNatural := quotientNat.natural at e, n
  have innerNaturals := divisors.subset_naturals at quotientNat(n, e)
  have columnValue := columnDivisorSums.value at productWeights(moebius, f), n,
    e with pairWeights
  have columns := columnTerms.function at productWeights(moebius, f), e with
    pairWeights, eMember
  have scaledValues : ∀ d ∈ divisors(quotientNat(n, e)),
      columnTerms(productWeights(moebius, f), e)(d) =
        integerMulValue(f(e), moebius(d))
    intros d hd
    have dMember := innerNaturals with d, hd
    have dNat := nat.omega.backward with dMember
    have dInteger := moebius.integer at d
    have termValue := columnTerms.value at productWeights(moebius, f), e, d
      with pairWeights, eMember
    have pairValue := productWeights.pair at moebius, f, d, e with
      moebiusFunction, weightsF, dMember, eMember
    rw [termValue, pairValue]
    ring
  have scaled := divisorSum.scale at quotientNat(n, e), moebius,
    columnTerms(productWeights(moebius, f), e), f(e) with moebiusFunction,
      columns, scaledValues
  rw [columnValue, scaled]
  rfl

theorem moebiusColumn.vanishes (f : Set) (n : Nat) (e : Set)
  requires weightsF : f : ℕ → ℤ
  requires nonzero : n ≠ 0
  requires member : e ∈ divisors(n)
  requires different : e ≠ n
  shows columnDivisorSums(productWeights(moebius, f), n)(e) = integerZero
proof
  have value := moebiusColumn.value with weightsF, nonzero, member
  have eNat := divisors.natural with member
  have factor := divisors.divides with member
  have eNonzero := divides.nonzero with factor, nonzero
  have quotientNonzero := divides.quotient_ne_zero with nonzero, factor
  have quotientNatural := quotientNat.natural at e, n
  have quotientNotOne : quotientNat(n, e) ≠ 1
    intro one
    have equation := divides.quotient_equation with eNonzero, factor
    rw one at equation
    have unit := mul.right_unit at e
    have same : e = n by cc
    exact different with same
  have characterization := LessEqual.one_iff_ne_zero at quotientNat(n, e)
  have positive := characterization.backward with quotientNonzero
  have alternatives := LessEqual.cases with positive
  have aboveOne : 1 < quotientNat(n, e)
    cases alternatives with strict | same
    . exact strict
    . exfalso
      apply quotientNotOne
      symm
      exact same
  have key := moebius.sum_divisors with aboveOne
  have naturals := divisors.subset_naturals at n
  have eMember := naturals with e, member
  have eInteger := FunctionBetween.apply_closed with weightsF, eMember
  rw [value, key]
  ring

theorem moebiusColumn.at_self (f : Set) (n : Nat)
  requires weightsF : f : ℕ → ℤ
  requires nonzero : n ≠ 0
  shows columnDivisorSums(productWeights(moebius, f), n)(n) = f(n)
proof
  have self := divisors.self_member at n
  have value := moebiusColumn.value with weightsF, nonzero, self
  have selfQuotient : quotientNat(n, n) = 1
    have unit := mul.right_unit at n
    have reduced := quotientNat.mul_left at n, 1 with nonzero
    cc
  have one := moebius.sum_one
  have nNat : Nat(n) by typecheck
  have nMember := nat.omega.forward with nNat
  have nInteger := FunctionBetween.apply_closed with weightsF, nMember
  have oneInteger := integerOne.integer
  rw [value, selfQuotient, one]
  ring

-- Möbius inversion: g(n) = ∑_{d ∣ n} f(d) for all n ≥ 1 gives
-- f(n) = ∑_{d ∣ n} μ(d) · g(n / d) for all n ≥ 1.
theorem moebius.inversion (f g : Set)
  requires weightsF : f : ℕ → ℤ
  requires weightsG : g : ℕ → ℤ
  requires sums : ∀ n ∈ ℕ, n ≠ 0 → g(n) = divisorSum(n, f)
  shows ∀ n ∈ ℕ, n ≠ 0 → f(n) = divisorSum(n, dirichletTerms(moebius, g, n))
proof
  intros n nMember nonzero
  have nNat := nat.omega.backward with nMember
  have moebiusFunction := moebius.function
  have pairWeights := productWeights.function at moebius, f with
    moebiusFunction, weightsF
  have naturals := divisors.subset_naturals at n
  have terms := dirichletTerms.function at moebius, g, n with moebiusFunction,
    weightsG
  have rows := rowDivisorSums.function at productWeights(moebius, f), n with
    pairWeights
  have columns := columnDivisorSums.function at productWeights(moebius, f), n
    with pairWeights
  have rowsAgree : ∀ d ∈ divisors(n), dirichletTerms(moebius, g, n)(d) =
      rowDivisorSums(productWeights(moebius, f), n)(d)
    intros d hd
    have dNat := divisors.natural with hd
    have dMember := naturals with d, hd
    have factor := divisors.divides with hd
    have quotientNonzero := divides.quotient_ne_zero with nonzero, factor
    have quotientNatural := quotientNat.natural at d, n
    have quotientMember := nat.omega.forward with quotientNatural
    have termValue := dirichletTerms.value at moebius, g, n, d with
      moebiusFunction, weightsG
    have rowValue := moebiusRow.value with weightsF, nonzero, hd
    have gValue := sums with quotientNat(n, d), quotientMember, quotientNonzero
    rw [termValue, rowValue, gValue]
    rfl
  have leftSum := divisorSum.congr with terms, rows, rowsAgree
  have interchange := divisorSum.interchange at productWeights(moebius, f), n
    with nonzero, pairWeights
  have vanishing : ∀ e ∈ divisors(n), e ≠ n →
      columnDivisorSums(productWeights(moebius, f), n)(e) = integerZero
    intros e he different
    exact moebiusColumn.vanishes with weightsF, nonzero, he, different
  have single := divisorSum.single with columns, vanishing
  have atSelf := moebiusColumn.at_self with weightsF, nonzero
  rw [leftSum, interchange, single, atSelf]
  rfl

-- Successive quotients by two divisors commute.
theorem quotientNat.swap (n d c : Nat)
  requires nonzero : n ≠ 0
  requires factor : d ∣ n
  requires inner : c ∣ quotientNat(n, d)
  shows quotientNat(quotientNat(n, d), c) = quotientNat(quotientNat(n, c), d)
proof
  have dNonzero := divides.nonzero with factor, nonzero
  have quotientNonzero := divides.quotient_ne_zero with nonzero, factor
  have cNonzero := divides.nonzero with inner, quotientNonzero
  have outerEquation := divides.quotient_equation with dNonzero, factor
  have innerEquation := divides.quotient_equation with cNonzero, inner
  have innerNatural := quotientNat.natural at c, quotientNat(n, d)
  have rearranged : d * (c * quotientNat(quotientNat(n, d), c)) =
      c * (d * quotientNat(quotientNat(n, d), c)) by ring
  have shape : n = c * (d * quotientNat(quotientNat(n, d), c)) by cc
  have outer := quotientNat.mul_left at c, d * quotientNat(quotientNat(n, d),
    c) with cNonzero
  have reduced := quotientNat.mul_left at d, quotientNat(quotientNat(n, d), c)
    with dNonzero
  cc

-- The weight (d, c) ↦ μ(c) · g((n / d) / c): the Möbius summands at n / d.
defgraph nestedTerms (g : Set) (n : Nat) (q ∈ ℕ × ℕ) : ℤ :=
  dirichletTerms(moebius, g, quotientNat(n,
    firstProjection(ℕ, ℕ)(q)))(secondProjection(ℕ, ℕ)(q))
requires weightsG : g : ℕ → ℤ
proof
  have moebiusFunction := moebius.function
  have firstFunction := firstProjection.function at ℕ, ℕ
  have secondFunction := secondProjection.function at ℕ, ℕ
  have firstMember := FunctionBetween.apply_closed with firstFunction, member
  have secondMember := FunctionBetween.apply_closed with secondFunction, member
  have firstNat := nat.omega.backward with firstMember
  have quotientNatural := quotientNat.natural at firstProjection(ℕ, ℕ)(q), n
  have terms := dirichletTerms.function at moebius, g, quotientNat(n,
    firstProjection(ℕ, ℕ)(q)) with moebiusFunction, weightsG
  exact FunctionBetween.apply_closed with terms, secondMember

theorem nestedTerms.pair (g : Set) (n : Nat) (d c : Set)
  requires weightsG : g : ℕ → ℤ
  requires dMember : d ∈ ℕ
  requires cMember : c ∈ ℕ
  shows nestedTerms(g, n)((d, c)) =
    dirichletTerms(moebius, g, quotientNat(n, d))(c)
proof
  have pairMember := product.pair with dMember, cMember
  have value := nestedTerms.value at g, n, (d, c) with weightsG, pairMember
  have firstValue := firstProjection.pair at ℕ, ℕ, d, c with dMember, cMember
  have secondValue := secondProjection.pair at ℕ, ℕ, d, c with dMember, cMember
  rw [value, firstValue, secondValue]
  rfl

-- The converse: f(n) = ∑_{d ∣ n} μ(d) · g(n / d) for all n ≥ 1 gives
-- g(n) = ∑_{d ∣ n} f(d) for all n ≥ 1.
theorem moebius.inversion_converse (f g : Set)
  requires weightsF : f : ℕ → ℤ
  requires weightsG : g : ℕ → ℤ
  requires inverted : ∀ n ∈ ℕ, n ≠ 0 →
    f(n) = divisorSum(n, dirichletTerms(moebius, g, n))
  shows ∀ n ∈ ℕ, n ≠ 0 → g(n) = divisorSum(n, f)
proof
  intros n nMember nonzero
  have nNat := nat.omega.backward with nMember
  have moebiusFunction := moebius.function
  have naturals := divisors.subset_naturals at n
  have summatoryFunction := summatory.function at g with weightsG
  have sumsG : ∀ m ∈ ℕ, m ≠ 0 → summatory(g)(m) = divisorSum(m, g)
    intros m hm mNonzero
    have mNat := nat.omega.backward with hm
    exact summatory.value at g, m with weightsG
  have forward := moebius.inversion with weightsG, summatoryFunction, sumsG
  have gValue := forward with n, nMember, nonzero
  have nested := nestedTerms.function at g, n with weightsG
  have reflected := complementaryWeights.function at n, f with weightsF
  have rows := rowDivisorSums.function at nestedTerms(g, n), n with nested
  have columns := columnDivisorSums.function at nestedTerms(g, n), n with
    nested
  have convolved := dirichletTerms.function at moebius, summatory(g), n with
    moebiusFunction, summatoryFunction
  have rowsAgree : ∀ d ∈ divisors(n), complementaryWeights(n, f)(d) =
      rowDivisorSums(nestedTerms(g, n), n)(d)
    intros d hd
    have dNat := divisors.natural with hd
    have dMember := naturals with d, hd
    have factor := divisors.divides with hd
    have quotientNonzero := divides.quotient_ne_zero with nonzero, factor
    have quotientNatural := quotientNat.natural at d, n
    have quotientMember := nat.omega.forward with quotientNatural
    have innerNaturals := divisors.subset_naturals at quotientNat(n, d)
    have reflectedValue := complementaryWeights.value at n, f, d with weightsF
    have fValue := inverted with quotientNat(n, d), quotientMember,
      quotientNonzero
    have rowValue := rowDivisorSums.value at nestedTerms(g, n), n, d with
      nested
    have rowTermsFunction := rowTerms.function at nestedTerms(g, n), d with
      nested, dMember
    have innerTerms := dirichletTerms.function at moebius, g, quotientNat(n, d)
      with moebiusFunction, weightsG
    have agree : ∀ c ∈ divisors(quotientNat(n, d)),
        rowTerms(nestedTerms(g, n), d)(c) =
          dirichletTerms(moebius, g, quotientNat(n, d))(c)
      intros c hc
      have cMember := innerNaturals with c, hc
      have termValue := rowTerms.value at nestedTerms(g, n), d, c with nested,
        dMember
      have pairValue := nestedTerms.pair at g, n, d, c with weightsG, dMember,
        cMember
      cc
    have same := divisorSum.congr at quotientNat(n, d),
      rowTerms(nestedTerms(g, n), d), dirichletTerms(moebius, g,
        quotientNat(n, d)) with rowTermsFunction, innerTerms, agree
    rw [reflectedValue, fValue, rowValue, same]
    rfl
  have leftSum := divisorSum.congr with reflected, rows, rowsAgree
  have reflect := divisorSum.reflect with nonzero, weightsF
  have interchange := divisorSum.interchange at nestedTerms(g, n), n with
    nonzero, nested
  have columnsAgree : ∀ c ∈ divisors(n),
      columnDivisorSums(nestedTerms(g, n), n)(c) =
        dirichletTerms(moebius, summatory(g), n)(c)
    intros c hc
    have cNat := divisors.natural with hc
    have cMember := naturals with c, hc
    have cFactor := divisors.divides with hc
    have cQuotientNonzero := divides.quotient_ne_zero with nonzero, cFactor
    have cQuotientNatural := quotientNat.natural at c, n
    have innerC := divisors.quotient_subset with nonzero, cFactor
    have columnValue := columnDivisorSums.value at nestedTerms(g, n), n, c
      with nested
    have columnTermsFunction := columnTerms.function at nestedTerms(g, n), c
      with nested, cMember
    have reflectedG := complementaryWeights.function at quotientNat(n, c), g
      with weightsG
    have cInteger := moebius.integer at c
    have scaledValues : ∀ d ∈ divisors(quotientNat(n, c)),
        columnTerms(nestedTerms(g, n), c)(d) = integerMulValue(moebius(c),
          complementaryWeights(quotientNat(n, c), g)(d))
      intros d hd
      have dInN := innerC with d, hd
      have dNat := divisors.natural with dInN
      have dMember := naturals with d, dInN
      have dFactor := divisors.divides with dInN
      have dQuotientNatural := quotientNat.natural at d, n
      have termValue := columnTerms.value at nestedTerms(g, n), c, d with
        nested, cMember
      have pairValue := nestedTerms.pair at g, n, d, c with weightsG, dMember,
        cMember
      have innerValue := dirichletTerms.value at moebius, g, quotientNat(n, d),
        c with moebiusFunction, weightsG
      have reflectedValue := complementaryWeights.value at quotientNat(n, c),
        g, d with weightsG
      have dDividesQuotient := divisors.divides with hd
      have characterizationC := divides.mul_iff_quotient at c, d, n with
        nonzero, cFactor
      have product := characterizationC.backward with dDividesQuotient
      have commuted := mul.commutative at c, d
      rw ← commuted at product
      have characterizationD := divides.mul_iff_quotient at d, c, n with
        nonzero, dFactor
      have cDividesQuotient := characterizationD.forward with product
      have swapped := quotientNat.swap with nonzero, dFactor, cDividesQuotient
      rw [termValue, pairValue, innerValue, reflectedValue, swapped]
      rfl
    have scaled := divisorSum.scale at quotientNat(n, c),
      complementaryWeights(quotientNat(n, c), g),
        columnTerms(nestedTerms(g, n), c), moebius(c) with reflectedG,
          columnTermsFunction, scaledValues
    have reflectG := divisorSum.reflect at quotientNat(n, c), g with
      cQuotientNonzero, weightsG
    have summatoryValue := summatory.value at g, quotientNat(n, c) with
      weightsG
    have convolvedValue := dirichletTerms.value at moebius, summatory(g), n, c
      with moebiusFunction, summatoryFunction
    rw [columnValue, scaled, reflectG, convolvedValue, summatoryValue]
    rfl
  have rightSum := divisorSum.congr with columns, convolved, columnsAgree
  rw [← reflect, leftSum, interchange, rightSum]
  exact gValue