tenshift-core 0.1.3

use crate::error::{Error, Result};
use crate::pipeline::{CollectorItem, CollectorStage};
use crate::sample::Sample;
use crate::transform::StatefulTransform;

/// Validate that `sample` has the same schema (field names, dtypes, and shapes)
/// as `reference`, comparing fields directly with zero allocation.
///
/// Field order is irrelevant: each of `sample`'s fields is looked up by name in
/// `reference`. Because field names are unique within a sample, equal field
/// counts plus every field of `sample` matching a `reference` field is a
/// bijection, so the schemas are equal.
///
/// The former implementation built a `SchemaSignature` (a fresh `Vec` of cloned
/// `String` names and `Vec<usize>` shapes, then sorted) for the reference AND
/// for every sample in every batch just to compare and discard - pure allocation
/// churn in the hot collector loop. This path allocates only when constructing
/// an error message.
fn schema_matches(reference: &Sample, sample: &Sample) -> Result<()> {
    if reference.len() != sample.len() {
        return Err(Error::CollateFailed {
            reason: format!(
                "schema mismatch: expected {} fields, got {}",
                reference.len(),
                sample.len()
            ),
        });
    }
    for (name, tensor) in sample.iter() {
        let Some(expected) = reference.get(name) else {
            return Err(Error::CollateFailed {
                reason: format!("schema mismatch: unexpected field '{name}'"),
            });
        };
        if expected.dtype() != tensor.dtype() {
            return Err(Error::CollateFailed {
                reason: format!(
                    "schema mismatch for field '{name}': expected dtype {:?}, got {:?}",
                    expected.dtype(),
                    tensor.dtype()
                ),
            });
        }
        if expected.shape() != tensor.shape() {
            return Err(Error::CollateFailed {
                reason: format!(
                    "schema mismatch for field '{name}': expected shape {:?}, got {:?}",
                    expected.shape(),
                    tensor.shape()
                ),
            });
        }
    }
    Ok(())
}

/// Validate that all samples in a batch have consistent schemas.
/// Returns an error if any sample deviates from the first sample's schema.
pub(crate) fn validate_batch_schema(batch: &[Sample]) -> Result<()> {
    if batch.len() < 2 {
        return Ok(());
    }
    let reference = &batch[0];
    for (idx, sample) in batch[1..].iter().enumerate() {
        if let Err(e) = schema_matches(reference, sample) {
            return Err(Error::CollateFailed {
                reason: format!("sample {} in batch: {}", idx + 1, e),
            });
        }
    }
    Ok(())
}

pub(crate) fn push_collector_stage(
    stage: &mut CollectorStage,
    item: CollectorItem,
    out: &mut Vec<CollectorItem>,
) -> Result<()> {
    match (stage, item) {
        (CollectorStage::Shuffle(shuffle), CollectorItem::Sample(sample)) => {
            for v in shuffle.push(sample) {
                out.push(CollectorItem::Sample(v));
            }
            Ok(())
        }
        (CollectorStage::Batch(batch), CollectorItem::Sample(sample)) => {
            let buffered = batch.push(sample);
            if !buffered.is_empty() {
                out.push(CollectorItem::Batch(buffered));
            }
            Ok(())
        }
        (CollectorStage::Shuffle(_), CollectorItem::Batch(_)) => Err(Error::InvalidConfig {
            reason: "shuffle cannot run after a batch stage".to_string(),
        }),
        (CollectorStage::Batch(_), CollectorItem::Batch(_)) => Err(Error::InvalidConfig {
            reason: "multiple batch stages are not supported".to_string(),
        }),
    }
}

pub(crate) fn finish_collector_stage(stage: &mut CollectorStage) -> Vec<CollectorItem> {
    match stage {
        CollectorStage::Shuffle(shuffle) => shuffle
            .finish()
            .into_iter()
            .map(CollectorItem::Sample)
            .collect(),
        CollectorStage::Batch(batch) => {
            let buffered = batch.finish();
            if buffered.is_empty() {
                Vec::new()
            } else {
                vec![CollectorItem::Batch(buffered)]
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::validate_batch_schema;
    use crate::error::Error;
    use crate::sample::{Sample, Tensor};

    fn sample_ab() -> Sample {
        Sample::new()
            .with("a", Tensor::f32(&[1.0, 2.0], vec![2]))
            .with("b", Tensor::i64(&[7], vec![1]))
    }

    #[test]
    fn validate_batch_schema_accepts_consistent_batch_regardless_of_field_order() {
        // Field insertion order differs between samples; schema comparison must be
        // order-independent (the old code sorted by name — the allocation-free
        // path looks up by name, so order still must not matter).
        let s1 = sample_ab();
        let s2 = Sample::new()
            .with("b", Tensor::i64(&[9], vec![1]))
            .with("a", Tensor::f32(&[3.0, 4.0], vec![2]));
        assert!(validate_batch_schema(&[s1, s2]).is_ok());
    }

    #[test]
    fn validate_batch_schema_rejects_field_count_mismatch() {
        let s1 = sample_ab();
        let s2 = Sample::new().with("a", Tensor::f32(&[1.0, 2.0], vec![2]));
        let err = validate_batch_schema(&[s1, s2]).unwrap_err();
        match err {
            Error::CollateFailed { reason } => {
                assert!(reason.contains("expected 2 fields, got 1"), "got: {reason}");
                assert!(reason.contains("sample 1 in batch"), "got: {reason}");
            }
            other => panic!("expected CollateFailed, got {other:?}"),
        }
    }

    #[test]
    fn validate_batch_schema_rejects_dtype_mismatch() {
        let s1 = sample_ab();
        // Field "b" is i64 in s1 but i32 here.
        let s2 = Sample::new()
            .with("a", Tensor::f32(&[1.0, 2.0], vec![2]))
            .with("b", Tensor::i32(&[7], vec![1]));
        let err = validate_batch_schema(&[s1, s2]).unwrap_err();
        match err {
            Error::CollateFailed { reason } => {
                assert!(reason.contains("field 'b'"), "got: {reason}");
                assert!(reason.contains("dtype"), "got: {reason}");
            }
            other => panic!("expected CollateFailed, got {other:?}"),
        }
    }

    #[test]
    fn validate_batch_schema_rejects_shape_mismatch() {
        let s1 = sample_ab();
        // Field "a" has shape [2] in s1 but [3] here.
        let s2 = Sample::new()
            .with("a", Tensor::f32(&[1.0, 2.0, 3.0], vec![3]))
            .with("b", Tensor::i64(&[7], vec![1]));
        let err = validate_batch_schema(&[s1, s2]).unwrap_err();
        match err {
            Error::CollateFailed { reason } => {
                assert!(reason.contains("field 'a'"), "got: {reason}");
                assert!(reason.contains("shape"), "got: {reason}");
            }
            other => panic!("expected CollateFailed, got {other:?}"),
        }
    }

    #[test]
    fn validate_batch_schema_rejects_renamed_field() {
        let s1 = sample_ab();
        // Same count/dtypes/shapes but "b" renamed to "c": must be caught by the
        // name lookup, not silently accepted.
        let s2 = Sample::new()
            .with("a", Tensor::f32(&[1.0, 2.0], vec![2]))
            .with("c", Tensor::i64(&[7], vec![1]));
        let err = validate_batch_schema(&[s1, s2]).unwrap_err();
        match err {
            Error::CollateFailed { reason } => {
                assert!(reason.contains("unexpected field 'c'"), "got: {reason}");
            }
            other => panic!("expected CollateFailed, got {other:?}"),
        }
    }
}