use ariadnetor_linalg::DenseHostOps;
use ariadnetor_tensor::{DenseTensor, DenseTensorData, MemoryOrder};

/// Build a `DenseTensor` from row-major data, reordered to column-major.
fn cm<T: Clone>(data: Vec<T>, shape: Vec<usize>) -> DenseTensor<T> {
    let rm = DenseTensorData::from_raw_parts(data, shape, MemoryOrder::RowMajor);
    let cm = ariadnetor_tensor::reorder_data(&rm, MemoryOrder::ColumnMajor);
    DenseTensor::from_data(cm)
}

/// Reorder a `DenseTensor` back to row-major for index-by-index assertions.
fn to_rm<T: Clone>(tensor: &DenseTensor<T>) -> DenseTensor<T> {
    let rm = ariadnetor_tensor::reorder_data(tensor.data(), MemoryOrder::RowMajor);
    DenseTensor::from_data(rm)
}

#[test]
fn test_permute_f64_2d() {
    let tensor = cm(vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]);

    let result = to_rm(&tensor.permute(&[1, 0]).unwrap());

    assert_eq!(result.shape(), &[3, 2]);
    // Transposed: [[1,4],[2,5],[3,6]]
    assert_eq!(result.get([0, 0]), 1.0);
    assert_eq!(result.get([0, 1]), 4.0);
    assert_eq!(result.get([1, 0]), 2.0);
    assert_eq!(result.get([1, 1]), 5.0);
    assert_eq!(result.get([2, 0]), 3.0);
    assert_eq!(result.get([2, 1]), 6.0);
}

#[test]
fn test_permute_f64_3d() {
    let data: Vec<f64> = (0..24).map(|i| i as f64).collect();
    let tensor = cm(data, vec![2, 3, 4]);

    let result = to_rm(&tensor.permute(&[2, 0, 1]).unwrap());
    let tensor_rm = to_rm(&tensor);

    assert_eq!(result.shape(), &[4, 2, 3]);
    assert_eq!(result.len(), 24);
    // input[0][0][0] = 0 -> output[0][0][0]
    assert_eq!(result.get([0, 0, 0]), tensor_rm.get([0, 0, 0]));
    // input[0][0][1] = 1 -> output[1][0][0]
    assert_eq!(result.get([1, 0, 0]), tensor_rm.get([0, 0, 1]));
}

#[test]
fn test_permute_f32_2d() {
    let tensor = cm(vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]);

    let result = to_rm(&tensor.permute(&[1, 0]).unwrap());

    assert_eq!(result.shape(), &[3, 2]);
    assert_eq!(result.get([0, 0]), 1.0f32);
    assert_eq!(result.get([0, 1]), 4.0f32);
    assert_eq!(result.get([1, 0]), 2.0f32);
}

#[test]
fn test_permute_complex_f64_2d() {
    use num_complex::Complex;

    let input = vec![
        Complex::new(1.0, 2.0),
        Complex::new(3.0, 4.0),
        Complex::new(5.0, 6.0),
        Complex::new(7.0, 8.0),
        Complex::new(9.0, 10.0),
        Complex::new(11.0, 12.0),
    ];
    let tensor = cm(input, vec![2, 3]);

    let result = to_rm(&tensor.permute(&[1, 0]).unwrap());

    assert_eq!(result.shape(), &[3, 2]);
    assert_eq!(result.get([0, 0]), Complex::new(1.0, 2.0));
    assert_eq!(result.get([0, 1]), Complex::new(7.0, 8.0));
    assert_eq!(result.get([1, 0]), Complex::new(3.0, 4.0));
    assert_eq!(result.get([1, 1]), Complex::new(9.0, 10.0));
}

#[test]
fn test_permute_empty_tensor() {
    let tensor = cm(Vec::<f64>::new(), vec![0, 3]);

    let result = tensor.permute(&[1, 0]).unwrap();

    assert_eq!(result.shape(), &[3, 0]);
    assert_eq!(result.len(), 0);
}
