import Mathlib.Analysis.SpecialFunctions.Log.Basic
import Mathlib.Data.Fintype.BigOperators
import Mathlib.Data.Real.Basic

/-
The Problem. Protocol and classical theorems need a shared information/analysis
interface, while concrete `Real.log`/spectral realizations should stay isolated.

Solution Structure.
1. Define compact operation interfaces (`Model`, `AnalysisModel`).
2. Define compact law interfaces (`Laws`, `AnalysisLaws`).
3. Re-export operations and laws through stable wrappers.
4. Keep concrete realizations in `ClassicalAnalysisInstance`.

## Trust Boundary

This file and `ClassicalAnalysisInstance.lean` form the explicit trust boundary
for all noncomputable real-analysis assumptions in the project. A reviewer can:
1. Verify the laws stated here match standard information theory
2. Check `ClassicalAnalysisInstance` proves each law from Mathlib
3. Run `lake build` to confirm type-correctness
4. Run `just escape` to confirm no escapes outside this boundary
-/

/-! # ClassicalAnalysisAPI

Abstract interface for information-theoretic and real-analysis operations.

## Primary References

- **Shannon (1948)**: "A Mathematical Theory of Communication." Bell System
  Technical Journal. Introduces Shannon entropy H(X) = -∑ p(x) log p(x).

- **Kullback & Leibler (1951)**: "On Information and Sufficiency." Annals of
  Mathematical Statistics. Introduces KL divergence D(P‖Q) = ∑ p(x) log(p(x)/q(x)).

- **Cover & Thomas (2006)**: "Elements of Information Theory" (2nd ed.).
  Standard textbook; theorem numbers referenced below.

## Expose

This module defines the abstract boundary consumed by the rest of the project:

- entropy interface:
  `EntropyAPI.Distribution`, `EntropyAPI.IsErasureKernel`,
  `EntropyAPI.Model`, `EntropyAPI.Laws`,
  `EntropyAPI.shannonEntropy`, `EntropyAPI.klDivergence`,
  `EntropyAPI.mutualInfo`
- entropy law wrappers:
  `EntropyAPI.shannon_entropy_nonneg`,
  `EntropyAPI.shannon_entropy_le_log_card`,
  `EntropyAPI.kl_divergence_nonneg`,
  `EntropyAPI.kl_divergence_eq_zero_iff`,
  `EntropyAPI.mutual_info_zero_of_erasure`
- extended analysis interface:
  `EntropyAPI.AnalysisModel`, `EntropyAPI.AnalysisLaws`,
  `EntropyAPI.exponentialTail`, `EntropyAPI.fluctuationScale`,
  `EntropyAPI.finiteAverage`, `EntropyAPI.normalizedCumulative`,
  `EntropyAPI.transitionMatrixComplex`,
  `EntropyAPI.nontrivialSpectrumModuli`,
  `EntropyAPI.secondSpectrumValue`,
  `EntropyAPI.spectralGapValue`

Concrete implementations live in `ClassicalAnalysisInstance`.
-/

set_option autoImplicit false

namespace EntropyAPI

open scoped BigOperators

/-! ## Core Types -/

/-- Finite probability distributions represented by a real-valued pmf. -/
structure Distribution (α : Type*) [Fintype α] where
  pmf : α → ℝ
  nonneg : ∀ a, 0 ≤ pmf a
  sum_one : ∑ a, pmf a = 1

/--
`IsErasureKernel labelDist joint` means observation carries no label
information: all mass is concentrated at one observation value.
-/
def IsErasureKernel {L O : Type} [DecidableEq O]
    (labelDist : L → ℝ) (joint : L × O → ℝ) : Prop :=
  ∃ o0, ∀ lo : L × O, joint lo = if lo.2 = o0 then labelDist lo.1 else 0

/-! ## Entropy Operations -/

/-- Information-theoretic operations required by downstream modules. -/
class Model where
  shannonEntropy : {α : Type*} → [Fintype α] → (α → ℝ) → ℝ
  klDivergence : {α : Type*} → [Fintype α] → (α → ℝ) → (α → ℝ) → ℝ
  mutualInfo :
    {α : Type*} → {β : Type*} → [Fintype α] → [Fintype β] →
    (α × β → ℝ) → ℝ

section Operations

variable [Model]

/-- Wrapper for Shannon entropy. -/
def shannonEntropy {α : Type*} [Fintype α] (p : α → ℝ) : ℝ :=
  Model.shannonEntropy p

/-- Wrapper for KL divergence. -/
def klDivergence {α : Type*} [Fintype α] (p q : α → ℝ) : ℝ :=
  Model.klDivergence p q

/-- Wrapper for mutual information. -/
def mutualInfo {α : Type*} {β : Type*} [Fintype α] [Fintype β]
    (pXY : α × β → ℝ) : ℝ :=
  Model.mutualInfo pXY

end Operations

/-! ## Entropy Laws

The following laws are the core information-theoretic properties required by
downstream modules. Each corresponds to a standard theorem in information theory.
-/

/-- Shannon entropy nonnegativity law.

**Statement**: H(X) ≥ 0 for any discrete random variable X.

**Intuition**: Each term p(x) log p(x) is nonpositive on [0,1] (since log p ≤ 0
when p ≤ 1), so negating the sum yields a nonnegative quantity.

**Reference**: Shannon (1948) §6; Cover & Thomas Theorem 2.1.1. -/
abbrev ShannonEntropyNonneg (M : Model) : Prop :=
  ∀ {α : Type*} [Fintype α], (d : Distribution α) →
    0 ≤ @Model.shannonEntropy M α _ d.pmf

/-- Shannon entropy cardinality upper bound law.

**Statement**: H(X) ≤ log|X| with equality iff X is uniformly distributed.

**Intuition**: Entropy measures uncertainty. Maximum uncertainty occurs when all
outcomes are equally likely. The uniform distribution achieves log|X| bits.
Formally, this follows from D(P‖U) ≥ 0 where U is uniform.

**Reference**: Cover & Thomas Theorem 2.6.4. -/
abbrev ShannonEntropyLeLogCard (M : Model) : Prop :=
  ∀ {α : Type*} [Fintype α] [Nonempty α], (d : Distribution α) →
    @Model.shannonEntropy M α _ d.pmf ≤ Real.log (Fintype.card α)

/-- KL divergence nonnegativity law (Gibbs' inequality).

**Statement**: D(P‖Q) ≥ 0 for distributions P, Q with P absolutely continuous
w.r.t. Q.

**Intuition**: KL divergence measures the "extra bits" needed to encode samples
from P using a code optimized for Q. You can never do better than the optimal
code for P, so the overhead is nonnegative.

**Proof sketch**: Apply Jensen's inequality to the convex function -log:
  D(P‖Q) = E_P[-log(Q/P)] ≥ -log(E_P[Q/P]) = -log(1) = 0

**Reference**: Kullback & Leibler (1951); Cover & Thomas Theorem 2.6.3. -/
abbrev KlDivergenceNonneg (M : Model) : Prop :=
  ∀ {α : Type*} [Fintype α], (p q : Distribution α) →
    (∀ a, p.pmf a ≠ 0 → q.pmf a ≠ 0) →
    0 ≤ @Model.klDivergence M α _ p.pmf q.pmf

/-- KL divergence zero characterization law.

**Statement**: D(P‖Q) = 0 if and only if P = Q (almost everywhere).

**Intuition**: Zero divergence means P and Q are indistinguishable from an
information-theoretic perspective. Strict convexity of -log implies Jensen's
inequality is tight only when the argument is constant a.e., forcing P = Q.

**Reference**: Cover & Thomas Theorem 2.6.3 (equality condition). -/
abbrev KlDivergenceEqZeroIff (M : Model) : Prop :=
  ∀ {α : Type*} [Fintype α], (p q : Distribution α) →
    (∀ a, p.pmf a ≠ 0 → q.pmf a ≠ 0) →
    (@Model.klDivergence M α _ p.pmf q.pmf = 0 ↔ p.pmf = q.pmf)

/-- Erasure implies zero mutual information law.

**Statement**: If the observation channel is constant (all labels map to the
same output), then I(Label; Observation) = 0.

**Intuition**: Mutual information I(X;Y) = H(Y) - H(Y|X) measures how much
knowing X reduces uncertainty about Y. If Y is deterministically fixed
regardless of X, then H(Y) = H(Y|X) = 0, so I(X;Y) = 0.

**Application**: This law is central to the blindness property in choreographic
projection. When a non-participant role's observation function is constant
across all branches, they learn nothing about which branch was taken.

**Reference**: Cover & Thomas §2.8 (data processing inequality, extreme case). -/
abbrev MutualInfoZeroOfErasure (M : Model) : Prop :=
  ∀ {L O : Type} [Fintype L] [Fintype O] [DecidableEq O],
    (labelDist : L → ℝ) →
    (∀ l, 0 ≤ labelDist l) →
    (∑ l, labelDist l = 1) →
    (joint : L × O → ℝ) →
    IsErasureKernel labelDist joint →
    @Model.mutualInfo M L O _ _ joint = 0

/-- Law-bearing entropy interface used by downstream modules. -/
class Laws extends Model where
  shannon_entropy_nonneg : ShannonEntropyNonneg toModel
  shannon_entropy_le_log_card : ShannonEntropyLeLogCard toModel
  kl_divergence_nonneg : KlDivergenceNonneg toModel
  kl_divergence_eq_zero_iff : KlDivergenceEqZeroIff toModel
  mutual_info_zero_of_erasure : MutualInfoZeroOfErasure toModel

/-- Promote `Laws` to `Model` so operation wrappers infer automatically. -/
instance (priority := 100) lawsToModel [Laws] : Model :=
  Laws.toModel

/-! ## Entropy Laws: Exported Theorems -/

section Laws

variable [Laws]

/-- Re-export Shannon entropy nonnegativity from `Laws`. -/
theorem shannon_entropy_nonneg {α : Type*} [Fintype α]
    (d : Distribution α) :
    0 ≤ shannonEntropy d.pmf :=
  Laws.shannon_entropy_nonneg (self := inferInstance) d

/-- Re-export Shannon entropy cardinality upper bound from `Laws`. -/
theorem shannon_entropy_le_log_card {α : Type*} [Fintype α] [Nonempty α]
    (d : Distribution α) :
    shannonEntropy d.pmf ≤ Real.log (Fintype.card α) :=
  Laws.shannon_entropy_le_log_card (self := inferInstance) d

/-- Re-export KL nonnegativity from `Laws`. -/
theorem kl_divergence_nonneg {α : Type*} [Fintype α]
    (p q : Distribution α)
    (habs : ∀ a, p.pmf a ≠ 0 → q.pmf a ≠ 0) :
    0 ≤ klDivergence p.pmf q.pmf :=
  Laws.kl_divergence_nonneg (self := inferInstance) p q habs

/-- Re-export zero-KL characterization from `Laws`. -/
theorem kl_divergence_eq_zero_iff {α : Type*} [Fintype α]
    (p q : Distribution α)
    (habs : ∀ a, p.pmf a ≠ 0 → q.pmf a ≠ 0) :
    klDivergence p.pmf q.pmf = 0 ↔ p.pmf = q.pmf :=
  Laws.kl_divergence_eq_zero_iff (self := inferInstance) p q habs

/-- Re-export erasure-implies-zero-MI from `Laws`. -/
theorem mutual_info_zero_of_erasure {L O : Type}
    [Fintype L] [Fintype O] [DecidableEq O]
    (labelDist : L → ℝ)
    (h_nn : ∀ l, 0 ≤ labelDist l)
    (h_sum : ∑ l, labelDist l = 1)
    (joint : L × O → ℝ)
    (hErase : IsErasureKernel labelDist joint) :
    mutualInfo joint = 0 :=
  Laws.mutual_info_zero_of_erasure
    (self := inferInstance) labelDist h_nn h_sum joint hErase

end Laws

/-! ## Extended Analysis Operations

These operations support concentration inequalities and spectral analysis of
Markov chains, used for quantitative protocol guarantees (mixing times,
convergence rates).
-/

/-- Extended operations for real/complex-analysis touchpoints.

These support:
- **Concentration bounds**: `exponentialTail` for sub-Gaussian tails
- **CLT scaling**: `fluctuationScale` = √n for normalized sums
- **Ergodic averages**: `finiteAverage`, `normalizedCumulative`
- **Spectral analysis**: transition matrix spectrum for mixing time bounds -/
class AnalysisModel where
  /-- Exponential tail bound 2·exp(-t²/2σ²) for sub-Gaussian concentration. -/
  exponentialTail : ℝ → ℝ → ℝ
  /-- Fluctuation scale √n for CLT normalization. -/
  fluctuationScale : Nat → ℝ
  /-- Arithmetic mean of a finite family. -/
  finiteAverage : {n : Nat} → (Fin n → ℝ) → ℝ
  /-- Normalized cumulative sum (∑ᵢ₌₀ᵗ (xᵢ - μ)) / √N. -/
  normalizedCumulative : (Nat → ℝ) → ℝ → Nat → Nat → ℝ
  /-- Lift real transition kernel to complex matrix for spectral analysis. -/
  transitionMatrixComplex :
    {State : Type*} → [Fintype State] → [DecidableEq State] →
      (State → State → ℝ) → Matrix State State ℂ
  /-- Moduli of non-trivial eigenvalues (excluding the Perron eigenvalue 1). -/
  nontrivialSpectrumModuli :
    {State : Type*} → [Fintype State] → [DecidableEq State] →
      (State → State → ℝ) → Set ℝ
  /-- Second-largest eigenvalue modulus (determines mixing rate). -/
  secondSpectrumValue :
    {State : Type*} → [Fintype State] → [DecidableEq State] →
      (State → State → ℝ) → ℝ
  /-- Spectral gap 1 - λ₂ (larger gap = faster mixing). -/
  spectralGapValue :
    {State : Type*} → [Fintype State] → [DecidableEq State] →
      (State → State → ℝ) → ℝ

section AnalysisOperations

variable [AnalysisModel]

/-- Wrapper for the exponential-tail form. -/
def exponentialTail (σ t : ℝ) : ℝ :=
  AnalysisModel.exponentialTail σ t

/-- Wrapper for the fluctuation scale. -/
def fluctuationScale (n : Nat) : ℝ :=
  AnalysisModel.fluctuationScale n

/-- Wrapper for finite averaging. -/
def finiteAverage {n : Nat} (x : Fin n → ℝ) : ℝ :=
  AnalysisModel.finiteAverage x

/-- Wrapper for normalized cumulative process. -/
def normalizedCumulative (x : Nat → ℝ) (μ : ℝ) (N t : Nat) : ℝ :=
  AnalysisModel.normalizedCumulative x μ N t

/-- Wrapper for the complex transition-matrix realization. -/
def transitionMatrixComplex {State : Type*} [Fintype State] [DecidableEq State]
    (kernel : State → State → ℝ) : Matrix State State ℂ :=
  AnalysisModel.transitionMatrixComplex kernel

/-- Wrapper for nontrivial spectrum moduli. -/
def nontrivialSpectrumModuli {State : Type*} [Fintype State] [DecidableEq State]
    (kernel : State → State → ℝ) : Set ℝ :=
  AnalysisModel.nontrivialSpectrumModuli kernel

/-- Wrapper for the second spectral value. -/
def secondSpectrumValue {State : Type*} [Fintype State] [DecidableEq State]
    (kernel : State → State → ℝ) : ℝ :=
  AnalysisModel.secondSpectrumValue kernel

/-- Wrapper for the spectral-gap value. -/
def spectralGapValue {State : Type*} [Fintype State] [DecidableEq State]
    (kernel : State → State → ℝ) : ℝ :=
  AnalysisModel.spectralGapValue kernel

end AnalysisOperations

/-! ## Extended Analysis Laws

Laws for the extended analysis interface. These ensure the concrete
implementations behave as expected for concentration and spectral arguments.
-/

/-- Laws for `AnalysisModel` ensuring semantic correctness.

These laws capture:
- Basic properties of exponential/Gaussian tails
- √n scaling for CLT normalization
- Permutation invariance of averages (exchangeability)
- Spectral gap definition consistency -/
class AnalysisLaws extends AnalysisModel where
  /-- Exponential tails are nonnegative (probability bound). -/
  exponential_tail_nonneg :
    ∀ (σ t : ℝ), 0 ≤ toAnalysisModel.exponentialTail σ t
  /-- At zero threshold, tail probability is 2 (two-sided bound covers all). -/
  exponential_tail_zero :
    ∀ (σ : ℝ), toAnalysisModel.exponentialTail σ 0 = 2
  /-- Fluctuation scale is positive for positive sample size. -/
  fluctuation_scale_pos :
    ∀ {n : Nat}, 0 < n → 0 < toAnalysisModel.fluctuationScale n
  /-- Fluctuation scale squared equals sample size (√n property). -/
  fluctuation_scale_sq :
    ∀ (n : Nat), (toAnalysisModel.fluctuationScale n) ^ 2 = n
  /-- Finite averages are permutation-invariant (exchangeability). -/
  finite_average_perm :
    ∀ {n : Nat}, (σ : Equiv.Perm (Fin n)) → (x : Fin n → ℝ) →
      toAnalysisModel.finiteAverage (fun i => x (σ i)) =
      toAnalysisModel.finiteAverage x
  /-- Average of a constant family is that constant. -/
  finite_average_const :
    ∀ {n : Nat}, (c : ℝ) → n ≠ 0 →
      toAnalysisModel.finiteAverage (fun _ : Fin n => c) = c
  /-- Centered constant sequence has zero cumulative fluctuation. -/
  normalized_cumulative_const_zero :
    ∀ (c : ℝ) (N t : Nat), N ≠ 0 →
      toAnalysisModel.normalizedCumulative (fun _ => c) c N t = 0
  /-- Spectral gap is 1 minus the second eigenvalue modulus. -/
  spectral_gap_eq :
    ∀ {State : Type*} [Fintype State] [DecidableEq State]
      (kernel : State → State → ℝ),
      toAnalysisModel.spectralGapValue kernel =
      1 - toAnalysisModel.secondSpectrumValue kernel

/-! ## Extended Analysis Laws: Exported Theorems -/

section AnalysisLaws

variable [AnalysisLaws]

/-- Re-export nonnegativity of the exponential-tail form. -/
theorem exponential_tail_nonneg (σ t : ℝ) : 0 ≤ exponentialTail σ t :=
  AnalysisLaws.exponential_tail_nonneg (self := inferInstance) σ t

/-- Re-export normalization at zero threshold for the exponential-tail form. -/
theorem exponential_tail_zero (σ : ℝ) : exponentialTail σ 0 = 2 :=
  AnalysisLaws.exponential_tail_zero (self := inferInstance) σ

/-- Re-export positivity of fluctuation scale for positive index. -/
theorem fluctuation_scale_pos {n : Nat} (hn : 0 < n) :
    0 < fluctuationScale n :=
  AnalysisLaws.fluctuation_scale_pos (self := inferInstance) hn

/-- Re-export squared-scale identity. -/
theorem fluctuation_scale_sq (n : Nat) :
    (fluctuationScale n) ^ 2 = n :=
  AnalysisLaws.fluctuation_scale_sq (self := inferInstance) n

/-- Re-export permutation invariance of finite averages. -/
theorem finite_average_perm {n : Nat} (σ : Equiv.Perm (Fin n)) (x : Fin n → ℝ) :
    finiteAverage (fun i => x (σ i)) = finiteAverage x :=
  AnalysisLaws.finite_average_perm (self := inferInstance) σ x

/-- Re-export constant-family identity for finite averages. -/
theorem finite_average_const {n : Nat} (c : ℝ) (hn : n ≠ 0) :
    finiteAverage (fun _ : Fin n => c) = c :=
  AnalysisLaws.finite_average_const (self := inferInstance) c hn

/-- Re-export zero fluctuation for centered constant cumulative process. -/
theorem normalized_cumulative_const_zero (c : ℝ) (N t : Nat) (hN : N ≠ 0) :
    normalizedCumulative (fun _ => c) c N t = 0 :=
  AnalysisLaws.normalized_cumulative_const_zero (self := inferInstance) c N t hN

/-- Re-export relation between spectral gap and second spectral value. -/
theorem spectral_gap_eq {State : Type*} [Fintype State] [DecidableEq State]
    (kernel : State → State → ℝ) :
    spectralGapValue kernel = 1 - secondSpectrumValue kernel :=
  AnalysisLaws.spectral_gap_eq (self := inferInstance) kernel

end AnalysisLaws

end EntropyAPI
