/-
  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

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

/-- Information entropy of a single residue site, from Void-X Eq. 3:
    S_i = -SUM_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 -> Real) : Real :=
  -sum 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 := Real.log 50

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

-- =================================================================
-- S2. 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 -> Pi or Lambda state
    Low B-factor  -> ordered        -> low entropy  -> Phi state

    The B-factor is already in the PDB file -- no ML model needed
    for the baseline entropy computation. -/
def entropyFromBFactor (bFactor : Real) (neighborBFactors : List Real) : Real :=
  -- 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 the clusters-by-entity-40 sequence
    cluster identity. Residues in the same cluster have similar
    structural contexts and thus similar entropy profiles. -/
def entropyFromCluster (clusterSize : Nat) (sequenceIdentity : Real) : Real :=
  -- 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)

-- =================================================================
-- S3. 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 x String x Real x List Real))
    : List (AminoAcidToken x Real x BindingSiteState) :=
  residues.map (fun (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 x Real x BindingSiteState)) : Real :=
  let entropies := profile.map (fun (_, 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 x Real x BindingSiteState))
    (globalMin globalMax : Real) : Real :=
  let avg := averageSiteEntropy profile
  100 * (1 - (avg - globalMin) / (globalMax - globalMin))

-- =================================================================
-- S4. 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 8x8 PIST adjacency matrix, which eigendecomposes to
    an 8D spectral profile -> Sidon address. -/
def entropyToSidonAddress (profile : List (AminoAcidToken x Real x BindingSiteState))
    : List Nat :=
  -- Extract top 8 entropy values (one per Hachimoji state category)
  let stateEntropies := List.filterMap (fun (_, e, s) =>
    match s with
    | .Phi => some (0, e) | .Lambda => some (1, e) | .Rho => some (2, e)
    | .Kappa => some (3, e) | .Omega => some (4, e) | .Sigma => some (5, e)
    | .Pi => some (6, e) | .Zeta => some (7, e)
  ) profile
  -- Map to Sidon powers {1, 2, 4, 8, 16, 32, 64, 128}
  -- weighted by entropy magnitude
  stateEntropies.map (fun (idx, e) =>
    Nat.pow 2 idx * (min (Nat.floor (e * 10)) 16)
  )

-- =================================================================
-- S5. 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 -> Real := fisherMetric50 distribution
  entropy : Real := siteEntropy distribution.val
  /-- Geodesic distance to another binding site manifold.
      The exact geodesic requires solving the geodesic equation on Delta^49
      under the Fisher metric, which has no closed form. We use the
      Fisher-Rao approximation (Bhattacharyya-based) as the default.
      TODO: Replace with exact geodesic solver when available. -/
  geodesicDistanceTo : BindingSiteManifold -> Real :=
    fun other => fisherRaoApprox self.distribution other.distribution

/-- Approximate Fisher-Rao distance between two binding sites
    using the Bhattacharyya coefficient (efficient approximation). -/
def fisherRaoApprox (p q : AminoAcidDistribution) : Real :=
  Real.sqrt (2 * Real.log (1 / sum 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.

    STATUS: PROOF BLOCKED on research-level Pinsker-type inequality.
    See entropy_lipschitz axiom below for the missing analytical gap.
    Once that axiom is proven, the proof proceeds by:
      1. Lipschitz bound gives |S(p) - S(q)| < 0.28 (from d_FR < 0.1)
      2. 0.28 < 0.3 = min threshold gap in entropyToHachimoji
      3. Case analysis on thresholds shows same classification. -/
theorem fisher_implies_similar_druggability (p q : AminoAcidDistribution)
    (sites : List (AminoAcidToken x Real x 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 -- BLOCKED: needs entropy_lipschitz + numerical bounds
        -- (see axiom and documentation below)
  where
    dominantState := fun d _ =>
      entropyToHachimoji (siteEntropy d.val) false true
    bothDruggable := fun d1 d2 _ =>
      let s1 := entropyToHachimoji (siteEntropy d1.val) false true
      let s2 := entropyToHachimoji (siteEntropy d2.val) false true
      (s1 = .Pi \/ s1 = .Lambda) /\ (s2 = .Pi \/ s2 = .Lambda)

-- =================================================================
-- S5.1 MISSING LEMMAS (blocking fisher_implies_similar_druggability)
-- =================================================================
-- OPEN PROBLEM: Pinsker-type inequality for Fisher-Rao metric.
--
-- The proof of fisher_implies_similar_druggability requires:
--   1. Entropy is Lipschitz w.r.t. Fisher-Rao distance:
--      |S(p) - S(q)| <= L * d_FR(p,q)  for some L < 4
--      (where L = sqrt(2*log(50)) ~ 2.80 suffices)
--   2. If d_FR(p,q) < 0.1 then |S(p) - S(q)| < 0.28 < 0.3
--      (the minimum threshold gap in entropyToHachimoji)
--   3. Classification stability: |e1 - e2| < 0.3 implies
--      entropyToHachimoji e1 = entropyToHachimoji e2 (case analysis)
--
-- The Lipschitz bound (1) is a research-level analytical result.
-- Once available, (2) is arithmetic and (3) is case analysis.
-- Marking as axiom to unblock downstream proofs when established.

/-- Entropy is Lipschitz continuous w.r.t. Fisher-Rao distance.
    CONSTANT: L = sqrt(2*log(50)) ~ 2.80.
    If d_FR(p,q) < 0.1 then |S(p)-S(q)| < 0.28 < 0.3 (min gap).
    STATUS: Research-level Pinsker-type inequality. -/
axiom entropy_lipschitz (p q : AminoAcidDistribution) :
    abs (siteEntropy p.val - siteEntropy q.val) <= Real.sqrt (2 * maxEntropy50) * fisherRaoApprox p q

end BindingSiteEntropy