tokitai-dl 0.1.2

//! Bridge between god-graph's [`TensorRingCompressor`] and a
//! tokitai-style Lean proof-obligation descriptor.
//!
//! ## Motivation (S6 — tensor ring compression + Lean proof bridge)
//!
//! When god-graph's `TensorRingCompressor::compress_graph` produces a
//! [`CompressionReport`], tokitai-dl wants to emit a corresponding
//! Lean proof-obligation descriptor that an external Lean toolchain
//! could consume and re-check after downstream integration. This is the
//! *emission* half of the bridge; the *reception* half lives in
//! `tokitai-operator`'s
//! `verify::theorems::TheoremBindingRegistry` (which is closed
//! CLAIMS-boundary code — see the notes on registration below).
//!
//! ## What is (and is not) implemented
//!
//! - **Emission**: a synthetic [`LeanProofObligation`] descriptor is built
//!   from a real `CompressionReport`, with the `compression_ratio` taken
//!   directly from the report's field (which itself is derived from
//!   `original_params / compressed_params` inside
//!   `TensorRingCompressor::compress_graph`). Emitting the descriptor does not
//!   prove the theorem.
//! - **Reception**: not implemented in this crate. A real registration
//!   in `tokitai-operator/src/verify/theorems.rs` would be a
//!   `TheoremBindingDescriptor` with
//!   `theorem_id = "tokitai.tensor_ring_compression_bound"`. That file
//!   lives behind the CLAIMS-boundary (`CLAIMS.md` is the source of
//!   truth, and editing it requires a P-number trace entry in
//!   `todo.json`); we deliberately do not touch it here. Downstream
//!   users can add a `TensorRingCompressionTheoremBinding` to that
//!   registry following the pattern of
//!   `padic_valuation_skip_theorem_descriptor` /
//!   `finite_sheaf_gluing_theorem_descriptor`.
//! - **Structural validation**:
//!   [`structurally_validate_obligation`] checks only local metadata
//!   (compression ratio > 1, loss bound finite, etc.). Real Lean
//!   verification is out of scope: the `lean` toolchain is not
//!   invoked by this crate. [`verify_obligation`] is retained as a
//!   compatibility wrapper for callers that used the historical name,
//!   but it does not perform Lean verification.
//!
//! ## Loss-bound honesty
//!
//! god-graph v0.6.0-alpha's [`CompressionReport`] does not expose a
//! per-layer or overall loss bound (it has `original_params`,
//! `compressed_params`, `compression_ratio`, and the per-layer
//! `LayerCompressionReport` list). The default `loss_bound` written
//! into the obligation is therefore `0.0`, with a clear comment
//! alongside the call site. If a future god-graph release adds a
//! Frobenius / spectral / operator-norm field, the
//! [`LeanBound::Frobenius`] / [`LeanBound::Spectral`] /
//! [`LeanBound::Operator`] variants are ready to carry it.
//!
//! This module is gated on both the `graph` and the `operator`
//! features: `graph` brings in `god-graph`'s
//! `transformer::optimization` re-exports, and `operator` keeps the
//! surface aligned with the other bridges in this crate.

use std::time::{SystemTime, UNIX_EPOCH};

use god_graph::transformer::optimization::TensorRingCompressor;
use god_graph::transformer::optimization::tensor_ring::CompressionReport;

use crate::BridgeError;

/// Stable identifier for the tensor-ring compression-bound theorem.
///
/// This is the *proposed* theorem id; the actual registration in
/// `tokitai-operator/src/verify/theorems.rs` is left to the user
/// (that file is behind the CLAIMS-boundary). Downstream code that
/// wants to consume this obligation should add a
/// `TheoremBindingDescriptor { theorem_id: TENSOR_RING_COMPRESSION_THEOREM_ID, .. }`
/// following the same pattern as `padic_valuation_skip_theorem_descriptor`.
///
/// # No example
///
/// String constant whose only purpose is to be embedded in the
/// `theorem_id` field of every [`LeanProofObligation`] emitted by
/// this bridge. A runnable example would just spell out the
/// constant's value (which the source already does); see the
/// example on [`LeanProofObligation`] for a usage site.
pub const TENSOR_RING_COMPRESSION_THEOREM_ID: &str = "tokitai.tensor_ring_compression_bound";

/// Default path to the Lean theorem file that is expected to contain a proof of
/// `TENSOR_RING_COMPRESSION_THEOREM_ID`.
///
/// Matches the convention used elsewhere in tokitai-operator
/// (`PADIC_VALUATION_SKIP_LEAN_FILE = "docs/theorems/padic_valuation_skip.lean"`,
/// `FINITE_SHEAF_GLUING_LEAN_FILE = "docs/theorems/finite_sheaf_gluing.lean"`).
///
/// # No example
///
/// String constant; its only consumer is the `replay_path` field
/// of every emitted [`LeanProofObligation`]. See
/// [`compression_to_lean_obligation`] for the call site.
pub const TENSOR_RING_COMPRESSION_LEAN_FILE: &str =
    "docs/theorems/tensor_ring_compression_bound.lean";

/// A Lean proof-obligation descriptor emitted by the S6 bridge.
///
/// One obligation corresponds to one [`CompressionReport`]; the
/// downstream `TheoremBindingRegistry` could resolve
/// [`theorem_id`](Self::theorem_id) to a Lean theorem and run an
/// external replay verifier at [`replay_path`](Self::replay_path).
/// This crate only constructs and structurally validates the descriptor; it
/// never invokes Lean, and an emitted obligation should not be read as a
/// completed proof.
///
/// # Example
///
/// ```rust
/// # #[cfg(all(feature = "graph", feature = "operator"))]
/// # fn demo() {
/// use tokitai_dl::lean_proof_bridge::{
///     LeanBound, LeanProofObligation, TENSOR_RING_COMPRESSION_LEAN_FILE,
///     TENSOR_RING_COMPRESSION_THEOREM_ID,
/// };
///
/// let obl = LeanProofObligation {
///     theorem_id: TENSOR_RING_COMPRESSION_THEOREM_ID.to_string(),
///     replay_path: TENSOR_RING_COMPRESSION_LEAN_FILE.to_string(),
///     compression_ratio: 2.0,
///     loss_bound: 0.0,
///     bound: LeanBound::Frobenius(0.0),
///     created_at: 0,
/// };
/// assert_eq!(obl.theorem_id, "tokitai.tensor_ring_compression_bound");
/// assert_eq!(obl.bound, LeanBound::Frobenius(0.0));
/// # }
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct LeanProofObligation {
    /// Stable theorem identifier. Resolved by
    /// `tokitai_operator::verify::theorems::TheoremBindingRegistry::lookup`.
    pub theorem_id: String,
    /// Path to the Lean theorem file associated with `theorem_id`.
    /// Documented but not loaded by this crate; the Lean toolchain is
    /// an external dependency.
    pub replay_path: String,
    /// Compression ratio reported by god-graph. Must be `> 1.0` for
    /// the obligation to be considered non-trivial by
    /// [`structurally_validate_obligation`].
    pub compression_ratio: f64,
    /// Loss bound written into the obligation. Defaults to `0.0`
    /// because god-graph's `CompressionReport` does not expose one
    /// (see the module-level "Loss-bound honesty" note).
    pub loss_bound: f64,
    /// Variant tag for the loss bound.
    pub bound: LeanBound,
    /// Unix time (seconds) at which the obligation was created. Used
    /// for staleness / replay-window checks downstream.
    pub created_at: u64,
}

/// Tag for the kind of loss bound carried by a
/// [`LeanProofObligation`].
///
/// The variants are ready to be filled in once god-graph exposes
/// per-layer Frobenius / spectral / operator-norm residuals; today
/// the bridge writes the *scalar* `loss_bound` field as
/// `0.0` and the `LeanBound` variant defaults to
/// [`LeanBound::Frobenius`] (purely a placeholder — no upstream
/// measurement exists).
///
/// # Example
///
/// ```rust
/// # #[cfg(all(feature = "graph", feature = "operator"))]
/// # fn demo() {
/// use tokitai_dl::lean_proof_bridge::LeanBound;
/// let b = LeanBound::Spectral(0.01);
/// assert_eq!(b, LeanBound::Spectral(0.01));
/// // The three variants are distinct even when their payloads agree.
/// assert_ne!(LeanBound::Frobenius(0.0), LeanBound::Operator(0.0));
/// # }
/// ```
#[derive(Debug, Clone, PartialEq)]
pub enum LeanBound {
    /// Frobenius-norm residual: `||A - A_approx||_F`.
    Frobenius(f64),
    /// Spectral-norm residual: `||A - A_approx||_2`.
    Spectral(f64),
    /// Operator-norm residual (same as spectral for finite matrices,
    /// but kept as a distinct variant for clarity).
    Operator(f64),
}

/// Build a [`LeanProofObligation`] from a real god-graph
/// [`CompressionReport`].
///
/// The `compression_ratio` is read directly from the report
/// (`report.compression_ratio`, which god-graph computes as
/// `original_params / compressed_params`). The `loss_bound` is set
/// to `0.0` because god-graph v0.6.0-alpha does not expose a
/// reconstruction-loss field; the chosen `LeanBound` variant is
/// `Frobenius(0.0)` to keep the obligation well-typed.
/// The returned value is metadata for an external proof workflow; this
/// function does not run Lean and does not establish the bound.
///
/// # Errors
///
/// - [`BridgeError::Upstream`] is never returned today, but the
///   `Result` signature is kept so future god-graph releases can
///   surface errors (e.g. a malformed ratio) without breaking the
///   public surface.
///
/// # Example
///
/// ```rust
/// # #[cfg(all(feature = "graph", feature = "operator"))]
/// # fn demo() {
/// use god_graph::transformer::optimization::tensor_ring::{
///     CompressionReport, LayerCompressionReport,
/// };
/// use tokitai_dl::lean_proof_bridge::{
///     compression_to_lean_obligation, structurally_validate_obligation,
/// };
///
/// // A synthetic 1024->256 CompressionReport (ratio = 4.0).
/// let report = CompressionReport {
///     original_params: 1024,
///     compressed_params: 256,
///     compression_ratio: 4.0,
///     layers: vec![LayerCompressionReport {
///         layer_name: "demo.weight".into(),
///         original_params: 1024,
///         compressed_params: 256,
///         compression_ratio: 4.0,
///         ranks: vec![8, 8, 8],
///     }],
/// };
/// let obl = compression_to_lean_obligation(report).unwrap();
/// assert!(obl.compression_ratio > 1.0);
/// assert!(structurally_validate_obligation(&obl));
/// # }
/// ```
pub fn compression_to_lean_obligation(
    compression: CompressionReport,
) -> Result<LeanProofObligation, BridgeError> {
    #[cfg(feature = "tracing")]
    tracing::trace!(
        "lean_proof_bridge::compression_to_lean_obligation: entry (compression_ratio={})",
        compression.compression_ratio
    );
    #[cfg(all(feature = "observability", not(feature = "tracing")))]
    log::trace!(
        "lean_proof_bridge::compression_to_lean_obligation: entry (compression_ratio={})",
        compression.compression_ratio
    );
    Ok(LeanProofObligation {
        theorem_id: TENSOR_RING_COMPRESSION_THEOREM_ID.to_string(),
        replay_path: TENSOR_RING_COMPRESSION_LEAN_FILE.to_string(),
        compression_ratio: compression.compression_ratio,
        // Loss bound is unknown upstream; see the module-level
        // "Loss-bound honesty" note.
        loss_bound: 0.0,
        bound: LeanBound::Frobenius(0.0),
        created_at: unix_now_secs(),
    })
}

/// Build a [`LeanProofObligation`] for a single tensor compressed by
/// a [`TensorRingCompressor`] without going through the full
/// `compress_graph` path.
///
/// This is the synthetic path used by tests and the example: it
/// runs `compressor.decompose(&tensor)`, extracts the per-tensor
/// compression ratio from the resulting `TensorRing`, and emits an
/// obligation descriptor without an upstream `CompressionReport`. It exists
/// so the bridge can be exercised on a single weight matrix (which
/// is the only case that doesn't need a full safetensors load).
/// It still does not run Lean or prove the compression bound.
///
/// # Errors
///
/// - [`BridgeError::Upstream`] wraps any god-graph
///   `TensorError` (e.g. dimension mismatch on a non-2D tensor).
///
/// # Example
///
/// ```rust
/// # #[cfg(all(feature = "graph", feature = "operator"))]
/// # fn demo() {
/// use god_graph::tensor::DenseTensor;
/// use god_graph::transformer::optimization::{CompressionConfig, TensorRingCompressor};
/// use tokitai_dl::lean_proof_bridge::{
///     structurally_validate_obligation, tensor_to_lean_obligation,
/// };
///
/// // A 64x64 F64 weight tensor compressed with target rank 2.
/// let config = CompressionConfig::new()
///     .with_target_ranks(vec![2])
///     .with_min_rank(2)
///     .with_max_rank(4);
/// let compressor = TensorRingCompressor::new(config);
/// let tensor = DenseTensor::from_vec(vec![1.0_f64; 64 * 64], vec![64, 64]);
/// let obl = tensor_to_lean_obligation(&compressor, &tensor).unwrap();
/// assert!(obl.compression_ratio > 1.0);
/// assert!(structurally_validate_obligation(&obl));
/// # }
/// ```
pub fn tensor_to_lean_obligation(
    compressor: &TensorRingCompressor,
    tensor: &god_graph::tensor::DenseTensor,
) -> Result<LeanProofObligation, BridgeError> {
    let ring = compressor
        .decompose(tensor)
        .map_err(|e| BridgeError::Upstream(format!("tensor ring decompose failed: {e}")))?;
    Ok(LeanProofObligation {
        theorem_id: TENSOR_RING_COMPRESSION_THEOREM_ID.to_string(),
        replay_path: TENSOR_RING_COMPRESSION_LEAN_FILE.to_string(),
        compression_ratio: ring.compression_ratio(),
        loss_bound: 0.0,
        bound: LeanBound::Frobenius(0.0),
        created_at: unix_now_secs(),
    })
}

/// Structurally validate a [`LeanProofObligation`].
///
/// Without a Lean toolchain in this crate's build graph, this
/// function does *not* invoke the Lean compiler. Instead it returns
/// `true` iff the obligation looks non-trivial:
///
/// - `compression_ratio` is finite and `> 1.0` (a ratio `<= 1.0`
///   means compression didn't actually compress, so the obligation
///   is meaningless; `NaN` is rejected by the finite check);
/// - `loss_bound` is finite and `>= 0.0` (`NaN`, `-infinity`, and
///   `+infinity` are all rejected);
/// - the [`LeanBound`] selector carries a value that is also finite
///   and `>= 0.0` (so an obligation whose tag is `Frobenius(+inf)`
///   is rejected even if the scalar field is `0.0`).
///
/// A real Lean verification step (e.g. shelling out to
/// `lean docs/theorems/tensor_ring_compression_bound.lean`) is out
/// of scope for this crate. Downstream code that wants a stronger
/// guarantee should add it on top of this structural check.
///
/// # Caveat
///
/// `loss_bound` is currently always `0.0` (see the module-level
/// "Loss-bound honesty" note). `structurally_validate_obligation`
/// therefore accepts every obligation that has a sensible compression
/// ratio. That is intentional: this is a structural metadata gate, not
/// a Lean proof check or soundness gate.
///
/// # Example
///
/// ```rust
/// # #[cfg(all(feature = "graph", feature = "operator"))]
/// # fn demo() {
/// use tokitai_dl::lean_proof_bridge::{
///     LeanBound, LeanProofObligation, TENSOR_RING_COMPRESSION_LEAN_FILE,
///     TENSOR_RING_COMPRESSION_THEOREM_ID, structurally_validate_obligation,
/// };
///
/// // ratio=2.0, loss_bound=0.01, LeanBound::Spectral(0.01) — all checks pass.
/// let good = LeanProofObligation {
///     theorem_id: TENSOR_RING_COMPRESSION_THEOREM_ID.to_string(),
///     replay_path: TENSOR_RING_COMPRESSION_LEAN_FILE.to_string(),
///     compression_ratio: 2.0,
///     loss_bound: 0.01,
///     bound: LeanBound::Spectral(0.01),
///     created_at: 0,
/// };
/// assert!(structurally_validate_obligation(&good));
///
/// // ratio=0.0 — rejected (compression did not compress).
/// let bad = LeanProofObligation {
///     compression_ratio: 0.0,
///     ..good.clone()
/// };
/// assert!(!structurally_validate_obligation(&bad));
/// # }
/// ```
pub fn structurally_validate_obligation(obl: &LeanProofObligation) -> bool {
    // --- compression ratio -------------------------------------------------
    if !obl.compression_ratio.is_finite() {
        return false;
    }
    if obl.compression_ratio <= 1.0 {
        return false;
    }
    // --- loss bound (scalar field) ----------------------------------------
    if !obl.loss_bound.is_finite() {
        return false;
    }
    if obl.loss_bound < 0.0 {
        return false;
    }
    // --- LeanBound selector (must agree with the scalar) -------------------
    let bound_value = match obl.bound {
        LeanBound::Frobenius(v) | LeanBound::Spectral(v) | LeanBound::Operator(v) => v,
    };
    if !bound_value.is_finite() {
        return false;
    }
    if bound_value < 0.0 {
        return false;
    }
    true
}

/// Compatibility wrapper for [`structurally_validate_obligation`].
///
/// The historical name is kept to avoid breaking callers. It performs
/// the same local structural validation and does **not** invoke Lean or
/// prove the theorem named by [`LeanProofObligation::theorem_id`].
pub fn verify_obligation(obl: &LeanProofObligation) -> bool {
    structurally_validate_obligation(obl)
}

/// Current Unix time in seconds, or `0` if the system clock is set
/// before the Unix epoch (defensive default — should not happen in
/// practice).
fn unix_now_secs() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}