-- Generated by the TypeDensity exporter from TypeDensity/Basic.lean (namespace TypeDensity); do not edit.
import Mathlib.FieldTheory.RatFunc.Basic
import Mathlib.MeasureTheory.Measure.Haar.Unique
import Mathlib.NumberTheory.ArithmeticFunction.Moebius
import Mathlib.NumberTheory.Padics.PadicIntegers
import Mathlib.RingTheory.IntegralClosure.IntegralRestrict
import Mathlib.RingTheory.MvPowerSeries.Basic
import Mathlib.RingTheory.RamificationInertia.Basic
import Mathlib.Topology.Algebra.Nonarchimedean.AdicTopology
set_option autoImplicit false

noncomputable section

namespace TypeDensity

open scoped BigOperators ENNReal

open IsLocalRing MeasureTheory

/-- A factorization type: a multiset of pairs `(e, f)` (ramification index, residue degree), one per
irreducible factor. -/
abbrev FactorType := Multiset (ℕ × ℕ)

/-- The degree `∑ e·f` of a factorization type. -/
def typeDegree (σ : FactorType) : ℕ := (σ.map fun z => z.1 * z.2).sum

/-- A type is valid if all its pairs have `e, f ≥ 1`. -/
def ValidType (σ : FactorType) : Prop := ∀ z ∈ σ, 0 < z.1 ∧ 0 < z.2

/-- Scaling the residue degrees by `δ` (the effect of descending along a degree-`δ` unramified
extension). -/
def scaleType (δ : ℕ) (σ : FactorType) : FactorType := σ.map fun z => (z.1, δ * z.2)

/-- Syntactic rational expressions in one variable (the residue cardinality `q`). -/
inductive Expr where
  | const (r : ℚ)
  | var
  | add (x y : Expr)
  | mul (x y : Expr)
  | inv (x : Expr)
  deriving Repr, DecidableEq

namespace Expr

/-- Negation `-x`. -/
def neg (x : Expr) : Expr := .mul (.const (-1)) x

/-- Subtraction `x - y`. -/
def sub (x y : Expr) : Expr := .add x y.neg

/-- The power `x ^ n`. -/
def pow (x : Expr) : ℕ → Expr
  | 0 => .const 1
  | n+1 => .mul (pow x n) x

/-- Evaluation at `q` (with `0⁻¹ = 0`). -/
def eval (q : ℚ) : Expr → ℚ
  | .const r => r
  | .var => q
  | .add x y => eval q x + eval q y
  | .mul x y => eval q x * eval q y
  | .inv x => (eval q x)⁻¹

/-- No denominator vanishes when evaluating at `q`. -/
def safe (q : ℚ) : Expr → Prop
  | .const _ => True
  | .var => True
  | .add x y => safe q x ∧ safe q y
  | .mul x y => safe q x ∧ safe q y
  | .inv x => safe q x ∧ eval q x ≠ 0

/-- Substitution of `y` for the variable. -/
def subst (x : Expr) (y : Expr) : Expr :=
  match x with
  | .const r => .const r
  | .var => y
  | .add a b => .add (subst a y) (subst b y)
  | .mul a b => .mul (subst a y) (subst b y)
  | .inv a => .inv (subst a y)

/-- A (numerator, denominator) pair of polynomials representing the expression. -/
def fraction : Expr → Polynomial ℚ × Polynomial ℚ
  | .const r => (Polynomial.C r, 1)
  | .var => (Polynomial.X, 1)
  | .add x y => let a := fraction x; let b := fraction y
               (a.1*b.2+b.1*a.2, a.2*b.2)
  | .mul x y => let a := fraction x; let b := fraction y
               (a.1*b.1, a.2*b.2)
  | .inv x => let a := fraction x; (a.2,a.1)

/-- The rational function `numerator / denominator` of the expression. -/
def rational (x : Expr) : RatFunc ℚ :=
  algebraMap (Polynomial ℚ) (RatFunc ℚ) x.fraction.1 /
    algebraMap (Polynomial ℚ) (RatFunc ℚ) x.fraction.2

/-- The sum of a list of expressions. -/
def sum (xs : List Expr) : Expr := xs.foldr .add (.const 0)

/-- The product of a list of expressions. -/
def prod (xs : List Expr) : Expr := xs.foldr .mul (.const 1)

/-- The expression is safe at every integer `q ≥ 2`. -/
def allSafe (x : Expr) : Prop := ∀ q : ℕ, 2 ≤ q → x.safe q

end Expr

/-- The rational function `f` has a representation `P / Q` with `Q(q) ≠ 0`. -/
def RegularAt (f : RatFunc ℚ) (q : ℚ) : Prop :=
  ∃ P Q : Polynomial ℚ, Q.eval q ≠ 0 ∧
    f = algebraMap (Polynomial ℚ) (RatFunc ℚ) P /
      algebraMap (Polynomial ℚ) (RatFunc ℚ) Q

/-- A law: a finite formal sum of factorization types with expression coefficients. -/
abbrev Law := List (FactorType × Expr)

/-- The unit law (the empty type with coefficient `1`). -/
def lawOne : Law := [(0, .const 1)]

/-- The law of a single factor of type `(e, f)`. -/
def lawAtom (e f : ℕ) : Law := [({(e,f)}, .const 1)]

/-- Multiplying every coefficient by `x`. -/
def lawScale (x : Expr) (L : Law) : Law := L.map fun z => (z.1,.mul x z.2)

/-- The product of laws (types add, coefficients multiply). -/
def lawMul (L M : Law) : Law := L.flatMap fun z => M.map fun w => (z.1+w.1,.mul z.2 w.2)

/-- The product of a list of laws. -/
def lawProd (Ls : List Law) : Law := Ls.foldr lawMul lawOne

/-- The total coefficient of the type `σ` in a law. -/
def lawCoeff (L : Law) (σ : FactorType) : Expr :=
  Expr.sum (L.map fun z => if z.1=σ then z.2 else .const 0)

/-- Descent along a degree-`δ` unramified extension: scale types by `δ` and substitute `q ↦ q^δ`. -/
def lawDescend (δ : ℕ) (L : Law) : Law :=
  L.map fun z => (scaleType δ z.1, z.2.subst (.pow .var δ))

/-- All coefficients of the law are safe at every `q ≥ 2`. -/
def LawSafe (L : Law) : Prop := ∀ z ∈ L, z.2.allSafe

/-- Recursion states: `single a c` is `S(a;c)` (one boundary component of multiplicity `a`, contact
`c`), `crossing a b c h` is `T(a,b;c,h)` (two crossing boundary components). -/
inductive State where
  | single (a c : ℕ)
  | crossing (a b c h : ℕ)
  deriving DecidableEq, Repr

namespace State

/-- All parameters are positive. -/
def Valid : State → Prop
  | .single a c => 0 < a ∧ 0 < c
  | .crossing a b c h => 0 < a ∧ 0 < b ∧ 0 < c ∧ 0 < h

/-- The degree `a·c` (resp. `a·c + b·h`). -/
def degree : State → ℕ
  | .single a c => a*c
  | .crossing a b c h => a*c+b*h

/-- The weight `a` (resp. `a + b`). -/
def weight : State → ℕ
  | .single a _ => a
  | .crossing a b _ _ => a+b

/-- `0` for single-boundary states, `1` for crossing states. -/
def tag : State → ℕ
  | .single _ _ => 0
  | .crossing _ _ _ _ => 1

/-- The well-founded rank `2d² + 2(d - w) + tag` driving the recursion. -/
def rank (s : State) : ℕ := 2*s.degree*s.degree + 2*(s.degree-s.weight) + s.tag

/-- Terminal states `S(a;1)`. -/
def terminal : State → Bool
  | .single _ 1 => true
  | _ => false

end State

/-- A residue-divisor pattern: multiplicities `m (i, j)` of closed points of degree `i + 1` with
contact `j + 1`. -/
abbrev Pattern (n : ℕ) := (Fin n × Fin n) → Fin (n+1)

/-- The weight `∑ (i + 1)(j + 1) m(i, j)` of a pattern. -/
def patternWeight {n : ℕ} (m : Pattern n) : ℕ :=
  ∑ ij : Fin n × Fin n, (ij.1.val+1)*(ij.2.val+1)*(m ij).val

/-- A tangent pattern of a state of degree `n`: the tangent degree `d`, the boundary contacts
`rt`, `rx`, and the pattern of the other points. -/
structure TangentPattern (n : ℕ) where
  d : Fin (n+1)
  rt : Fin (n+1)
  rx : Fin (n+1)
  other : Pattern n
  deriving DecidableEq, Fintype

/-- The admissible tangent patterns of a state (Proposition 5.1). -/
def Allowed (s : State) (P : TangentPattern s.degree) : Prop :=
  0<P.d.val ∧ P.rt.val+P.rx.val+patternWeight P.other=P.d.val ∧
  match s with
  | .single _ c => P.d.val≤c ∧ P.rx.val=0 ∧ (0<P.rt.val ↔ P.d.val<c)
  | .crossing _ _ c h => P.d.val≤c ∧ P.d.val≤h ∧
      (0<P.rt.val ↔ P.d.val<c) ∧ (0<P.rx.val ↔ P.d.val<h)

/-- The children `(δ, S(a; j + 1))` contributed by the non-boundary points of a pattern. -/
def otherChildren {n : ℕ} (a : ℕ) (m : Pattern n) : List (ℕ × State) :=
  (Finset.univ : Finset (Fin n × Fin n)).toList.flatMap fun ij =>
    List.replicate (m ij).val (ij.1.val+1, .single a (ij.2.val+1))

/-- All children (residue degree, child state) of a state under a tangent pattern. -/
def children (s : State) (P : TangentPattern s.degree) : List (ℕ × State) :=
  match s with
  | .single a c =>
    (if P.rt.val=0 then [] else [(1,.crossing a a P.rt.val (c-P.d.val))]) ++
      otherChildren a P.other
  | .crossing a b c h =>
    (if P.rt.val=0 then [] else [(1,.crossing (a+b) a P.rt.val (c-P.d.val))]) ++
    (if P.rx.val=0 then [] else [(1,.crossing (a+b) b P.rx.val (h-P.d.val))]) ++
      otherChildren (a+b) P.other

/-- Self patterns return to the same single-boundary state. -/
def SelfPattern (s : State) (P : TangentPattern s.degree) : Prop :=
  ∃ a c, s = .single a c ∧ children s P = [(1,s)]

open scoped Classical in
/-- The finite set of allowed tangent patterns. -/
noncomputable def allowedPatterns (s : State) : Finset (TangentPattern s.degree) :=
  (Finset.univ).filter (Allowed s)

/-- `d(d+1)/2`. -/
def tri (d : ℕ) : ℕ := d*(d+1)/2

/-- `d(d-1)/2`. -/
def chooseTwo (d : ℕ) : ℕ := d*(d-1)/2

/-- The expression `q - 1`. -/
def eminusOne : Expr := .sub .var (.const 1)

/-- The volume `q^{-c}` (resp. `(q - 1) q^{-(c+h)}`) of a state. -/
def stateVolume : State → Expr
  | .single _ c => .inv (.pow .var c)
  | .crossing _ _ c h => .mul eminusOne (.inv (.pow .var (c+h)))

/-- The self-return weight `q · q^{-tri c}` of `S(a;c)`. -/
def returnWeight (c : ℕ) : Expr := .mul .var (.inv (.pow .var (tri c)))

/-- The probability of the tangent divisor degree `d` (equation (7.2)). -/
def divisorProbability (s : State) (d : ℕ) : Expr :=
  match s with
  | .single _ c => .mul (.inv (.pow .var (tri d))) (.pow eminusOne (if d<c then 1 else 0))
  | .crossing _ _ c h =>
    let j := (if d<c then 1 else 0) + (if d<h then 1 else 0)
    .mul (.inv (.pow .var (chooseTwo d)))
      (if j=0 then .inv eminusOne else .pow eminusOne (j-1))

/-- The number `(1/δ) ∑_{j ∣ δ} μ(j) q^{δ/j}` of closed points of degree `δ` of `ℙ¹` minus `∞`
(necklace polynomial). -/
def pointCount (δ : ℕ) : Expr :=
  .mul (.const (1/(δ:ℚ))) (Expr.sum ((Nat.divisors δ).toList.map fun j =>
    .mul (.const ((ArithmeticFunction.moebius j : ℤ) : ℚ)) (.pow .var (δ/j))))

/-- Available points of degree `δ` (degree one: `q + offset`). -/
def availablePoints (offset : ℤ) (δ : ℕ) : Expr :=
  if δ=1 then .add .var (.const offset) else pointCount δ

/-- The falling factorial `x (x - 1) ⋯ (x - m + 1)`. -/
def falling (x : Expr) (m : ℕ) : Expr :=
  Expr.prod ((List.range m).map fun j => .sub x (.const j))

/-- The number of residue divisors with pattern `m` (choices of distinct points per degree). -/
def divisorCount {n : ℕ} (offset : ℤ) (m : Pattern n) : Expr :=
  Expr.prod ((Finset.univ : Finset (Fin n)).toList.map fun i =>
    .mul (falling (availablePoints offset (i.val+1)) (∑ j : Fin n, (m (i,j)).val))
      (.const (1 / (∏ j : Fin n, ((Nat.factorial (m (i,j)).val : ℕ) : ℚ)))))

/-- The number of tangent divisors with the pattern of `P`. -/
def tangentCount (s : State) (P : TangentPattern s.degree) : Expr :=
  divisorCount (match s with | .single _ _ => 0 | .crossing _ _ _ _ => -1) P.other

/-- `1 + q + ⋯ + q^n`. -/
def projectiveNormalizer (n : ℕ) : Expr :=
  Expr.sum ((List.range (n+1)).map fun i => .pow .var i)

open scoped Classical in
/-- The measured recursion of §8 with explicit fuel. -/
noncomputable def stateLawFuel : ℕ → State → Law
  | 0, _ => []
  | k+1, s =>
    match s with
    | .single a 1 => lawAtom a 1
    | _ =>
      let terms := (allowedPatterns s).toList.flatMap fun P =>
        if SelfPattern s P then [] else
          lawScale (.mul (tangentCount s P) (divisorProbability s P.d.val))
            (lawProd ((children s P).map fun z => lawDescend z.1 (stateLawFuel k z.2)))
      match s with
      | .single _ c => lawScale (.inv (.sub (.const 1) (returnWeight c))) terms
      | .crossing _ _ _ _ => terms

/-- The law of a state (fuel `rank + 1` suffices). -/
noncomputable def stateLaw (s : State) : Law := stateLawFuel (s.rank+1) s

/-- The law of the root cell `S(1;r)` (`lawOne` for `r = 0`). -/
noncomputable def rootLaw (r : ℕ) : Law := if r=0 then lawOne else stateLaw (.single 1 r)

/-- Residue patterns of total weight `n`. -/
noncomputable def residuePatterns (n : ℕ) : Finset (Pattern n) :=
  Finset.univ.filter fun m => patternWeight m=n

/-- The global law: average over residue divisors of the products of descended child laws. -/
noncomputable def globalLaw (projective : Bool) (n : ℕ) : Law :=
  lawScale (.inv (if projective then projectiveNormalizer n else .pow .var n))
    ((residuePatterns n).toList.flatMap fun m =>
      lawScale (divisorCount (if projective then 1 else 0) m)
        (lawProd ((otherChildren 1 m).map fun z =>
          lawDescend z.1 (stateLaw z.2))))

/-- The universal monic density expression `A_n,σ(q)`. -/
noncomputable def monicExpr (n : ℕ) (σ : FactorType) : Expr := lawCoeff (globalLaw false n) σ

/-- The universal projective density expression `B_n,σ(q)`. -/
noncomputable def projectiveExpr (n : ℕ) (σ : FactorType) : Expr := lawCoeff (globalLaw true n) σ

section Classical

variable (O : Type) [CommRing O] [IsDomain O]

/-- The fraction field of `O`. -/
abbrev BaseField := FractionRing O

/-- The invariant `(e, f)` of an irreducible `g`: ramification index and inertia degree of the
maximal ideal of the integral closure of `O` in `Frac(O)[X]/(g)` when that closure is local (junk
value `(0, 0)` otherwise). -/
noncomputable def factorInvariant (g : Polynomial (BaseField O)) : ℕ × ℕ := by
  classical
  let B := integralClosure O (AdjoinRoot g)
  exact if h : IsLocalRing B then
    letI : IsLocalRing B := h
    ((maximalIdeal B).ramificationIdx O, (maximalIdeal B).inertiaDeg O)
  else (0,0)

/-- `F` is the product of distinct monic irreducibles whose invariants form the type `σ`. -/
def HasPolynomialType (F : Polynomial (BaseField O)) (σ : FactorType) : Prop :=
  ∃ gs : List (Polynomial (BaseField O)),
    (∀ g ∈ gs, g.Monic ∧ Irreducible g) ∧ gs.Pairwise (fun g h => g ≠ h) ∧
    F=gs.prod ∧ σ=(gs.map (factorInvariant O) : Multiset (ℕ × ℕ))

/-- The monic polynomial `X^n + ∑ aᵢ Xⁱ`. -/
def monicPolynomial {n : ℕ} (a : Fin n → O) : Polynomial (BaseField O) :=
  Polynomial.X^n + ∑ i : Fin n,
    Polynomial.C (algebraMap O (BaseField O) (a i)) * Polynomial.X^i.val

/-- The polynomial `∑ aᵢ Xⁱ` (the dehomogenized binary form). -/
def affinePolynomial {n : ℕ} (a : Fin (n + 1) → O) : Polynomial (BaseField O) :=
  ∑ i : Fin (n+1), Polynomial.C (algebraMap O (BaseField O) (a i)) * Polynomial.X^i.val

/-- Some coefficient is a unit. -/
def primitive {n : ℕ} (a : Fin (n + 1) → O) : Prop := ∃ i, IsUnit (a i)

/-- The binary form with coefficients `a` has type `σ` (a missing top degree contributes the point
at infinity, of type `(1, 1)`). -/
def HasBinaryType {n : ℕ} (a : Fin (n + 1) → O) (σ : FactorType) : Prop :=
  let F := affinePolynomial O a
  F≠0 ∧ F.natDegree≤n ∧ n-F.natDegree≤1 ∧
    ∃ τ, HasPolynomialType O (Polynomial.C F.leadingCoeff⁻¹ * F) τ ∧
      σ=τ+Multiset.replicate (n-F.natDegree) (1,1)

end Classical

section Measure

variable (O : Type) [CommRing O] [IsDomain O]
  [TopologicalSpace O] [IsTopologicalRing O] [CompactSpace O] [T2Space O]
  [SecondCountableTopology O]
  [MeasurableSpace O] [BorelSpace O]

/-- The normalized Haar (probability) measure on `Oⁿ`. -/
noncomputable def coefficientMeasure (n : ℕ) : Measure (Fin n → O) :=
  (Measure.addHaar (G := Fin n → O) Set.univ)⁻¹ • Measure.addHaar (G := Fin n → O)

/-- The probability that a random monic polynomial of degree `n` over `O` has type `σ`. -/
noncomputable def monicDensity (n : ℕ) (σ : FactorType) : ℝ≥0∞ :=
  coefficientMeasure O n {a | HasPolynomialType O (monicPolynomial O a) σ}

/-- The probability that a random primitive binary form of degree `n` over `O` has type `σ`. -/
noncomputable def projectiveDensity (n : ℕ) (σ : FactorType) : ℝ≥0∞ :=
  coefficientMeasure O (n+1) {a | primitive O a ∧ HasBinaryType O a σ} /
    coefficientMeasure O (n+1) {a | primitive O a}

end Measure

/-- The value `x(q)` of an expression, as an extended nonnegative real. -/
abbrev densityValue (x : Expr) (q : ℕ) : ℝ≥0∞ := ENNReal.ofReal ((x.eval q : ℚ) : ℝ)

/-- `q` is a prime power `p^r`, `r ≥ 1`. -/
def PrimePower (q : ℕ) : Prop := ∃ p r : ℕ, p.Prime ∧ 0<r ∧ q=p^r

/-- Both universal expressions are regular at every prime power. -/
def FormulaRegular (n : ℕ) (σ : FactorType) : Prop :=
  ∀ q : ℕ, PrimePower q →
    RegularAt (monicExpr n σ).rational q ∧ RegularAt (projectiveExpr n σ).rational q

section PadicIntegers

variable (p : ℕ) [Fact p.Prime] (K : Type) [Field K] [Algebra ℤ_[p] K]

/-- The ring of integers of `K`: the integral closure of `ℤ_[p]` in `K`, as a type. -/
abbrev padicIntegers : Type := ↥(integralClosure ℤ_[p] K)

/-- `K/ℚ_p` is unramified (for the ring of integers `padicIntegers p K`) iff `p` stays
irreducible. -/
def PadicUnramified : Prop := Irreducible (p : padicIntegers p K)

end PadicIntegers

section RouteF

variable (O : Type) [CommRing O] [IsDomain O]

/-- `HasPolynomialType` together with finiteness of the integral closure in each root field. -/
def HasPolynomialTypeF (F : Polynomial (BaseField O)) (σ : FactorType) : Prop :=
  ∃ gs : List (Polynomial (BaseField O)),
    (∀ g ∈ gs, g.Monic ∧ Irreducible g) ∧ gs.Pairwise (fun g h => g ≠ h) ∧
    F=gs.prod ∧ σ=(gs.map (factorInvariant O) : Multiset (ℕ × ℕ)) ∧
    ∀ g ∈ gs, Module.Finite O (integralClosure O (AdjoinRoot g))

/-- `HasBinaryType` with `HasPolynomialTypeF` in place of `HasPolynomialType`. -/
def HasBinaryTypeF {n : ℕ} (a : Fin (n + 1) → O) (σ : FactorType) : Prop :=
  let F := affinePolynomial O a
  F≠0 ∧ F.natDegree≤n ∧ n-F.natDegree≤1 ∧
    ∃ τ, HasPolynomialTypeF O (Polynomial.C F.leadingCoeff⁻¹ * F) τ ∧
      σ=τ+Multiset.replicate (n-F.natDegree) (1,1)

variable [TopologicalSpace O] [IsTopologicalRing O] [CompactSpace O] [T2Space O]
  [SecondCountableTopology O] [MeasurableSpace O] [BorelSpace O]

/-- The monic density for the route-F predicate. -/
noncomputable def monicDensityF (n : ℕ) (σ : FactorType) : ℝ≥0∞ :=
  coefficientMeasure O n {a | HasPolynomialTypeF O (monicPolynomial O a) σ}

/-- The projective density for the route-F predicate. -/
noncomputable def projectiveDensityF (n : ℕ) (σ : FactorType) : ℝ≥0∞ :=
  coefficientMeasure O (n+1) {a | primitive O a ∧ HasBinaryTypeF O a σ} /
    coefficientMeasure O (n+1) {a | primitive O a}

end RouteF

end TypeDensity

end