cobre-io 0.15.0

//! Input and output record types for value-function artifact serialization.
//!
//! Input types (`PolicyCutRecord`, `PolicyBasisRecord`, `StageStatesPayload`,
//! `StageCutsPayload`) borrow from caller-owned buffers; owned output types
//! (`Owned*`, `*ReadResult`, `PolicyCheckpoint`) own their vectors. All use
//! generic names to maintain infrastructure crate genericity; conversion from
//! algorithm-specific types is the calling crate's responsibility. Field names
//! correspond to the tables in `schemas/policy.fbs`.

/// Current on-disk value-function artifact format version.
///
/// [`CheckpointManifest::format_version`] must equal this;
/// [`crate::read_policy_checkpoint`] rejects any other value — and absence —
/// with a named error before parsing any payload, so a pre-marker artifact is
/// cleanly rejected, never read positionally.
pub const FORMAT_VERSION: u32 = 1;

/// Sentinel [`EntitySlot::delivery_date`] value for a slot with no
/// delivery/arrival calendar semantics; also the value a reader yields when the
/// field is absent from a pre-`id:5` buffer (forward-compatible default).
pub const ENTITY_SLOT_DELIVERY_DATE_SENTINEL: i32 = i32::MIN;

/// One per-slot entity-identity record for a state-vector dimension.
///
/// `entity_type` is the raw discriminant byte of the `EntityType` enum in
/// `schemas/policy.fbs`; [`EntitySlot::family`] reads it as the typed
/// [`StateFamily`].
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct EntitySlot {
    /// Raw [`StateFamily`] discriminant byte; see [`EntitySlot::family`].
    pub entity_type: u8,
    /// Owning entity's id; `int32` because a sentinel id can be `-1`.
    pub entity_id: i32,
    /// Secondary index within the owning entity (per-type meaning is the caller's).
    pub subindex: u32,
    /// Whether the owning entity was operationally active at this slot's stage.
    pub was_active: bool,
    /// Canonical absolute delivery/arrival calendar date for this slot, encoded
    /// `YYYYMMDD` (`year * 10000 + month * 100 + day`);
    /// [`ENTITY_SLOT_DELIVERY_DATE_SENTINEL`] when the slot has no delivery
    /// semantics. Which calendar date maps to a slot is the calling crate's
    /// responsibility, as with `subindex`.
    pub delivery_date: i32,
}

/// State-vector dimension class of an [`EntitySlot`] — the typed Rust view of
/// the `EntityType` enum in `schemas/policy.fbs`. The wire representation stays
/// the raw [`EntitySlot::entity_type`] byte; this enum is the checked reading of
/// it, so the discriminants MUST match the schema.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[repr(u8)]
pub enum StateFamily {
    /// Reservoir storage volume (`subindex` is `0`).
    HydroStorage = 0,
    /// Hydro inflow AR lag (`subindex` is the 1-based AR lag order).
    HydroInflowLag = 1,
    /// Anticipated thermal commitment (`subindex` is the ring-buffer slot).
    AnticipatedThermalState = 2,
    /// Water in-transit bucket (`entity_id` is the downstream hydro, `subindex`
    /// the maturity lag).
    HydroTransitBucket = 3,
}

impl StateFamily {
    /// The raw `EntityType` discriminant byte for this family.
    #[must_use]
    pub const fn code(self) -> u8 {
        self as u8
    }

    /// The family for a raw `EntityType` byte, or `None` for a discriminant no
    /// `schemas/policy.fbs` `EntityType` variant defines.
    #[must_use]
    pub const fn from_code(code: u8) -> Option<Self> {
        match code {
            0 => Some(Self::HydroStorage),
            1 => Some(Self::HydroInflowLag),
            2 => Some(Self::AnticipatedThermalState),
            3 => Some(Self::HydroTransitBucket),
            _ => None,
        }
    }
}

impl EntitySlot {
    /// The typed [`StateFamily`] of this slot, or `None` when [`Self::entity_type`]
    /// is a byte no `EntityType` variant defines.
    #[must_use]
    pub fn family(&self) -> Option<StateFamily> {
        StateFamily::from_code(self.entity_type)
    }
}

/// One affine-piece record for value-function artifact serialization.
///
/// `'a` borrows the coefficient slice without copying (vectors can be large).
#[derive(Debug, Clone)]
pub struct PolicyCutRecord<'a> {
    /// Unique identifier for this piece across all iterations.
    pub cut_id: u64,
    /// LP row position (required for artifact reproducibility).
    pub slot_index: u32,
    /// Training iteration that generated this piece.
    pub iteration: u32,
    /// Forward pass index within the generating iteration.
    pub forward_pass_index: u32,
    /// Pre-computed affine intercept.
    pub intercept: f64,
    /// Gradient coefficients, length must equal `state_dimension`.
    ///
    /// Positional only: index `i` is the i-th state-vector dimension, whose
    /// identity is carried by slot `i` of the co-located [`EntitySlot`] manifest
    /// (`entity_manifest`); no labels are stored inline.
    pub coefficients: &'a [f64],
    /// Whether this piece is currently active in the LP.
    pub is_active: bool,
}

/// One stage's solver basis for value-function artifact serialization.
#[derive(Debug, Clone)]
pub struct PolicyBasisRecord<'a> {
    /// Stage index (0-based).
    pub stage_id: u32,
    /// Training iteration that produced this basis.
    pub iteration: u32,
    /// One status code per LP column (variable). Encoding is solver-specific.
    pub column_status: &'a [u8],
    /// One status code per LP row (constraint). Encoding is solver-specific.
    pub row_status: &'a [u8],
    /// Number of trailing rows in `row_status` that correspond to affine-piece rows.
    pub num_cut_rows: u32,
}

/// Sentinel [`StageStatesPayload::node_id`]/[`StageStatesReadResult::node_id`]
/// value for a policy-graph node identity absent from the write path (a
/// caller that never resolved one) or from a pre-`id:5` buffer
/// (forward-compatible default).
pub const STAGE_STATES_NODE_ID_SENTINEL: i32 = -1;

/// Sentinel [`StageCutsPayload::node_id`]/[`StageCutsReadResult::node_id`] value
/// for a pool with no single owning node (a shared pool, never a boundary
/// source) or a pre-`id:8` buffer (forward-compatible default).
pub const STAGE_CUTS_NODE_ID_SENTINEL: i32 = -1;

/// Sentinel [`StageCutsPayload::graph_stage_id`]/[`StageCutsReadResult::graph_stage_id`]
/// value for an unresolved owning-stage key or a pre-`id:8` buffer
/// (forward-compatible default).
pub const STAGE_CUTS_GRAPH_STAGE_ID_SENTINEL: i32 = -1;

/// Payload for writing per-stage visited states to a value-function artifact.
///
/// The `data` slice contains the flat state vectors (row-major, each of length
/// `state_dimension`). The total number of stored states is `count`.
#[derive(Debug, Clone)]
pub struct StageStatesPayload<'a> {
    /// Study stage index (0-based).
    pub stage_id: u32,
    /// Policy-graph node identity (the declared node id on a branching graph;
    /// [`STAGE_STATES_NODE_ID_SENTINEL`] when absent). Distinct from
    /// `stage_id` the moment a graph carries more than one node per stage.
    pub node_id: i32,
    /// Length of each state vector.
    pub state_dimension: u32,
    /// Number of states stored.
    pub count: u32,
    /// Flat data buffer: `count * state_dimension` f64 elements.
    pub data: &'a [f64],
    /// Per-slot entity identity; length equals `state_dimension` when populated.
    /// An empty slice means no manifest is written.
    pub entity_manifest: &'a [EntitySlot],
}

/// Per-pool affine-piece data payload for [`crate::write_policy_checkpoint`],
/// grouping the arguments of [`crate::serialize_stage_cuts`].
#[derive(Debug)]
pub struct StageCutsPayload<'a> {
    /// Pool id (0-based) — the storage-unit key naming this payload's file
    /// `cuts/<pool>.bin`. Equals the stage index on a chain.
    pub stage_id: u32,
    /// Number of state variables; determines coefficient vector length per piece.
    pub state_dimension: u32,
    /// Total preallocated affine-piece slots in the pool.
    pub capacity: u32,
    /// Number of slots `[0..warm_start_count)` loaded from a prior artifact.
    pub warm_start_count: u32,
    /// Slice of affine-piece records to serialize; length equals `populated_count`.
    pub cuts: &'a [PolicyCutRecord<'a>],
    /// Indices of pieces currently active in the LP.
    pub active_cut_indices: &'a [u32],
    /// Number of filled slots in the pool.
    pub populated_count: u32,
    /// Per-slot entity identity; length equals `state_dimension` when populated.
    /// An empty slice means no manifest is written.
    pub entity_manifest: &'a [EntitySlot],
    /// Objective cost-scale factor the writing study resolved; the provenance
    /// marker making each affine piece scale-independent at rest.
    pub cost_scale_factor: f64,
    /// Owning node's policy-graph id, or [`STAGE_CUTS_NODE_ID_SENTINEL`] for a
    /// shared pool.
    pub node_id: i32,
    /// Graph-stage id of the node(s) owning this pool — the boundary-resolution
    /// key; [`STAGE_CUTS_GRAPH_STAGE_ID_SENTINEL`] when unresolved.
    pub graph_stage_id: i32,
}

/// One node of the value-function artifact's graph manifest: its declared id,
/// the stage it sits at, and the pool whose payload carries its affine pieces.
///
/// `pool_id` **is** the node → pool map: a node references one pool, and a
/// reader resolves node `id`'s pieces as `pool_id`'s payload (leaf nodes sharing
/// a pool all name the same `pool_id`).
#[derive(Debug, Clone)]
pub struct ManifestNode {
    /// Declared node id.
    pub id: i32,
    /// Stage id this node sits at.
    pub stage_id: i32,
    /// Pool whose payload holds this node's affine pieces.
    pub pool_id: u32,
}

/// One directed edge of the graph manifest, with its transition probability.
#[derive(Debug, Clone)]
pub struct ManifestEdge {
    /// Source node id.
    pub source_id: i32,
    /// Target node id.
    pub target_id: i32,
    /// Transition probability `P(source -> target)`.
    pub probability: f64,
}

/// The graph manifest: the node list (each node carrying its own node → pool
/// assignment), the edge list, and the pool-set size.
///
/// This is the identity source the positional format never had — a reader
/// resolves a node's payload through it (node `n`'s pieces are `pool(n)`'s
/// payload), rather than trusting a filename.
#[derive(Debug, Clone, Default)]
pub struct GraphManifest {
    /// Number of distinct pools (the pool-set size).
    pub n_pools: u32,
    /// Every node, in canonical order, each with its stage and pool.
    pub nodes: Vec<ManifestNode>,
    /// Every directed edge with its transition probability.
    pub edges: Vec<ManifestEdge>,
}

/// Producer-namespaced metadata: everything specific to how the artifact was
/// produced (the training algorithm's own recorded state). Segregated from the
/// neutral core carried on [`CheckpointManifest`], whose doc states the
/// segregation rationale.
#[derive(Debug, Clone)]
pub struct ProducerBlock {
    /// Number of training iterations completed at write time.
    pub completed_iterations: u32,
    /// Lower bound value after the final completed iteration.
    pub final_lower_bound: f64,
    /// Last iteration's upper bound, if available (the final value, not a
    /// min-tracked best).
    pub best_upper_bound: Option<f64>,
    /// Maximum number of iterations configured for the run.
    pub max_iterations: u32,
    /// Number of forward passes per iteration.
    pub forward_passes: u32,
    /// Number of pieces loaded from a previous artifact at run start.
    pub warm_start_cuts: u32,
    /// Per-pool warm-start piece counts, in pool-id order.
    ///
    /// When non-empty, supersedes [`warm_start_cuts`] for per-pool accuracy.
    ///
    /// [`warm_start_cuts`]: Self::warm_start_cuts
    pub warm_start_counts: Vec<u32>,
    /// RNG seed used by the scenario sampler.
    ///
    /// Per-draw seeds are derived from `(rng_seed, iteration, scenario, stage)`,
    /// so resume needs only the seed — no accumulated RNG state is persisted.
    pub rng_seed: u64,
    /// Total visited states across all nodes.
    pub total_visited_states: u64,
    /// Block mode the artifact was trained under: the shared lowercase mode when
    /// every study stage agrees, else `"mixed"`.
    pub training_block_mode: String,
    /// Per-study-stage training block modes, in study-stage order.
    ///
    /// Populated only for mixed-mode studies.
    pub training_block_mode_per_stage: Vec<String>,
    /// Objective cost-scale factor the writing study resolved
    /// (`modeling.cost_scale_factor`) — the provenance marker that makes piece
    /// `coefficients`/`intercept` scale-independent at rest (canonical
    /// currency units, not the writer's internal scaled cost space).
    ///
    /// Absent when unmarked; a missing marker is interpreted as
    /// scaled-at-`1_000_000.0`, the constant every unmarked artifact was
    /// unconditionally written under.
    pub cost_scale_factor: Option<f64>,
}

/// Study-global checkpoint metadata carried on the `FlatBuffers`
/// `CheckpointManifest` root at `manifest.bin`: the neutral core
/// (`format_version`, `cobre_version`, `created_at`, `num_stages`, the
/// [`GraphManifest`] descriptors) plus the namespaced [`ProducerBlock`]. Read
/// first by [`crate::read_policy_checkpoint`], whose version gate rejects a
/// stale `format_version` before any payload is parsed.
///
/// The neutral core describes the artifact itself; the algorithm's own recorded
/// state lives under the namespaced [`producer`] block, so a reader that does
/// not know the producer reads the core from the core's own vocabulary.
///
/// [`producer`]: Self::producer
#[derive(Debug, Clone)]
pub struct CheckpointManifest {
    /// On-disk format version; must equal [`FORMAT_VERSION`] on read.
    pub format_version: u32,
    /// Cobre crate version that wrote this checkpoint.
    pub cobre_version: String,
    /// ISO 8601 timestamp when the checkpoint was written.
    pub created_at: String,
    /// Number of stages the graph manifest spans.
    pub num_stages: u32,
    /// Graph manifest: node list, edge list, node → pool map, and pool-set size.
    pub graph_manifest: GraphManifest,
    /// Producer-namespaced metadata (the training algorithm's own state).
    pub producer: ProducerBlock,
}

// ── Owned output types for deserialization ───────────────────────────────────

/// Owned version of [`PolicyCutRecord`] returned by [`crate::deserialize_stage_cuts`].
///
/// Unlike [`PolicyCutRecord<'a>`], this type owns its `coefficients` vector so it
/// can be returned from a deserialization function that does not borrow from the
/// input buffer.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct OwnedPolicyCutRecord {
    /// Unique identifier for this piece across all iterations.
    pub cut_id: u64,
    /// LP row position (required for artifact reproducibility).
    pub slot_index: u32,
    /// Training iteration that generated this piece.
    pub iteration: u32,
    /// Forward pass index within the generating iteration.
    pub forward_pass_index: u32,
    /// Pre-computed affine intercept.
    pub intercept: f64,
    /// Gradient coefficients; positional per the [`PolicyCutRecord::coefficients`] contract.
    pub coefficients: Vec<f64>,
    /// Whether this piece is currently active in the LP.
    pub is_active: bool,
}

/// Owned version of [`PolicyBasisRecord`] returned by [`crate::deserialize_stage_basis`].
///
/// Unlike [`PolicyBasisRecord<'a>`], this type owns its status byte vectors so it
/// can be returned from a deserialization function.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct OwnedPolicyBasisRecord {
    /// Stage index (0-based).
    pub stage_id: u32,
    /// Training iteration that produced this basis.
    pub iteration: u32,
    /// One status code per LP column (variable). Encoding is solver-specific.
    pub column_status: Vec<u8>,
    /// One status code per LP row (constraint). Encoding is solver-specific.
    pub row_status: Vec<u8>,
    /// Number of trailing rows in `row_status` that correspond to affine-piece rows.
    pub num_cut_rows: u32,
}

/// Stage-level metadata and affine-piece records returned by [`crate::deserialize_stage_cuts`].
///
/// Contains the stage-level fields stored in the `StageCuts` root table plus the
/// vector of deserialized affine-piece records.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct StageCutsReadResult {
    /// Pool id (0-based), as written by the pool-keyed payload.
    pub stage_id: u32,
    /// Number of state variables; equals the length of each piece's `coefficients` vector.
    pub state_dimension: u32,
    /// Total preallocated affine-piece slots in the pool.
    pub capacity: u32,
    /// Number of slots loaded from a prior artifact.
    pub warm_start_count: u32,
    /// Number of filled slots in the pool.
    pub populated_count: u32,
    /// Deserialized affine-piece records.
    pub cuts: Vec<OwnedPolicyCutRecord>,
    /// Per-slot entity identity; empty when the field is absent from the buffer.
    pub entity_manifest: Vec<EntitySlot>,
    /// Cost-scale provenance factor; `None` when absent from a pre-`id:8` buffer.
    pub cost_scale_factor: Option<f64>,
    /// Owning node's policy-graph id; [`STAGE_CUTS_NODE_ID_SENTINEL`] for a shared
    /// pool or a pre-`id:8` buffer.
    pub node_id: i32,
    /// Graph-stage id key; [`STAGE_CUTS_GRAPH_STAGE_ID_SENTINEL`] when unresolved
    /// or absent from a pre-`id:8` buffer.
    pub graph_stage_id: i32,
}

/// Owned version of [`StageStatesPayload`] returned by [`crate::deserialize_stage_states`].
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct StageStatesReadResult {
    /// Study stage index (0-based).
    pub stage_id: u32,
    /// Policy-graph node identity; [`STAGE_STATES_NODE_ID_SENTINEL`] when the
    /// field is absent from the buffer (a pre-`id:5` artifact).
    pub node_id: i32,
    /// Length of each state vector.
    pub state_dimension: u32,
    /// Number of states stored.
    pub count: u32,
    /// Flat data buffer (owned).
    pub data: Vec<f64>,
    /// Per-slot entity identity; empty when the field is absent from the buffer.
    pub entity_manifest: Vec<EntitySlot>,
}

/// Complete deserialized value-function artifact returned by [`crate::read_policy_checkpoint`].
#[derive(Debug, Clone)]
pub struct PolicyCheckpoint {
    /// Checkpoint manifest read from `manifest.bin`.
    pub metadata: CheckpointManifest,
    /// Per-pool affine-piece collections, sorted by pool id.
    pub stage_cuts: Vec<StageCutsReadResult>,
    /// Per-stage solver bases, sorted by `stage_id`.
    pub stage_bases: Vec<OwnedPolicyBasisRecord>,
    /// Per-stage visited states, sorted by `stage_id`.
    ///
    /// Empty when the artifact was written without visited states.
    pub stage_states: Vec<StageStatesReadResult>,
}

#[cfg(test)]
mod tests {
    use super::{ENTITY_SLOT_DELIVERY_DATE_SENTINEL, EntitySlot, StateFamily};

    #[test]
    fn state_family_codes_match_policy_fbs_entity_type() {
        assert_eq!(StateFamily::HydroStorage.code(), 0);
        assert_eq!(StateFamily::HydroInflowLag.code(), 1);
        assert_eq!(StateFamily::AnticipatedThermalState.code(), 2);
        assert_eq!(StateFamily::HydroTransitBucket.code(), 3);
    }

    #[test]
    fn state_family_from_code_round_trips_and_rejects_unknown() {
        for family in [
            StateFamily::HydroStorage,
            StateFamily::HydroInflowLag,
            StateFamily::AnticipatedThermalState,
            StateFamily::HydroTransitBucket,
        ] {
            assert_eq!(StateFamily::from_code(family.code()), Some(family));
        }
        assert_eq!(StateFamily::from_code(4), None);
        assert_eq!(StateFamily::from_code(u8::MAX), None);
    }

    #[test]
    fn entity_slot_family_reads_the_raw_byte() {
        let slot = EntitySlot {
            entity_type: StateFamily::AnticipatedThermalState.code(),
            entity_id: 7,
            subindex: 0,
            was_active: true,
            delivery_date: ENTITY_SLOT_DELIVERY_DATE_SENTINEL,
        };
        assert_eq!(slot.family(), Some(StateFamily::AnticipatedThermalState));
    }
}