use super::*;
use burn_tensor::TensorData;

#[test]
fn test_equal_inf() {
    let data_1 = TensorData::from([[0.0, 1.0, 2.0], [f32::INFINITY, 4.0, f32::NEG_INFINITY]]);
    let data_2 = TensorData::from([[1.0, 1.0, 1.0], [f32::INFINITY, 3.0, f32::NEG_INFINITY]]);
    let device = Default::default();
    let tensor_1 = TestTensor::<2>::from_data(data_1, &device);
    let tensor_2 = TestTensor::<2>::from_data(data_2, &device);

    let data_actual_cloned = tensor_1.clone().equal(tensor_2.clone());
    let data_actual_inplace = tensor_1.equal(tensor_2);

    let data_expected = TensorData::from([[false, true, false], [true, false, true]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_not_equal_inf() {
    let data_1 = TensorData::from([[0.0, 1.0, 2.0], [3.0, f32::INFINITY, 5.0]]);
    let data_2 = TensorData::from([[1.0, 1.0, 1.0], [f32::INFINITY, 3.0, f32::NEG_INFINITY]]);
    let device = Default::default();
    let tensor_1 = TestTensor::<2>::from_data(data_1, &device);
    let tensor_2 = TestTensor::<2>::from_data(data_2, &device);

    let data_actual_cloned = tensor_1.clone().not_equal(tensor_2.clone());
    let data_actual_inplace = tensor_1.not_equal(tensor_2);

    let data_expected = TensorData::from([[true, false, true], [true, true, true]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_equal() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);
    let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 5.0]]);

    let data_actual_cloned = tensor_1.clone().equal(tensor_2.clone());
    let data_actual_inplace = tensor_1.equal(tensor_2);

    let data_expected = TensorData::from([[false, true, false], [false, false, true]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_not_equal() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);
    let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 5.0]]);

    let data_actual_cloned = tensor_1.clone().not_equal(tensor_2.clone());
    let data_actual_inplace = tensor_1.not_equal(tensor_2);

    let data_expected = TensorData::from([[true, false, true], [true, true, false]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_equal_elem() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 2.0, 5.0]]);

    let data_actual_cloned = tensor_1.clone().equal_elem(2);
    let data_actual_inplace = tensor_1.equal_elem(2);

    let data_expected = TensorData::from([[false, false, true], [false, true, false]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_not_equal_elem() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 2.0, 5.0]]);

    let data_actual_cloned = tensor_1.clone().not_equal_elem(2);
    let data_actual_inplace = tensor_1.not_equal_elem(2);

    let data_expected = TensorData::from([[true, true, false], [true, false, true]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn greater_elem() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);

    let data_actual_cloned = tensor_1.clone().greater_elem(4);
    let data_actual_inplace = tensor_1.greater_elem(4);

    let data_expected = TensorData::from([[false, false, false], [false, false, true]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_greater_equal_elem() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);

    let data_actual_cloned = tensor_1.clone().greater_equal_elem(4.0);
    let data_actual_inplace = tensor_1.greater_equal_elem(4.0);

    let data_expected = TensorData::from([[false, false, false], [false, true, true]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_greater() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);
    let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 50.0]]);

    let data_actual_cloned = tensor_1.clone().greater(tensor_2.clone());
    let data_actual_inplace = tensor_1.greater(tensor_2);

    let data_expected = TensorData::from([[false, false, true], [false, true, false]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_greater_equal() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);
    let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 50.0]]);

    let data_actual_cloned = tensor_1.clone().greater_equal(tensor_2.clone());
    let data_actual_inplace = tensor_1.greater_equal(tensor_2);

    let data_expected = TensorData::from([[false, true, true], [false, true, false]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_lower_elem() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);

    let data_actual_cloned = tensor_1.clone().lower_elem(4.0);
    let data_actual_inplace = tensor_1.lower_elem(4.0);

    let data_expected = TensorData::from([[true, true, true], [true, false, false]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_lower_equal_elem() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);

    let data_actual_cloned = tensor_1.clone().lower_equal_elem(4.0);
    let data_actual_inplace = tensor_1.lower_equal_elem(4.0);

    let data_expected = TensorData::from([[true, true, true], [true, true, false]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_lower() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);
    let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 50.0]]);

    let data_actual_cloned = tensor_1.clone().lower(tensor_2.clone());
    let data_actual_inplace = tensor_1.lower(tensor_2);

    let data_expected = TensorData::from([[true, false, false], [true, false, true]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_lower_equal() {
    let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);
    let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 50.0]]);

    let data_actual_cloned = tensor_1.clone().lower_equal(tensor_2.clone());
    let data_actual_inplace = tensor_1.lower_equal(tensor_2);

    let data_expected = TensorData::from([[true, true, false], [true, false, true]]);
    data_expected.assert_eq(&data_actual_cloned.into_data(), false);
    data_expected.assert_eq(&data_actual_inplace.into_data(), false);
}

#[test]
fn test_greater_broadcast() {
    // Test broadcasting with shape [1, 4] vs [4, 4]
    let device = Default::default();
    let data_1 = TensorData::from([[1.0, 2.0, 3.0, 4.0]]);
    let data_2 = TensorData::from([
        [0.5, 1.5, 2.5, 3.5],
        [1.5, 2.5, 3.5, 4.5],
        [2.5, 3.5, 4.5, 5.5],
        [3.5, 4.5, 5.5, 6.5],
    ]);
    let tensor_1 = TestTensor::<2>::from_data(data_1, &device);
    let tensor_2 = TestTensor::<2>::from_data(data_2, &device);

    let result = tensor_1.greater(tensor_2);

    let expected = TensorData::from([
        [true, true, true, true],
        [false, false, false, false],
        [false, false, false, false],
        [false, false, false, false],
    ]);
    expected.assert_eq(&result.into_data(), false);
}

#[test]
fn test_greater_equal_broadcast() {
    // Test broadcasting with shape [4, 1] vs [1, 4]
    let device = Default::default();
    let data_1 = TensorData::from([[1.0], [2.0], [3.0], [4.0]]);
    let data_2 = TensorData::from([[1.0, 2.0, 3.0, 4.0]]);
    let tensor_1 = TestTensor::<2>::from_data(data_1, &device);
    let tensor_2 = TestTensor::<2>::from_data(data_2, &device);

    let result = tensor_1.greater_equal(tensor_2);

    let expected = TensorData::from([
        [true, false, false, false],
        [true, true, false, false],
        [true, true, true, false],
        [true, true, true, true],
    ]);
    expected.assert_eq(&result.into_data(), false);
}

#[test]
fn test_lower_broadcast() {
    // Test broadcasting mimicking CLIP pattern: [1, 5] vs [5, 1]
    let device = Default::default();
    let data_1 = TensorData::from([[0.0, 1.0, -1.0, 2.0, -2.0]]);
    let data_2 = TensorData::from([[0.5], [1.5], [-0.5], [-1.5], [2.5]]);
    let tensor_1 = TestTensor::<2>::from_data(data_1, &device);
    let tensor_2 = TestTensor::<2>::from_data(data_2, &device);

    let result = tensor_1.lower(tensor_2);

    let expected = TensorData::from([
        [true, false, true, false, true],
        [true, true, true, false, true],
        [false, false, true, false, true],
        [false, false, false, false, true],
        [true, true, true, true, true],
    ]);
    expected.assert_eq(&result.into_data(), false);
}

#[test]
fn test_lower_equal_broadcast() {
    // Test broadcasting with shape [1, 1] vs [2, 4]
    let device = Default::default();
    let data_1 = TensorData::from([[2.5]]);
    let data_2 = TensorData::from([[1.0, 2.0, 3.0, 4.0], [2.0, 2.5, 3.0, 3.5]]);
    let tensor_1 = TestTensor::<2>::from_data(data_1, &device);
    let tensor_2 = TestTensor::<2>::from_data(data_2, &device);

    let result = tensor_1.lower_equal(tensor_2);

    let expected = TensorData::from([[false, false, true, true], [false, true, true, true]]);
    expected.assert_eq(&result.into_data(), false);
}

#[test]
fn test_equal_broadcast() {
    // Test broadcasting with different ranks
    let device = Default::default();
    let data_1 = TensorData::from([[2.0], [3.0], [4.0]]);
    let data_2 = TensorData::from([[2.0, 3.0, 4.0, 2.0]]);
    let tensor_1 = TestTensor::<2>::from_data(data_1, &device);
    let tensor_2 = TestTensor::<2>::from_data(data_2, &device);

    let result = tensor_1.equal(tensor_2);

    let expected = TensorData::from([
        [true, false, false, true],
        [false, true, false, false],
        [false, false, true, false],
    ]);
    expected.assert_eq(&result.into_data(), false);
}

#[test]
fn test_not_equal_broadcast() {
    // Test broadcasting with shape [3, 1] vs [1, 3]
    let device = Default::default();
    let data_1 = TensorData::from([[1.0], [2.0], [3.0]]);
    let data_2 = TensorData::from([[1.0, 2.0, 3.0]]);
    let tensor_1 = TestTensor::<2>::from_data(data_1, &device);
    let tensor_2 = TestTensor::<2>::from_data(data_2, &device);

    let result = tensor_1.not_equal(tensor_2);

    let expected = TensorData::from([
        [false, true, true],
        [true, false, true],
        [true, true, false],
    ]);
    expected.assert_eq(&result.into_data(), false);
}

#[test]
fn test_greater_transposed() {
    // [[1,2],[3,4]] transposed -> [[1,3],[2,4]]; > [[2,2],[2,2]] = [[F,T],[F,T]]
    let lhs = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).transpose();
    let rhs = TestTensor::<2>::from([[2.0, 2.0], [2.0, 2.0]]);

    let result = lhs.greater(rhs);

    result
        .into_data()
        .assert_eq(&TensorData::from([[false, true], [false, true]]), false);
}

#[test]
fn test_equal_flipped_1d() {
    // [1,2,3,4] flipped -> [4,3,2,1]; == [4,2,2,1] = [T,F,T,T]
    let lhs = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]).flip([0]);
    let rhs = TestTensor::<1>::from([4.0, 2.0, 2.0, 1.0]);

    let result = lhs.equal(rhs);

    result
        .into_data()
        .assert_eq(&TensorData::from([true, false, true, true]), false);
}

#[test]
fn test_lower_flipped_2d() {
    // [[1,2],[3,4]] axis-0 flipped -> [[3,4],[1,2]]; < [[2,5],[2,1]] = [[F,T],[T,F]]
    let lhs = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).flip([0]);
    let rhs = TestTensor::<2>::from([[2.0, 5.0], [2.0, 1.0]]);

    let result = lhs.lower(rhs);

    result
        .into_data()
        .assert_eq(&TensorData::from([[false, true], [true, false]]), false);
}

#[test]
fn test_greater_elem_flipped() {
    // [1,2,3,4] flipped -> [4,3,2,1]; > 2.5 = [T,T,F,F]
    let lhs = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]).flip([0]);

    let result = lhs.greater_elem(2.5);

    result
        .into_data()
        .assert_eq(&TensorData::from([true, true, false, false]), false);
}

#[test]
fn test_equal_both_transposed() {
    // Both transposed. [[1,2],[3,4]]^T == [[1,3],[2,4]]^T = ?
    let lhs = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).transpose();
    let rhs = TestTensor::<2>::from([[1.0, 3.0], [2.0, 4.0]]).transpose();

    let result = lhs.equal(rhs);

    result
        .into_data()
        .assert_eq(&TensorData::from([[true, false], [false, true]]), false);
}

#[test]
fn test_not_equal_narrowed() {
    // [1,2,3,4,5,6] narrowed to [2,3,4,5]; != [2,2,4,4] = [F,T,F,T]
    let lhs = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).narrow(0, 1, 4);
    let rhs = TestTensor::<1>::from([2.0, 2.0, 4.0, 4.0]);

    let result = lhs.not_equal(rhs);

    result
        .into_data()
        .assert_eq(&TensorData::from([false, true, false, true]), false);
}

#[test]
fn test_lower_flipped_both_axes() {
    // [[1,2],[3,4]] flipped on both axes -> [[4,3],[2,1]]; < [[3,3],[3,3]] = [[F,F],[T,T]]
    let lhs = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).flip([0, 1]);
    let rhs = TestTensor::<2>::from([[3.0, 3.0], [3.0, 3.0]]);

    let result = lhs.lower(rhs);

    result
        .into_data()
        .assert_eq(&TensorData::from([[false, false], [true, true]]), false);
}
