/-
  BindingSiteEntropy.lean — Information Entropy for Protein Binding Sites

  Computes the information entropy profile of a binding site using
  the Fisher information metric (guaranteed unique by Chentsov).
  This is the direct protein-structure analog of the spectral profile
  pipeline (eigensolid_pipeline.py) for equations.

  References:
    - Yang, Yuan, Chou 2025 (Void-X): Eq. 3 (information entropy)
    - Giani, Win, Conti 2025: quantum discrimination via PVGS
    - Research-Stack library/ChentsovFinite.lean: metric uniqueness
-/}

import Mathlib
import BindingSiteHachimoji

namespace BindingSiteEntropy

open BindingSiteHachimoji

-- =================================================================
-- §1. INFORMATION ENTROPY PER RESIDUE SITE (Void-X Eq. 3)
-- =================================================================

/-- Information entropy of a single residue site, from Void-X Eq. 3:
    S_i = -Σ_j p(a_j | context) log p(a_j | context)

    where p(a_j | context) is the conditional probability of atom
    type a_j given the structural context (neighboring atoms).

    In Void-X, this is computed from the diffusion model's output
    distribution over 50 atom types. Here we formalize it as a
    probability distribution over the Hachimoji state space. -/
def siteEntropy (probDist : Fin 50 → ℝ) : ℝ :=
  -∑ i, if probDist i > 0 then probDist i * Real.log (probDist i) else 0

/-- The maximum possible entropy for 50 states (uniform distribution).
    S_max = log(50) ≈ 3.912. -/
def maxEntropy50 : ℝ := Real.log 50

/-- Normalized entropy: S* = S_i / S_max ∈ [0, 1].
    This is what Void-X uses for the bindability score. -/
def normalizedEntropy (probDist : Fin 50 → ℝ) : ℝ :=
  siteEntropy probDist / maxEntropy50

-- =================================================================
-- §2. ENTROPY FROM PDB STRUCTURE
-- =================================================================

/-- Extract the amino acid distribution at a residue position from
    a PDB structure. This reads the B-factors (temperature factors)
    as a proxy for positional uncertainty, which maps to entropy.

    High B-factor → high uncertainty → high entropy → Π or Λ state
    Low B-factor  → ordered        → low entropy  → Φ state

    The B-factor is already in the PDB file — no ML model needed
    for the baseline entropy computation. -/
def entropyFromBFactor (bFactor : ℝ) (neighborBFactors : List ℝ) : ℝ :=
  -- Local average B-factor normalizes by neighborhood context
  let localAvg := (bFactor + List.sum neighborBFactors) / (1 + neighborBFactors.length)
  -- Map to [0, 1]: higher B-factor = higher entropy
  Real.log (1 + localAvg) / Real.log (1 + 100)
  -- Dividing by log(101) since B-factors typically range 0-100

/-- Alternative: entropy from theclusters-by-entity-40 sequence
    cluster identity. Residues in the same cluster have similar
    structural contexts and thus similar entropy profiles. -/
def entropyFromCluster (clusterSize : ℕ) (sequenceIdentity : ℝ) : ℝ :=
  -- High sequence identity within cluster → low entropy (conserved)
  -- Large cluster size → high diversity → higher entropy
  let diversity := Real.log (1 + clusterSize)
  let conservation := sequenceIdentity
  diversity * (1 - conservation)

-- =================================================================
-- §3. BINDING SITE ENTROPY PROFILE
-- =================================================================

/-- Compute the full entropy profile of a binding site from a
    sequence of residue data (PDB-derived or Void-X-generated).

    This is the protein-structure analog of `spectralProfile` in
    eigensolid_pipeline.py. -/
def bindingSiteEntropyProfile (residues : List (String × String × ℝ × List ℝ))
    : List (AminoAcidToken × ℝ × BindingSiteState) :=
  residues.map (λ (residueType, mod, bFactor, neighborBFs) =>
    let token := residueToToken residueType mod
    let entropy := entropyFromBFactor bFactor neighborBFs
    let state := entropyToHachimoji entropy false true
    (token, entropy, state)
  )

/-- Average entropy of a binding site (the main metric from Void-X). -/
def averageSiteEntropy (profile : List (AminoAcidToken × ℝ × BindingSiteState)) : ℝ :=
  let entropies := profile.map (λ (_, e, _) => e)
  match entropies with
  | [] => 0
  | es => List.sum es / es.length

/-- Bindability score B* from Yang et al. 2025 (Fig. S8):
    B* = 100 × [1 - (S* - min(S)) / (max(S) - min(S))]

    High B* = low entropy relative to the protein surface =
    ordered pocket suitable for ligand binding. -/
def bindabilityScore (profile : List (AminoAcidToken × ℝ × BindingSiteState))
    (globalMin globalMax : ℝ) : ℝ :=
  let avg := averageSiteEntropy profile
  100 * (1 - (avg - globalMin) / (globalMax - globalMin))

-- =================================================================
-- §4. SIDON ADDRESS FOR BINDING SITE (from existing library)
-- =================================================================

/-- A binding site gets a Sidon address from its entropy profile,
    exactly like an equation gets a Sidon address from its spectral
    profile (eigensolid_pipeline.py).

    The 8 dominant entropy values are the "observables" that feed
    into the 8×8 PIST adjacency matrix, which eigendecomposes to
    an 8D spectral profile → Sidon address. -/
def entropyToSidonAddress (profile : List (AminoAcidToken × ℝ × BindingSiteState))
    : List ℕ :=
  -- Extract top 8 entropy values (one per Hachimoji state category)
  let stateEntropies := List.filterMap (λ (_, e, s) =>
    match s with
    | .Φ => some (0, e) | .Λ => some (1, e) | .Ρ => some (2, e)
    | .Κ => some (3, e) | .Ω => some (4, e) | .Σ => some (5, e)
    | .Π => some (6, e) | .Ζ => some (7, e)
  ) profile
  -- Map to Sidon powers {1, 2, 4, 8, 16, 32, 64, 128}
  -- weighted by entropy magnitude
  stateEntropies.map (λ (idx, e) =>
    Nat.pow 2 idx * (min (Nat.floor (e * 10)) 16)
  )

-- =================================================================
-- §5. FISHER METRIC ON BINDING SITE MANIFOLD
-- =================================================================

/-- The binding site manifold: probability distributions over
    residue tokens, equipped with the Fisher metric.
    Geodesics on this manifold are evolutionarily optimal paths. -/
structure BindingSiteManifold where
  distribution : AminoAcidDistribution
  metric : Fin 50 → Fin 50 → ℝ := fisherMetric50 distribution
  entropy : ℝ := siteEntropy distribution.val
  geodesicDistance : BindingSiteManifold → ℝ := sorry
  -- Geodesic distance requires solving the geodesic equation on Δ^49.
  -- This is computationally intensive; use approximation for now.

/-- Approximate Fisher-Rao distance between two binding sites
    using the Bhattacharyya coefficient (efficient approximation). -/
def fisherRaoApprox (p q : AminoAcidDistribution) : ℝ :=
  Real.sqrt (2 * Real.log (1 / ∑ i, Real.sqrt (p.val i * q.val i)))

/-- Theorem: nearby binding sites (small Fisher distance) have
    similar druggability profiles. This is what enables
    transfer learning across protein families. -/
theorem fisher_implies_similar_druggability (p q : AminoAcidDistribution)
    (sites : List (AminoAcidToken × ℝ × BindingSiteState))
    (h : fisherRaoApprox p q < 0.1) :
    -- Sites with similar distributions have similar dominant states
    (dominantState p sites = dominantState q sites) ∨
    (bothDruggable p q sites) := by
  sorry -- Proof: relies on continuity of entropy w.r.t. Fisher metric
        -- and the classification threshold structure of entropyToHachimoji.
  where
    dominantState := λ d _ =>
      entropyToHachimoji (siteEntropy d.val) false true
    bothDruggable := λ d1 d2 _ =>
      let s1 := entropyToHachimoji (siteEntropy d1.val) false true
      let s2 := entropyToHachimoji (siteEntropy d2.val) false true
      (s1 = .Π ∨ s1 = .Λ) ∧ (s2 = .Π ∨ s2 = .Λ)

end BindingSiteEntropy