use smallvec::SmallVec;

pub type Dims = SmallVec<[usize; 4]>;

/// Compute contiguous strides for a given shape (row-major / C order).
/// Example: shape [2, 3, 4] => strides [12, 4, 1]
pub fn contiguous_strides(shape: &[usize]) -> Dims {
    let mut strides = SmallVec::with_capacity(shape.len());
    strides.resize(shape.len(), 0);
    if shape.is_empty() {
        return strides;
    }
    strides[shape.len() - 1] = 1;
    for i in (0..shape.len() - 1).rev() {
        strides[i] = strides[i + 1] * shape[i + 1];
    }
    strides
}

/// Check if the given strides represent contiguous (row-major) layout for the shape.
/// A stride mismatch on a dimension of size 1 is allowed because that
/// dimension is never stepped.
pub fn is_contiguous(shape: &[usize], strides: &[usize]) -> bool {
    if shape.is_empty() {
        return true;
    }
    let ndim = shape.len();
    let mut expected_stride = 1usize;
    for d in (0..ndim).rev() {
        if shape[d] != 1 && strides[d] != expected_stride {
            return false;
        }
        expected_stride *= shape[d];
    }
    true
}

/// Total number of elements given a shape.
pub fn num_elements(shape: &[usize]) -> usize {
    shape.iter().product()
}

/// Compute the shape after broadcasting two shapes together (NumPy rules).
/// Returns None if shapes are not broadcastable.
pub fn broadcast_shape(a: &[usize], b: &[usize]) -> Option<Dims> {
    let ndim = a.len().max(b.len());
    let mut result = SmallVec::with_capacity(ndim);
    for i in 0..ndim {
        let da = if i < ndim - a.len() { 1 } else { a[i - (ndim - a.len())] };
        let db = if i < ndim - b.len() { 1 } else { b[i - (ndim - b.len())] };
        if da == db {
            result.push(da);
        } else if da == 1 {
            result.push(db);
        } else if db == 1 {
            result.push(da);
        } else {
            return None;
        }
    }
    Some(result)
}

/// Compute broadcast strides: for dimensions where size is 1 but output is >1, stride becomes 0.
pub fn broadcast_strides(shape: &[usize], strides: &[usize], target_shape: &[usize]) -> Dims {
    let ndim = target_shape.len();
    let offset = ndim - shape.len();
    let mut result = SmallVec::with_capacity(ndim);
    for i in 0..ndim {
        if i < offset {
            result.push(0);
        } else {
            let orig_idx = i - offset;
            if shape[orig_idx] == 1 && target_shape[i] > 1 {
                result.push(0);
            } else {
                result.push(strides[orig_idx]);
            }
        }
    }
    result
}

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

    #[test]
    fn test_contiguous_strides() {
        assert_eq!(contiguous_strides(&[2, 3, 4]).as_slice(), &[12, 4, 1]);
        assert_eq!(contiguous_strides(&[5]).as_slice(), &[1]);
        assert_eq!(contiguous_strides(&[2, 3]).as_slice(), &[3, 1]);
    }

    #[test]
    fn test_is_contiguous() {
        assert!(is_contiguous(&[2, 3], &[3, 1]));
        assert!(!is_contiguous(&[3, 2], &[1, 3])); // transposed
    }

    #[test]
    fn test_broadcast_shape() {
        assert_eq!(broadcast_shape(&[3, 1], &[1, 4]).unwrap().as_slice(), &[3, 4]);
        assert_eq!(broadcast_shape(&[2, 3, 4], &[4]).unwrap().as_slice(), &[2, 3, 4]);
        assert_eq!(broadcast_shape(&[1], &[5, 3]).unwrap().as_slice(), &[5, 3]);
        assert!(broadcast_shape(&[3], &[4]).is_none());
    }

    #[test]
    fn test_broadcast_strides() {
        // [3,1] with strides [1,1] broadcast to [3,4]
        assert_eq!(broadcast_strides(&[3, 1], &[1, 1], &[3, 4]).as_slice(), &[1, 0]);
    }
}