//! Utility functions for tensor operations that aren't available in rtx-tensor
//! Following strict TDD approach - implementing missing functionality
//!
//! This module includes tensor indexing utilities needed for cross-validation splits.

use crate::{AutoMLError, AutoMLResult};
use rtx_tensor::{Result, Tensor};

/// Find minimum value in a tensor
/// Since rtx-tensor doesn't have min(), we implement it using max with negation
pub fn find_min(tensor: &Tensor) -> Result<f32> {
    // Negate the tensor, find max, then negate result
    let negated = tensor.mul_scalar(-1.0)?;
    let max_of_negated = negated.max()?.to_scalar::<f32>()?;
    Ok(-max_of_negated)
}

/// Find minimum value in a tensor and return as Tensor
pub fn tensor_min(tensor: &Tensor) -> Result<Tensor> {
    let min_val = find_min(tensor)?;
    Tensor::scalar(min_val, tensor.dtype(), tensor.device())
}

/// Index a tensor by a list of sample indices (for CV train/test splits)
///
/// This function selects rows from a tensor based on the provided indices,
/// useful for creating train/validation splits from cross-validation.
///
/// # Arguments
/// * `tensor` - The input tensor with shape [n_samples, ...]
/// * `indices` - The indices of samples to select
///
/// # Returns
/// A new tensor containing only the selected samples
///
/// # Example
/// ```ignore
/// let x = Tensor::randn(&[100, 10], &device)?;
/// let train_idx = vec![0, 2, 4, 6, 8];
/// let x_train = index_tensor(&x, &train_idx)?;
/// // x_train has shape [5, 10]
/// ```
pub fn index_tensor(tensor: &Tensor, indices: &[usize]) -> AutoMLResult<Tensor> {
    if indices.is_empty() {
        return Err(AutoMLError::ValidationError(
            "Cannot index tensor with empty indices".to_string(),
        ));
    }

    let shape = tensor.shape();
    let n_samples = shape.dims()[0];

    // Validate indices
    for &idx in indices {
        if idx >= n_samples {
            return Err(AutoMLError::ValidationError(format!(
                "Index {idx} out of bounds for tensor with {n_samples} samples"
            )));
        }
    }

    // Use index_select on dimension 0 (samples dimension)
    // rtx-tensor's index_select takes &[usize] directly
    tensor
        .index_select(0, indices)
        .map_err(|e| AutoMLError::ModelError(e.to_string()))
}

/// Split tensor data into train and validation sets based on indices
///
/// # Arguments
/// * `x` - Features tensor with shape [n_samples, n_features]
/// * `y` - Labels tensor with shape [n_samples] or [n_samples, n_outputs]
/// * `train_idx` - Indices for training samples
/// * `val_idx` - Indices for validation samples
///
/// # Returns
/// Tuple of (x_train, y_train, x_val, y_val)
pub fn split_train_val(
    x: &Tensor,
    y: &Tensor,
    train_idx: &[usize],
    val_idx: &[usize],
) -> AutoMLResult<(Tensor, Tensor, Tensor, Tensor)> {
    let x_train = index_tensor(x, train_idx)?;
    let y_train = index_tensor(y, train_idx)?;
    let x_val = index_tensor(x, val_idx)?;
    let y_val = index_tensor(y, val_idx)?;

    Ok((x_train, y_train, x_val, y_val))
}

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

    #[test]
    fn test_find_min_basic() {
        let device = Device::cpu();
        let data = vec![3.0, 1.0, 4.0, 1.5, 9.0, 2.6];
        let tensor = Tensor::from_slice(&data, &[6], &device).unwrap();

        let min_val = find_min(&tensor).unwrap();
        assert_eq!(min_val, 1.0);
    }

    #[test]
    fn test_find_min_with_negatives() {
        let device = Device::cpu();
        let data = vec![-5.0, 3.0, -1.0, 4.0];
        let tensor = Tensor::from_slice(&data, &[4], &device).unwrap();

        let min_val = find_min(&tensor).unwrap();
        assert_eq!(min_val, -5.0);
    }

    #[test]
    fn test_tensor_min_returns_tensor() {
        let device = Device::cpu();
        let data = vec![3.0, 1.0, 4.0];
        let tensor = Tensor::from_slice(&data, &[3], &device).unwrap();

        let min_tensor = tensor_min(&tensor).unwrap();
        let min_val = min_tensor.to_scalar::<f32>().unwrap();
        assert_eq!(min_val, 1.0);
    }
}