//! PreparedDataset -- wraps a [`DataSplit`] and provides batch iteration
//! for feeding into ML training loops.
//!
//! The [`BatchIterator`] yields [`Batch`] structs of a configurable size,
//! iterating over the underlying feature vectors without direct indexing.

use super::manager::{DataSplit, FEATURE_DIM};

// ---------------------------------------------------------------------------
// Batch
// ---------------------------------------------------------------------------

/// A single batch of feature vectors ready for training.
#[derive(Debug, Clone)]
pub struct Batch {
    /// The feature vectors in this batch.
    pub features: Vec<[f64; FEATURE_DIM]>,
    /// Zero-based index of this batch within the iterator.
    pub batch_index: usize,
    /// Total number of batches the iterator will produce.
    pub total_batches: usize,
}

// ---------------------------------------------------------------------------
// BatchIterator
// ---------------------------------------------------------------------------

/// Iterator that yields [`Batch`] chunks over a borrowed slice of features.
#[derive(Debug)]
pub struct BatchIterator {
    /// Owned copy of the feature data being iterated.
    data: Vec<[f64; FEATURE_DIM]>,
    /// Number of items per batch (always >= 1).
    batch_size: usize,
    /// Current batch index (0-based).
    current: usize,
    /// Pre-computed total number of batches.
    total_batches: usize,
}

impl BatchIterator {
    /// Create a new iterator over `data` with the given `batch_size`.
    ///
    /// The batch size is clamped to a minimum of 1 to avoid division-by-zero.
    fn new(data: Vec<[f64; FEATURE_DIM]>, batch_size: usize) -> Self {
        let bs = batch_size.max(1);
        let total_batches = if data.is_empty() {
            0
        } else {
            (data.len() + bs - 1) / bs
        };
        Self {
            data,
            batch_size: bs,
            current: 0,
            total_batches,
        }
    }
}

impl Iterator for BatchIterator {
    type Item = Batch;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current >= self.total_batches {
            return None;
        }

        let start = self.current * self.batch_size;
        let end = ((self.current + 1) * self.batch_size).min(self.data.len());
        let features = self.data.get(start..end)?.to_vec(); // cpu-side batch slice

        let batch = Batch {
            features,
            batch_index: self.current,
            total_batches: self.total_batches,
        };
        self.current += 1;
        Some(batch)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.total_batches.saturating_sub(self.current);
        (remaining, Some(remaining))
    }
}

impl ExactSizeIterator for BatchIterator {}

// ---------------------------------------------------------------------------
// PreparedDataset
// ---------------------------------------------------------------------------

/// A prepared dataset ready for ML training, wrapping a [`DataSplit`] and
/// providing batch iteration over each partition.
#[derive(Debug, Clone)]
pub struct PreparedDataset {
    split: DataSplit,
}

impl PreparedDataset {
    /// Wrap an existing [`DataSplit`] into a prepared dataset.
    pub fn new(split: DataSplit) -> Self {
        Self { split }
    }

    /// Borrow the underlying [`DataSplit`].
    pub fn split(&self) -> &DataSplit {
        &self.split
    }

    /// Number of feature vectors in the training partition.
    pub fn train_len(&self) -> usize {
        self.split.train.len()
    }

    /// Number of feature vectors in the validation partition.
    pub fn val_len(&self) -> usize {
        self.split.val.len()
    }

    /// Number of feature vectors in the test partition.
    pub fn test_len(&self) -> usize {
        self.split.test.len()
    }

    /// Iterate over training data in batches of `batch_size`.
    pub fn train_batches(&self, batch_size: usize) -> BatchIterator {
        BatchIterator::new(self.split.train.clone(), batch_size)
    }

    /// Iterate over validation data in batches of `batch_size`.
    pub fn val_batches(&self, batch_size: usize) -> BatchIterator {
        BatchIterator::new(self.split.val.clone(), batch_size)
    }

    /// Iterate over test data in batches of `batch_size`.
    pub fn test_batches(&self, batch_size: usize) -> BatchIterator {
        BatchIterator::new(self.split.test.clone(), batch_size)
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    /// Build a [`DataSplit`] with the given partition sizes.
    fn make_split(train_n: usize, val_n: usize, test_n: usize) -> DataSplit {
        let make_features = |n: usize, offset: usize| -> Vec<[f64; FEATURE_DIM]> {
            (0..n)
                .map(|i| {
                    let mut f = [0.0_f64; FEATURE_DIM];
                    if let Some(v) = f.get_mut(0) {
                        *v = (offset + i) as f64;
                    }
                    f
                })
                .collect()
        };

        DataSplit {
            name: "test-split".into(),
            train: make_features(train_n, 0),
            val: make_features(val_n, train_n),
            test: make_features(test_n, train_n + val_n),
            total_bars: train_n + val_n + test_n,
            symbols: vec!["SYM".into()],
        }
    }

    #[test]
    fn test_batch_iterator_correct_count() {
        let split = make_split(100, 0, 0);
        let ds = PreparedDataset::new(split);
        let batches: Vec<Batch> = ds.train_batches(32).collect();
        // 100 / 32 = 3 full + 1 partial = 4
        assert_eq!(batches.len(), 4);
    }

    #[test]
    fn test_batch_sizes() {
        let split = make_split(100, 0, 0);
        let ds = PreparedDataset::new(split);
        let batches: Vec<Batch> = ds.train_batches(32).collect();

        // First three batches should have 32 items each
        let b0 = batches.get(0);
        let b1 = batches.get(1);
        let b2 = batches.get(2);
        let b3 = batches.get(3);

        assert_eq!(b0.map(|b| b.features.len()), Some(32));
        assert_eq!(b1.map(|b| b.features.len()), Some(32));
        assert_eq!(b2.map(|b| b.features.len()), Some(32));
        // Last batch has the remaining 4 items
        assert_eq!(b3.map(|b| b.features.len()), Some(4));
    }

    #[test]
    fn test_batch_total_count_field() {
        let split = make_split(100, 0, 0);
        let ds = PreparedDataset::new(split);
        let batches: Vec<Batch> = ds.train_batches(32).collect();

        for batch in &batches {
            assert_eq!(batch.total_batches, 4);
        }
    }

    #[test]
    fn test_val_and_test_batches() {
        let split = make_split(0, 15, 15);
        let ds = PreparedDataset::new(split);

        let val_batches: Vec<Batch> = ds.val_batches(10).collect();
        let test_batches: Vec<Batch> = ds.test_batches(10).collect();

        // 15 / 10 = 1 full + 1 partial = 2
        assert_eq!(val_batches.len(), 2);
        assert_eq!(test_batches.len(), 2);

        assert_eq!(val_batches.get(0).map(|b| b.features.len()), Some(10));
        assert_eq!(val_batches.get(1).map(|b| b.features.len()), Some(5));
        assert_eq!(test_batches.get(0).map(|b| b.features.len()), Some(10));
        assert_eq!(test_batches.get(1).map(|b| b.features.len()), Some(5));
    }

    #[test]
    fn test_empty_split() {
        let split = make_split(0, 0, 0);
        let ds = PreparedDataset::new(split);

        let train_batches: Vec<Batch> = ds.train_batches(32).collect();
        let val_batches: Vec<Batch> = ds.val_batches(32).collect();
        let test_batches: Vec<Batch> = ds.test_batches(32).collect();

        assert!(train_batches.is_empty());
        assert!(val_batches.is_empty());
        assert!(test_batches.is_empty());
    }

    #[test]
    fn test_batch_size_one() {
        let split = make_split(5, 0, 0);
        let ds = PreparedDataset::new(split);
        let batches: Vec<Batch> = ds.train_batches(1).collect();

        assert_eq!(batches.len(), 5);
        for batch in &batches {
            assert_eq!(batch.features.len(), 1);
        }
    }

    #[test]
    fn test_lens() {
        let split = make_split(70, 15, 15);
        let ds = PreparedDataset::new(split);

        assert_eq!(ds.train_len(), 70);
        assert_eq!(ds.val_len(), 15);
        assert_eq!(ds.test_len(), 15);
    }
}
