tenferro-einsum 0.9.0

use num_complex::Complex64;
use tenferro_cpu::CpuEngine;
use tenferro_tensor::{
    ContractionScalar, DType, DotGeneralAccumulation, Tensor, TensorRead, TensorView,
    TensorViewMut, TensorWrite, TypedTensor, TypedTensorView, TypedTensorViewMut, TypedTensorWrite,
};

use crate::{
    parse_einsum_subscripts, ConcreteEinsumPlan, EinsumAxis, EinsumNotation, Error,
    TensorEinsumExt, TensorEinsumIntoExt, TensorReadEinsumExt, TensorReadEinsumIntoExt,
    TensorTensordotExt, TypedTensorEinsumExt, TypedTensorEinsumIntoExt, TypedTensorReadEinsumExt,
    TypedTensorReadEinsumIntoExt, TypedTensorTensordotExt,
};

fn assert_f64_tensor(tensor: &Tensor, shape: &[usize], expected: &[f64]) {
    assert_eq!(tensor.shape(), shape);
    assert_eq!(tensor.as_slice::<f64>().unwrap(), expected);
}

#[test]
fn public_tensor_einsum_ext_executes_dtype_erased_inputs() {
    let mut backend = CpuEngine::serial();
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();

    let result = {
        let session = &mut backend;
        [&lhs, &rhs].einsum("ij,jk->ik", session)
    }
    .unwrap();

    assert_f64_tensor(&result, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);
}

#[test]
fn binary_dot_shape_check_uses_original_operand_order_for_unequal_ranks() {
    let mut backend = CpuEngine::serial();
    let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
    let rhs = Tensor::from_vec_col_major(vec![4, 2, 3], vec![1.0_f64; 24]).unwrap();

    let output = {
        let session = &mut backend;
        [&lhs, &rhs].einsum("ij,pqj->pqi", session)
    }
    .unwrap();
    assert_f64_tensor(&output, &[4, 2, 2], &[3.0; 16]);
}

#[test]
fn singleton_contracting_extent_uses_general_broadcast_fallback() {
    let mut backend = CpuEngine::serial();
    let lhs = Tensor::from_vec_col_major(vec![2, 1], vec![2.0_f64; 2]).unwrap();
    let rhs = Tensor::from_vec_col_major(vec![4, 2, 3], vec![1.0_f64; 24]).unwrap();

    let output = {
        let session = &mut backend;
        [&lhs, &rhs].einsum("ij,pqj->pqi", session)
    }
    .unwrap();
    assert_f64_tensor(&output, &[4, 2, 2], &[6.0; 16]);
}

#[test]
fn binary_read_into_rejects_extra_inputs_without_mutating_output() {
    let mut backend = CpuEngine::serial();
    let lhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
    let rhs = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
    let extra = Tensor::from_vec_col_major(vec![2], vec![7.0_f64; 2]).unwrap();
    let subscripts = crate::EinsumSubscripts::new(&[&[0, 1], &[2, 1]], &[2, 0]);

    let mut owned_out = Tensor::from_vec_col_major(vec![2, 2], vec![9.0_f64; 4]).unwrap();
    let owned = {
        let session = &mut backend;
        {
            [&lhs, &rhs, &extra].einsum_into_subscripts(
                &subscripts,
                session,
                TensorWrite::from_tensor(&mut owned_out),
            )
        }
    };
    assert!(owned.is_err());
    assert_f64_tensor(&owned_out, &[2, 2], &[9.0; 4]);

    let reads = [
        TensorRead::from_tensor(&lhs),
        TensorRead::from_tensor(&rhs),
        TensorRead::from_tensor(&extra),
    ];
    let mut read_out = Tensor::from_vec_col_major(vec![2, 2], vec![8.0_f64; 4]).unwrap();
    let read = {
        let session = &mut backend;
        {
            reads.einsum_read_into_subscripts(
                &subscripts,
                session,
                TensorWrite::from_tensor(&mut read_out),
            )
        }
    };
    assert!(read.is_err());
    assert_f64_tensor(&read_out, &[2, 2], &[8.0; 4]);

    let typed_lhs = TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0; 6]).unwrap();
    let typed_rhs = TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0; 6]).unwrap();
    let typed_extra = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![7.0; 2]).unwrap();
    let typed_views = [
        typed_lhs.as_view(),
        typed_rhs.as_view(),
        typed_extra.as_view(),
    ];
    let mut typed_out = TypedTensor::<f64>::from_vec_col_major(vec![2, 2], vec![6.0; 4]).unwrap();
    let typed = {
        let session = &mut backend;
        {
            typed_views.einsum_read_into_subscripts(&subscripts, session, &mut typed_out)
        }
    };
    assert!(typed.is_err());
    assert_eq!(typed_out.as_slice().unwrap(), &[6.0; 4]);
}

#[test]
fn concrete_einsum_ellipsis_supports_diagonal_and_zero_rank_cases() {
    let mut backend = CpuEngine::serial();
    let input =
        Tensor::from_vec_col_major(vec![2, 3, 3], (0..18).map(f64::from).collect::<Vec<_>>())
            .unwrap();
    let vector = Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap();

    let (diagonal, zero_rank) = (|| {
        let session = &mut backend;
        {
            let diagonal = [&input].einsum("...ii->...i", session)?;
            let zero_rank = [&vector].einsum("...i->...i", session)?;
            Ok::<_, crate::Error>((diagonal, zero_rank))
        }
    })()
    .unwrap();

    assert_f64_tensor(&diagonal, &[2, 3], &[0.0, 1.0, 8.0, 9.0, 16.0, 17.0]);
    assert_f64_tensor(&zero_rank, &[3], &[1.0, 2.0, 3.0]);
}

#[test]
fn concrete_einsum_ellipsis_supports_broadcast_and_programmatic_notation() {
    let mut backend = CpuEngine::serial();
    let lhs = Tensor::from_vec_col_major(vec![2, 2, 3], vec![1.0_f64; 12]).unwrap();
    let rhs = Tensor::from_vec_col_major(vec![1, 3, 2], vec![1.0_f64; 6]).unwrap();
    let notation = crate::EinsumNotation::new(
        &[
            &[
                crate::EinsumAxis::Ellipsis,
                crate::EinsumAxis::Label(0),
                crate::EinsumAxis::Label(1),
            ],
            &[
                crate::EinsumAxis::Ellipsis,
                crate::EinsumAxis::Label(1),
                crate::EinsumAxis::Label(2),
            ],
        ],
        &[
            crate::EinsumAxis::Ellipsis,
            crate::EinsumAxis::Label(0),
            crate::EinsumAxis::Label(2),
        ],
    );

    let (string_result, programmatic_result) = (|| {
        let session = &mut backend;
        {
            let string_result = [&lhs, &rhs].einsum("...ij,...jk->...ik", session)?;
            let programmatic_result = [&lhs, &rhs].einsum_notation(&notation, session)?;
            Ok::<_, crate::Error>((string_result, programmatic_result))
        }
    })()
    .unwrap();

    assert_f64_tensor(&string_result, &[2, 2, 2], &[3.0; 8]);
    assert_f64_tensor(&programmatic_result, &[2, 2, 2], &[3.0; 8]);
}

#[test]
fn concrete_programmatic_notation_covers_read_into_and_prepared_surfaces() {
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let notation = EinsumNotation::new(
        &[
            &[
                EinsumAxis::Ellipsis,
                EinsumAxis::Label(0),
                EinsumAxis::Label(1),
            ],
            &[
                EinsumAxis::Ellipsis,
                EinsumAxis::Label(1),
                EinsumAxis::Label(2),
            ],
        ],
        &[
            EinsumAxis::Ellipsis,
            EinsumAxis::Label(0),
            EinsumAxis::Label(2),
        ],
    );
    let mut backend = CpuEngine::serial();

    let mut erased_out = Tensor::from_vec_col_major(vec![2, 2], vec![0.0_f64; 4]).unwrap();
    {
        let session = &mut backend;
        {
            [&lhs, &rhs].einsum_into_notation(
                &notation,
                session,
                TensorWrite::from_tensor(&mut erased_out),
            )
        }
    }
    .unwrap();
    assert_f64_tensor(&erased_out, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);

    let typed_lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let typed_rhs =
        TypedTensor::<f64>::from_vec_col_major(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let typed = {
        let session = &mut backend;
        {
            [&typed_lhs, &typed_rhs].einsum_notation(&notation, session)
        }
    }
    .unwrap();
    assert_eq!(typed.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);

    let typed_views = [typed_lhs.as_view(), typed_rhs.as_view()];
    let typed_read = {
        let session = &mut backend;
        typed_views.einsum_read_notation(&notation, session)
    }
    .unwrap();
    assert_eq!(typed_read.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);

    let mut typed_out = TypedTensor::<f64>::from_vec_col_major(vec![2, 2], vec![0.0; 4]).unwrap();
    {
        let session = &mut backend;
        {
            [&typed_lhs, &typed_rhs].einsum_into_notation(&notation, session, &mut typed_out)
        }
    }
    .unwrap();
    assert_eq!(typed_out.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);

    let mut typed_read_out =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 2], vec![0.0; 4]).unwrap();
    {
        let session = &mut backend;
        {
            typed_views.einsum_read_into_notation(&notation, session, &mut typed_read_out)
        }
    }
    .unwrap();
    assert_eq!(
        typed_read_out.as_slice().unwrap(),
        &[22.0, 28.0, 49.0, 64.0]
    );

    let reads = [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)];
    let read_result = {
        let session = &mut backend;
        reads.einsum_read_notation(&notation, session)
    }
    .unwrap();
    assert_f64_tensor(&read_result, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);
    let mut read_out = Tensor::from_vec_col_major(vec![2, 2], vec![0.0_f64; 4]).unwrap();
    {
        let session = &mut backend;
        {
            reads.einsum_read_into_notation(
                &notation,
                session,
                TensorWrite::from_tensor(&mut read_out),
            )
        }
    }
    .unwrap();
    assert_f64_tensor(&read_out, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);

    ConcreteEinsumPlan::prepare_notation([&lhs, &rhs], &notation).unwrap();
    ConcreteEinsumPlan::prepare_typed_notation::<f64, _>([&typed_lhs, &typed_rhs], &notation)
        .unwrap();
    ConcreteEinsumPlan::prepare_read_notation(reads, &notation).unwrap();
}

#[test]
fn concrete_tensordot_matches_einsum_for_erased_and_typed_tensors() {
    let mut backend = CpuEngine::serial();
    let lhs = Tensor::from_vec_col_major([2], vec![1.0_f64, 2.0]).unwrap();
    let rhs = Tensor::from_vec_col_major([2], vec![3.0_f64, 4.0]).unwrap();
    let erased = {
        let session = &mut backend;
        {
            lhs.tensordot(&rhs, crate::TensorDotAxes::Count(1), session)
        }
    }
    .unwrap();
    assert_eq!(erased.as_slice::<f64>().unwrap(), &[11.0]);

    let lhs = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![1.0, 2.0]).unwrap();
    let rhs = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![3.0, 4.0]).unwrap();
    let typed = {
        let session = &mut backend;
        {
            lhs.tensordot(&rhs, crate::TensorDotAxes::Count(1), session)
        }
    }
    .unwrap();
    assert_eq!(typed.as_slice().unwrap(), &[11.0]);
}

#[test]
fn public_tensor_einsum_ext_accepts_slice_and_integer_subscripts() {
    let mut backend = CpuEngine::serial();
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let subscripts = parse_einsum_subscripts("ij,jk->ik").unwrap();
    let inputs = vec![&lhs, &rhs];

    let slice_result = {
        let session = &mut backend;
        inputs.as_slice().einsum_subscripts(&subscripts, session)
    }
    .unwrap();
    let array_result = {
        let session = &mut backend;
        [&lhs, &rhs].einsum_subscripts(&subscripts, session)
    }
    .unwrap();

    assert_f64_tensor(&slice_result, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);
    assert_f64_tensor(&array_result, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);
}

#[test]
fn public_typed_tensor_einsum_ext_preserves_complex_dtype() {
    let mut backend = CpuEngine::serial();
    let lhs = TypedTensor::<Complex64>::from_vec_col_major(
        vec![2, 2],
        vec![
            Complex64::new(1.0, 1.0),
            Complex64::new(2.0, -1.0),
            Complex64::new(3.0, 0.0),
            Complex64::new(4.0, 2.0),
        ],
    )
    .unwrap();
    let rhs = TypedTensor::<Complex64>::from_vec_col_major(
        vec![2, 1],
        vec![Complex64::new(5.0, 0.0), Complex64::new(6.0, -1.0)],
    )
    .unwrap();

    let result = {
        let session = &mut backend;
        [&lhs, &rhs].einsum("ij,jk->ik", session)
    }
    .unwrap();
    let subscripts = parse_einsum_subscripts("ij,jk->ik").unwrap();
    let integer_result = {
        let session = &mut backend;
        [&lhs, &rhs].einsum_subscripts(&subscripts, session)
    }
    .unwrap();
    let plan = ConcreteEinsumPlan::prepare_typed_subscripts([&lhs, &rhs], &subscripts).unwrap();
    let planned_result = {
        let session = &mut backend;
        plan.execute_typed([&lhs, &rhs], session)
    }
    .unwrap();

    assert_eq!(result.shape(), &[2, 1]);
    assert_eq!(
        result.as_slice().unwrap(),
        &[Complex64::new(23.0, 2.0), Complex64::new(36.0, 3.0)]
    );
    assert_eq!(
        integer_result.as_slice().unwrap(),
        &[Complex64::new(23.0, 2.0), Complex64::new(36.0, 3.0)]
    );
    assert_eq!(
        planned_result.as_slice().unwrap(),
        &[Complex64::new(23.0, 2.0), Complex64::new(36.0, 3.0)]
    );
}

#[test]
fn public_typed_tensor_einsum_ext_accepts_slice_and_integer_subscripts() {
    let mut backend = CpuEngine::serial();
    let lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let rhs =
        TypedTensor::<f64>::from_vec_col_major(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let subscripts = parse_einsum_subscripts("ij,jk->ik").unwrap();
    let inputs = vec![&lhs, &rhs];

    let slice_result = {
        let session = &mut backend;
        inputs.as_slice().einsum_subscripts(&subscripts, session)
    }
    .unwrap();
    let array_result = {
        let session = &mut backend;
        [&lhs, &rhs].einsum_subscripts(&subscripts, session)
    }
    .unwrap();

    assert_eq!(slice_result.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);
    assert_eq!(array_result.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);
}

#[test]
fn public_typed_tensor_read_einsum_ext_accepts_borrowed_strided_complex_views() {
    let mut backend = CpuEngine::serial();
    let matrix_data = [
        Complex64::new(1.0, 0.0),
        Complex64::new(2.0, 0.0),
        Complex64::new(3.0, 0.0),
        Complex64::new(4.0, 0.0),
        Complex64::new(5.0, 0.0),
        Complex64::new(6.0, 0.0),
    ];
    let vector_data = [
        Complex64::new(10.0, 0.0),
        Complex64::new(20.0, 0.0),
        Complex64::new(30.0, 0.0),
    ];
    let matrix_view = TypedTensorView::from_slice([2, 3], [3, 1], 0, &matrix_data).unwrap();
    let vector_view = TypedTensorView::from_col_major(&[3], &vector_data).unwrap();
    let mut out =
        TypedTensor::<Complex64>::from_vec_col_major(vec![2], vec![Complex64::new(0.0, 0.0); 2])
            .unwrap();

    let result = {
        let session = &mut backend;
        {
            [matrix_view.clone(), vector_view.clone()].einsum_read("ij,j->i", session)
        }
    }
    .unwrap();
    {
        let session = &mut backend;
        {
            [matrix_view, vector_view].einsum_read_into("ij,j->i", session, &mut out)
        }
    }
    .unwrap();

    assert_eq!(
        result.as_slice().unwrap(),
        &[Complex64::new(140.0, 0.0), Complex64::new(320.0, 0.0)]
    );
    assert_eq!(
        out.as_slice().unwrap(),
        &[Complex64::new(140.0, 0.0), Complex64::new(320.0, 0.0)]
    );
}

#[test]
fn public_tensor_read_einsum_ext_accepts_strided_views() {
    let mut backend = CpuEngine::serial();
    let matrix_data = [1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0];
    let vector = Tensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0]).unwrap();
    let row_major_view = TypedTensorView::from_slice([2, 3], [3, 1], 0, &matrix_data).unwrap();
    let inputs = [
        TensorRead::from_view(TensorView::F64(row_major_view)),
        TensorRead::from_tensor(&vector),
    ];

    let result = {
        let session = &mut backend;
        inputs.einsum_read("ij,j->i", session)
    }
    .unwrap();

    assert_f64_tensor(&result, &[2], &[140.0, 320.0]);
}

#[test]
fn public_tensor_read_einsum_ext_accepts_slice_and_integer_subscripts() {
    let mut backend = CpuEngine::serial();
    let matrix_data = [1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0];
    let vector = Tensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0]).unwrap();
    let row_major_view = TypedTensorView::from_slice([2, 3], [3, 1], 0, &matrix_data).unwrap();
    let inputs = [
        TensorRead::from_view(TensorView::F64(row_major_view)),
        TensorRead::from_tensor(&vector),
    ];
    let subscripts = parse_einsum_subscripts("ij,j->i").unwrap();

    let slice_result = {
        let session = &mut backend;
        {
            inputs
                .as_slice()
                .einsum_read_subscripts(&subscripts, session)
        }
    }
    .unwrap();
    let array_result = {
        let session = &mut backend;
        inputs.einsum_read_subscripts(&subscripts, session)
    }
    .unwrap();

    assert_f64_tensor(&slice_result, &[2], &[140.0, 320.0]);
    assert_f64_tensor(&array_result, &[2], &[140.0, 320.0]);
}

#[test]
fn public_einsum_into_writes_dynamic_typed_and_read_outputs() {
    let mut backend = CpuEngine::serial();
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();

    let mut dynamic_out = Tensor::from_vec_col_major(vec![2, 2], vec![-1.0_f64; 4]).unwrap();
    {
        let session = &mut backend;
        {
            [&lhs, &rhs].einsum_into(
                "ij,jk->ik",
                session,
                TensorWrite::from_tensor(&mut dynamic_out),
            )
        }
    }
    .unwrap();
    assert_f64_tensor(&dynamic_out, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);

    let typed_lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let typed_rhs =
        TypedTensor::<f64>::from_vec_col_major(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let mut typed_out = TypedTensor::<f64>::from_vec_col_major(vec![2, 2], vec![-1.0; 4]).unwrap();
    {
        let session = &mut backend;
        {
            [&typed_lhs, &typed_rhs].einsum_into("ij,jk->ik", session, &mut typed_out)
        }
    }
    .unwrap();
    assert_eq!(typed_out.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);

    let mut typed_strided_data = [-1.0_f64; 8];
    {
        let out_view =
            TypedTensorViewMut::from_slice([2, 2], [1, 3], 1, &mut typed_strided_data).unwrap();
        {
            let session = &mut backend;
            {
                [&typed_lhs, &typed_rhs].einsum_into(
                    "ij,jk->ik",
                    session,
                    TypedTensorWrite::from_view(out_view),
                )
            }
        }
        .unwrap();
    }
    assert_eq!(
        typed_strided_data,
        [-1.0, 22.0, 28.0, -1.0, 49.0, 64.0, -1.0, -1.0]
    );

    let mut strided_data = [-1.0_f64; 8];
    {
        let out_view = TensorViewMut::F64(
            TypedTensorViewMut::from_slice([2, 2], [1, 3], 1, &mut strided_data).unwrap(),
        );
        let inputs = [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)];
        {
            let session = &mut backend;
            {
                inputs.einsum_read_into("ij,jk->ik", session, TensorWrite::from_view(out_view))
            }
        }
        .unwrap();
    }
    assert_eq!(
        strided_data,
        [-1.0, 22.0, 28.0, -1.0, 49.0, 64.0, -1.0, -1.0]
    );
}

#[test]
fn public_einsum_read_into_swaps_binary_operands_for_exact_output_order() {
    let mut backend = CpuEngine::serial();
    let lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let rhs = TypedTensor::<f64>::from_vec_col_major(
        vec![4, 3],
        vec![
            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
        ],
    )
    .unwrap();
    let mut output = TypedTensor::<f64>::from_vec_col_major(vec![4, 2], vec![-1.0; 8]).unwrap();

    {
        let session = &mut backend;
        {
            [lhs.as_view(), rhs.as_view()].einsum_read_into("ij,pj->pi", session, &mut output)
        }
    }
    .unwrap();

    assert_eq!(
        output.as_slice().unwrap(),
        &[61.0, 70.0, 79.0, 88.0, 76.0, 88.0, 100.0, 112.0]
    );
}

#[test]
fn typed_view_read_into_compact_planner_handles_notation_and_parsed_labels() {
    let mut backend = CpuEngine::serial();
    let lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let rhs = TypedTensor::<f64>::from_vec_col_major(
        vec![4, 3],
        vec![
            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
        ],
    )
    .unwrap();
    let notation = EinsumNotation::new(
        &[
            &[EinsumAxis::Label(105), EinsumAxis::Label(106)],
            &[EinsumAxis::Label(112), EinsumAxis::Label(106)],
        ],
        &[EinsumAxis::Label(112), EinsumAxis::Label(105)],
    );
    let parsed = crate::EinsumSubscripts::new(&[&[105, 106], &[112, 106]], &[112, 105]);
    let mut notation_out =
        TypedTensor::<f64>::from_vec_col_major(vec![4, 2], vec![-1.0; 8]).unwrap();
    let mut parsed_out = TypedTensor::<f64>::from_vec_col_major(vec![4, 2], vec![-1.0; 8]).unwrap();

    (|| {
        let session = &mut backend;
        {
            [lhs.as_view(), rhs.as_view()].einsum_read_into_notation(
                &notation,
                session,
                &mut notation_out,
            )?;
            [lhs.as_view(), rhs.as_view()].einsum_read_into_subscripts(
                &parsed,
                session,
                &mut parsed_out,
            )
        }
    })()
    .unwrap();

    let expected = [61.0, 70.0, 79.0, 88.0, 76.0, 88.0, 100.0, 112.0];
    assert_eq!(notation_out.as_slice().unwrap(), &expected);
    assert_eq!(parsed_out.as_slice().unwrap(), &expected);
}

#[test]
fn swapped_binary_einsum_preserves_complex_values() {
    let mut backend = CpuEngine::serial();
    let lhs = TypedTensor::<Complex64>::from_vec_col_major(
        vec![2, 2],
        vec![
            Complex64::new(1.0, 1.0),
            Complex64::new(2.0, 0.0),
            Complex64::new(3.0, 0.0),
            Complex64::new(4.0, -1.0),
        ],
    )
    .unwrap();
    let rhs = TypedTensor::<Complex64>::from_vec_col_major(
        vec![2, 2],
        vec![
            Complex64::new(1.0, 0.0),
            Complex64::new(0.0, 2.0),
            Complex64::new(3.0, 0.0),
            Complex64::new(4.0, 0.0),
        ],
    )
    .unwrap();
    let mut output = TypedTensor::<Complex64>::from_vec_col_major(
        vec![2, 2],
        vec![Complex64::new(-1.0, 0.0); 4],
    )
    .unwrap();

    {
        let session = &mut backend;
        {
            [lhs.as_view(), rhs.as_view()].einsum_read_into("ij,pj->pi", session, &mut output)
        }
    }
    .unwrap();

    assert_eq!(
        output.as_slice().unwrap(),
        &[
            Complex64::new(10.0, 1.0),
            Complex64::new(10.0, 2.0),
            Complex64::new(14.0, -3.0),
            Complex64::new(16.0, 0.0),
        ]
    );
}

#[test]
fn swapped_binary_einsum_writes_strided_output_without_reordering_payload() {
    let mut backend = CpuEngine::serial();
    let lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let rhs = TypedTensor::<f64>::from_vec_col_major(
        vec![4, 3],
        vec![
            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
        ],
    )
    .unwrap();
    let mut data = [-1.0; 10];
    let out = TypedTensorViewMut::from_slice([4, 2], [1, 5], 1, &mut data).unwrap();

    {
        let session = &mut backend;
        {
            [lhs.as_view(), rhs.as_view()].einsum_read_into(
                "ij,pj->pi",
                session,
                TypedTensorWrite::from_view(out),
            )
        }
    }
    .unwrap();

    assert_eq!(
        data,
        [-1.0, 61.0, 70.0, 79.0, 88.0, -1.0, 76.0, 88.0, 100.0, 112.0]
    );
}

#[test]
fn public_einsum_into_preserves_complex_dtype() {
    let mut backend = CpuEngine::serial();
    let lhs = TypedTensor::<Complex64>::from_vec_col_major(
        vec![2, 2],
        vec![
            Complex64::new(1.0, 1.0),
            Complex64::new(2.0, -1.0),
            Complex64::new(3.0, 0.0),
            Complex64::new(4.0, 2.0),
        ],
    )
    .unwrap();
    let rhs = TypedTensor::<Complex64>::from_vec_col_major(
        vec![2, 1],
        vec![Complex64::new(5.0, 0.0), Complex64::new(6.0, -1.0)],
    )
    .unwrap();
    let mut out =
        TypedTensor::<Complex64>::from_vec_col_major(vec![2, 1], vec![Complex64::new(0.0, 0.0); 2])
            .unwrap();

    {
        let session = &mut backend;
        [&lhs, &rhs].einsum_into("ij,jk->ik", session, &mut out)
    }
    .unwrap();

    assert_eq!(
        out.as_slice().unwrap(),
        &[Complex64::new(23.0, 2.0), Complex64::new(36.0, 3.0)]
    );
}

#[test]
fn public_einsum_into_rejects_output_shape_and_dtype_mismatch() {
    let mut backend = CpuEngine::serial();
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let mut wrong_shape = Tensor::from_vec_col_major(vec![4], vec![0.0_f64; 4]).unwrap();
    let shape_err = {
        let session = &mut backend;
        {
            [&lhs, &rhs].einsum_into(
                "ij,jk->ik",
                session,
                TensorWrite::from_tensor(&mut wrong_shape),
            )
        }
    }
    .unwrap_err();
    assert!(matches!(
        shape_err,
        Error::Validation {
            op: "TensorEinsumIntoExt::einsum_into",
            source: tenferro_tensor::ValidationError::ShapeMismatch(_),
        }
    ));

    let mut wrong_dtype = Tensor::from_vec_col_major(vec![2, 2], vec![0.0_f32; 4]).unwrap();
    let dtype_err = {
        let session = &mut backend;
        {
            [&lhs, &rhs].einsum_into(
                "ij,jk->ik",
                session,
                TensorWrite::from_tensor(&mut wrong_dtype),
            )
        }
    }
    .unwrap_err();
    assert!(matches!(
        dtype_err,
        Error::Tensor(tenferro_tensor::Error::Validation {
            op: "TensorEinsumIntoExt::einsum_into",
            source: tenferro_tensor::ValidationError::DTypeMismatch { .. },
        })
    ));
}

#[test]
fn einsum_into_gemm_fast_path_uses_exact_output_operand_order() {
    let source = include_str!("eager.rs");
    let start = source
        .find("pub(crate) fn eager_einsum_exec_read_into")
        .expect("missing eager_einsum_exec_read_into");
    let tail = &source[start..];
    let end = tail
        .find("pub(crate) fn eager_einsum_exec_read_into_accum")
        .expect("missing following function");
    let body = &tail[..end];

    assert!(
        include_str!("eager.rs").contains("try_build_exact_output_binary_dot_plan"),
        "binary read-into must plan an exact output order, including operand swapping"
    );
    assert!(
        body.contains("binary_dot_plan_for_read_into"),
        "tree-based read-into must share the exact-output planner"
    );
}

#[test]
fn ordinary_binary_read_into_dispatches_before_tree_preparation() {
    let source = include_str!("concrete.rs");
    let typed = source
        .split("TypedTensorReadEinsumIntoExt<T> for [TypedTensorView")
        .nth(1)
        .expect("non-AD typed-view implementation");
    assert!(
        typed.find("parse_fast_ascii_binary_labels").unwrap()
            < typed.find("parse_einsum_notation").unwrap(),
        "the measured non-AD string API must dispatch before the generic parser"
    );
    let start = source
        .find("fn tensor_read_einsum_into_subscripts(")
        .expect("missing read-into entry point");
    let tail = &source[start..];
    let end = tail
        .find("fn validate_output(")
        .expect("missing following function");
    let body = &tail[..end];
    let dispatch = body
        .find("read_binary_dot_config")
        .expect("ordinary binary read-into should use the compact config path");
    let tree = body
        .find("prepare_subscripts_internal")
        .expect("general einsum planning fallback should remain");
    assert!(
        dispatch < tree,
        "compact binary dispatch must precede tree creation"
    );
    assert!(
        !body.contains("binary_dot_plan_for_read_into"),
        "concrete read-into must not build result/target label vectors"
    );
}

#[test]
fn einsum_into_gemm_fast_path_dispatches_to_backend_into_before_owned_fallback() {
    let source = include_str!("eager.rs");
    let start = source
        .find("pub(crate) fn eager_einsum_exec_read_into")
        .expect("missing eager_einsum_exec_read_into");
    let tail = &source[start..];
    let end = tail
        .find("pub(crate) fn eager_einsum_exec_read_into_accum")
        .expect("missing following function");
    let body = &tail[..end];

    let into_call = body
        .find("execute_binary_dot_read_into")
        .expect("GEMM-compatible einsum_into must call backend read-into");
    let fallback = body
        .find("eager_einsum_exec_values_with_output(exec, values, tree, Some(out))")
        .expect("general einsum_into should walk the tree with the destination");

    assert!(
        into_call < fallback,
        "GEMM-compatible einsum_into should try backend read-into before the tree walk"
    );
}

#[test]
fn prepared_read_into_reuses_borrowed_metadata_and_binary_plan() {
    let source = include_str!("concrete.rs");
    let start = source
        .find("pub fn execute_read_into<'a, I>(")
        .expect("missing prepared read-into method");
    let tail = &source[start..];
    let end = tail
        .find("pub fn execute_read_into_accum<'a, I>(")
        .expect("missing following method");
    let body = &tail[..end];

    assert!(body.contains("validate_read_inputs(inputs"));
    assert!(body.contains("validate_cached_output(&out"));
    assert!(body.contains("if let Some(binary_dot) = &self.binary_dot"));
    assert!(!body.contains("read_input_specs(inputs)"));
}

#[test]
fn concrete_einsum_plan_executes_without_replanning_contract() {
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs2 =
        Tensor::from_vec_col_major(vec![3, 2], vec![2.0_f64, 0.0, 1.0, 3.0, 4.0, 5.0]).unwrap();
    let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik").unwrap();

    let mut backend = CpuEngine::serial();
    let (first, second) = (|| {
        let session = &mut backend;
        {
            let first = plan.execute([&lhs, &rhs], session)?;
            let second = plan.execute([&lhs, &rhs2], session)?;
            Ok::<_, crate::Error>((first, second))
        }
    })()
    .unwrap();

    assert_f64_tensor(&first, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);
    assert_f64_tensor(&second, &[2, 2], &[7.0, 10.0, 40.0, 52.0]);
}

#[test]
fn concrete_einsum_plan_execute_into_writes_reused_outputs() {
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs2 =
        Tensor::from_vec_col_major(vec![3, 2], vec![2.0_f64, 0.0, 1.0, 3.0, 4.0, 5.0]).unwrap();
    let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik").unwrap();
    let mut backend = CpuEngine::serial();
    let mut out = Tensor::from_vec_col_major(vec![2, 2], vec![-1.0_f64; 4]).unwrap();

    let first = (|| {
        let session = &mut backend;
        {
            plan.execute_into([&lhs, &rhs], session, TensorWrite::from_tensor(&mut out))?;
            let first = out.duplicate().unwrap();
            plan.execute_into([&lhs, &rhs2], session, TensorWrite::from_tensor(&mut out))?;
            Ok::<_, crate::Error>(first)
        }
    })()
    .unwrap();
    assert_f64_tensor(&first, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);
    assert_f64_tensor(&out, &[2, 2], &[7.0, 10.0, 40.0, 52.0]);
}

#[test]
fn concrete_einsum_plan_execute_read_into_accum_updates_outputs() {
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let mut backend = CpuEngine::serial();

    let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik").unwrap();
    let mut out = Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64; 4]).unwrap();
    let fallback_plan = ConcreteEinsumPlan::prepare([&lhs], "ij->").unwrap();
    let mut scalar_out = Tensor::from_vec_col_major(vec![], vec![2.0_f64]).unwrap();
    let accumulation =
        DotGeneralAccumulation::scaled(ContractionScalar::F64(0.5), ContractionScalar::F64(2.0))
            .unwrap();
    (|| {
        let session = &mut backend;
        {
            plan.execute_read_into_accum(
                [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)],
                session,
                DotGeneralAccumulation::add_to(DType::F64).unwrap(),
                TensorWrite::from_tensor(&mut out),
            )?;
            fallback_plan.execute_read_into_accum(
                [TensorRead::from_tensor(&lhs)],
                session,
                accumulation,
                TensorWrite::from_tensor(&mut scalar_out),
            )
        }
    })()
    .unwrap();
    assert_f64_tensor(&out, &[2, 2], &[23.0, 29.0, 50.0, 65.0]);
    assert_f64_tensor(&scalar_out, &[], &[14.5]);
}

#[test]
fn concrete_einsum_plan_execute_typed_and_read_into_outputs() {
    let lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let rhs =
        TypedTensor::<f64>::from_vec_col_major(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let plan = ConcreteEinsumPlan::prepare_typed([&lhs, &rhs], "ij,jk->ik").unwrap();
    let mut backend = CpuEngine::serial();

    let mut typed_out = TypedTensor::<f64>::from_vec_col_major(vec![2, 2], vec![-1.0; 4]).unwrap();
    let mut typed_strided_data = [-1.0_f64; 8];
    (|| {
        let session = &mut backend;
        {
            plan.execute_typed_into([&lhs, &rhs], session, &mut typed_out)?;
            let out_view =
                TypedTensorViewMut::from_slice([2, 2], [1, 3], 1, &mut typed_strided_data).unwrap();
            plan.execute_typed_into([&lhs, &rhs], session, TypedTensorWrite::from_view(out_view))
        }
    })()
    .unwrap();
    assert_eq!(typed_out.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);
    assert_eq!(
        typed_strided_data,
        [-1.0, 22.0, 28.0, -1.0, 49.0, 64.0, -1.0, -1.0]
    );

    let lhs_erased = Tensor::from(lhs.duplicate().unwrap());
    let rhs_erased = Tensor::from(rhs.duplicate().unwrap());
    let mut strided_data = [-1.0_f64; 8];
    {
        let session = &mut backend;
        {
            let out_view = TensorViewMut::F64(
                TypedTensorViewMut::from_slice([2, 2], [1, 3], 1, &mut strided_data).unwrap(),
            );
            plan.execute_read_into(
                [
                    TensorRead::from_tensor(&lhs_erased),
                    TensorRead::from_tensor(&rhs_erased),
                ],
                session,
                TensorWrite::from_view(out_view),
            )
        }
    }
    .unwrap();
    assert_eq!(
        strided_data,
        [-1.0, 22.0, 28.0, -1.0, 49.0, 64.0, -1.0, -1.0]
    );
}

#[test]
fn prepared_binary_read_into_reuses_swapped_dot_metadata() {
    let mut backend = CpuEngine::serial();
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs = Tensor::from_vec_col_major(
        vec![4, 3],
        vec![
            1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
        ],
    )
    .unwrap();
    let reads = [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)];
    let plan = ConcreteEinsumPlan::prepare_read(&reads, "ij,pj->pi").unwrap();
    let mut out = Tensor::from_vec_col_major(vec![4, 2], vec![-1.0; 8]).unwrap();

    {
        let session = &mut backend;
        {
            plan.execute_read_into(&reads, session, TensorWrite::from_tensor(&mut out))
        }
    }
    .unwrap();

    assert_f64_tensor(
        &out,
        &[4, 2],
        &[61.0, 70.0, 79.0, 88.0, 76.0, 88.0, 100.0, 112.0],
    );
}

#[test]
fn concrete_einsum_plan_execute_into_rejects_incompatible_output() {
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik").unwrap();
    let mut backend = CpuEngine::serial();
    let mut wrong_shape = Tensor::from_vec_col_major(vec![4], vec![0.0_f64; 4]).unwrap();

    let err = {
        let session = &mut backend;
        {
            plan.execute_into(
                [&lhs, &rhs],
                session,
                TensorWrite::from_tensor(&mut wrong_shape),
            )
        }
    }
    .unwrap_err();

    assert!(matches!(
        err,
        Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            source: tenferro_tensor::ValidationError::ShapeMismatch(_),
        }
    ));
}

#[test]
fn concrete_einsum_plan_prepares_integer_and_read_variants() {
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let typed_lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let typed_rhs =
        TypedTensor::<f64>::from_vec_col_major(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let read_inputs = [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)];
    let subscripts = parse_einsum_subscripts("ij,jk->ik").unwrap();

    let tensor_plan = ConcreteEinsumPlan::prepare_subscripts([&lhs, &rhs], &subscripts).unwrap();
    let typed_plan =
        ConcreteEinsumPlan::prepare_typed_subscripts([&typed_lhs, &typed_rhs], &subscripts)
            .unwrap();
    let read_string_plan =
        ConcreteEinsumPlan::prepare_read(read_inputs.as_slice(), "ij,jk->ik").unwrap();
    let read_integer_plan =
        ConcreteEinsumPlan::prepare_read_subscripts(read_inputs.as_slice(), &subscripts).unwrap();
    let debug = format!("{tensor_plan:?}");

    let mut backend = CpuEngine::serial();
    let tensor_result = {
        let session = &mut backend;
        tensor_plan.execute([&lhs, &rhs], session)
    }
    .unwrap();
    let (typed_result, read_string_result, read_integer_result) = (|| {
        let session = &mut backend;
        {
            let typed_result = typed_plan.execute_typed([&typed_lhs, &typed_rhs], session)?;
            let read_string_result =
                read_string_plan.execute_read(read_inputs.as_slice(), session)?;
            let read_integer_result =
                read_integer_plan.execute_read(read_inputs.as_slice(), session)?;
            Ok::<_, crate::Error>((typed_result, read_string_result, read_integer_result))
        }
    })()
    .unwrap();

    assert!(debug.contains("ConcreteEinsumPlan"));
    assert!(debug.contains("F64"));
    assert_f64_tensor(&tensor_result, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);
    assert_eq!(typed_result.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);
    assert_f64_tensor(&read_string_result, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);
    assert_f64_tensor(&read_integer_result, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);
}

#[test]
fn concrete_einsum_plan_executes_read_and_typed_inputs() {
    let lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let rhs =
        TypedTensor::<f64>::from_vec_col_major(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let plan = ConcreteEinsumPlan::prepare_typed([&lhs, &rhs], "ij,jk->ik").unwrap();
    let lhs_erased = Tensor::from(lhs.duplicate().unwrap());
    let rhs_erased = Tensor::from(rhs.duplicate().unwrap());

    let mut backend = CpuEngine::serial();
    let (typed, read) = (|| {
        let session = &mut backend;
        {
            let typed = plan.execute_typed([&lhs, &rhs], session)?;
            let read = plan.execute_read(
                [
                    TensorRead::from_tensor(&lhs_erased),
                    TensorRead::from_tensor(&rhs_erased),
                ],
                session,
            )?;
            Ok::<_, crate::Error>((typed, read))
        }
    })()
    .unwrap();

    assert_eq!(typed.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);
    assert_f64_tensor(&read, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);
}

#[test]
fn concrete_einsum_plan_rejects_shape_and_dtype_mismatches() {
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let wrong_shape = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
    let wrong_dtype = Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f32; 6]).unwrap();
    let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik").unwrap();

    let mut backend = CpuEngine::serial();
    let (shape_err, dtype_err) = {
        let session = &mut backend;
        {
            let shape_err = plan.execute([&lhs, &wrong_shape], session).unwrap_err();
            let dtype_err = plan.execute([&lhs, &wrong_dtype], session).unwrap_err();
            (shape_err, dtype_err)
        }
    };

    assert!(matches!(
        shape_err,
        Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            source: tenferro_tensor::ValidationError::ShapeMismatch(_),
        }
    ));
    assert!(matches!(
        dtype_err,
        Error::Tensor(tenferro_tensor::Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            source: tenferro_tensor::ValidationError::DTypeMismatch { .. },
        })
    ));
}

#[test]
fn concrete_einsum_public_api_reports_parse_and_input_count_errors() {
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let mut backend = CpuEngine::serial();

    let parse_err = {
        let session = &mut backend;
        [&lhs, &rhs].einsum("ij,(jk)->ik", session)
    }
    .unwrap_err();
    let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik").unwrap();
    let count_err = {
        let session = &mut backend;
        plan.execute([&lhs], session)
    }
    .unwrap_err();

    assert!(matches!(parse_err, Error::InvalidSubscripts { .. }));
    assert!(matches!(
        count_err,
        Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            ..
        }
    ));
}

#[test]
fn concrete_einsum_typed_result_reports_defensive_dtype_mismatch() {
    let tensor = Tensor::from_vec_col_major(vec![1], vec![1.0_f32]).unwrap();

    let err = crate::concrete::into_typed_result::<f64>(tensor, "test typed result").unwrap_err();

    assert!(matches!(
        err,
        Error::Tensor(tenferro_tensor::Error::Validation {
            op: "test typed result",
            source: tenferro_tensor::ValidationError::DTypeMismatch { .. },
        })
    ));
}

// ---------------------------------------------------------------------------
// Direct session tests for the ConcreteEinsumPlan execute surface
// ---------------------------------------------------------------------------

fn f64_plan_fixture() -> (Tensor, Tensor, ConcreteEinsumPlan) {
    let lhs =
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let rhs =
        Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ik").unwrap();
    (lhs, rhs, plan)
}

#[test]
fn concrete_einsum_plan_execute_in_session_matches_expected_values() {
    let (lhs, rhs, plan) = f64_plan_fixture();
    let mut backend = CpuEngine::serial();

    let in_session = {
        let session = &mut backend;
        plan.execute([&lhs, &rhs], session)
    }
    .unwrap();
    assert_f64_tensor(&in_session, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);

    let reads = [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)];
    let in_session = {
        let session = &mut backend;
        plan.execute_read(reads.clone(), session)
    }
    .unwrap();
    assert_f64_tensor(&in_session, &[2, 2], &[22.0, 28.0, 49.0, 64.0]);

    let typed_lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let typed_rhs =
        TypedTensor::<f64>::from_vec_col_major(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let typed_plan =
        ConcreteEinsumPlan::prepare_typed([&typed_lhs, &typed_rhs], "ij,jk->ik").unwrap();
    let in_session = {
        let session = &mut backend;
        typed_plan.execute_typed([&typed_lhs, &typed_rhs], session)
    }
    .unwrap();
    assert_eq!(in_session.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);
}

#[test]
fn concrete_einsum_plan_execute_in_session_rejects_bad_inputs() {
    let (lhs, _rhs, plan) = f64_plan_fixture();
    let wrong_shape = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
    let wrong_dtype = Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f32; 6]).unwrap();
    let mut backend = CpuEngine::serial();

    let in_session = {
        let session = &mut backend;
        plan.execute([&lhs, &wrong_shape], session)
    }
    .unwrap_err();
    assert!(matches!(
        in_session,
        Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            source: tenferro_tensor::ValidationError::ShapeMismatch(_),
        }
    ));

    let in_session = {
        let session = &mut backend;
        plan.execute([&lhs, &wrong_dtype], session)
    }
    .unwrap_err();
    assert!(matches!(
        in_session,
        Error::Tensor(tenferro_tensor::Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            source: tenferro_tensor::ValidationError::DTypeMismatch { .. },
        })
    ));

    let in_session = {
        let session = &mut backend;
        plan.execute([&lhs], session)
    }
    .unwrap_err();
    assert!(matches!(
        in_session,
        Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            ..
        }
    ));
}

#[test]
fn concrete_einsum_plan_execute_typed_in_session_validates_dtype() {
    let typed_lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let typed_rhs =
        TypedTensor::<f64>::from_vec_col_major(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let plan = ConcreteEinsumPlan::prepare_typed([&typed_lhs, &typed_rhs], "ij,jk->ik").unwrap();
    let mut backend = CpuEngine::serial();

    let in_session = {
        let session = &mut backend;
        plan.execute_typed([&typed_lhs, &typed_rhs], session)
    }
    .unwrap();
    assert_eq!(in_session.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);

    // Requesting a scalar type different from the prepared dtype must fail
    // through validate_inputs' dtype contract.
    let f32_lhs = TypedTensor::<f32>::from_vec_col_major(vec![2, 3], vec![1.0_f32; 6]).unwrap();
    let f32_rhs = TypedTensor::<f32>::from_vec_col_major(vec![3, 2], vec![1.0_f32; 6]).unwrap();
    let in_session = {
        let session = &mut backend;
        plan.execute_typed([&f32_lhs, &f32_rhs], session)
    }
    .unwrap_err();
    assert!(matches!(
        in_session,
        Error::Tensor(tenferro_tensor::Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            source: tenferro_tensor::ValidationError::DTypeMismatch { .. },
        })
    ));
}

#[test]
fn concrete_einsum_plan_execute_into_in_session_validates_output() {
    let (lhs, rhs, plan) = f64_plan_fixture();
    let mut backend = CpuEngine::serial();

    // execute_into: incompatible output shape.
    let mut wrong_shape = Tensor::from_vec_col_major(vec![4], vec![0.0_f64; 4]).unwrap();
    let in_session = {
        let session = &mut backend;
        {
            plan.execute_into(
                [&lhs, &rhs],
                session,
                TensorWrite::from_tensor(&mut wrong_shape),
            )
        }
    }
    .unwrap_err();
    assert!(matches!(
        in_session,
        Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            source: tenferro_tensor::ValidationError::ShapeMismatch(_),
        }
    ));

    // execute_read_into: incompatible output dtype.
    let mut wrong_dtype_out = Tensor::from_vec_col_major(vec![2, 2], vec![0.0_f32; 4]).unwrap();
    let in_session = {
        let session = &mut backend;
        {
            plan.execute_read_into(
                [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)],
                session,
                TensorWrite::from_tensor(&mut wrong_dtype_out),
            )
        }
    }
    .unwrap_err();
    assert!(matches!(
        in_session,
        Error::Tensor(tenferro_tensor::Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            source: tenferro_tensor::ValidationError::DTypeMismatch { .. },
        })
    ));

    // execute_typed_into: incompatible output shape.
    let typed_lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let typed_rhs =
        TypedTensor::<f64>::from_vec_col_major(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let typed_plan =
        ConcreteEinsumPlan::prepare_typed([&typed_lhs, &typed_rhs], "ij,jk->ik").unwrap();
    let mut typed_out = TypedTensor::<f64>::from_vec_col_major(vec![3], vec![0.0; 3]).unwrap();
    let in_session = {
        let session = &mut backend;
        {
            typed_plan.execute_typed_into([&typed_lhs, &typed_rhs], session, &mut typed_out)
        }
    }
    .unwrap_err();
    assert!(matches!(
        in_session,
        Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            source: tenferro_tensor::ValidationError::ShapeMismatch(_),
        }
    ));

    // execute_read_into_accum: incompatible output shape.
    let mut accum_out = Tensor::from_vec_col_major(vec![4], vec![0.0_f64; 4]).unwrap();
    let accumulation = DotGeneralAccumulation::add_to(DType::F64).unwrap();
    let in_session = {
        let session = &mut backend;
        {
            plan.execute_read_into_accum(
                [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)],
                session,
                accumulation,
                TensorWrite::from_tensor(&mut accum_out),
            )
        }
    }
    .unwrap_err();
    assert!(matches!(
        in_session,
        Error::Validation {
            op: "ConcreteEinsumPlan::execute",
            source: tenferro_tensor::ValidationError::ShapeMismatch(_),
        }
    ));

    // Valid into execution writes the expected values.
    let mut out = Tensor::from_vec_col_major(vec![2, 2], vec![-1.0_f64; 4]).unwrap();
    {
        let session = &mut backend;
        {
            plan.execute_into([&lhs, &rhs], session, TensorWrite::from_tensor(&mut out))
        }
    }
    .unwrap();
    assert_eq!(out.as_slice::<f64>().unwrap(), &[22.0, 28.0, 49.0, 64.0]);
}

/// A non-`Send` output adapter: `PhantomData<Rc<()>>` makes the type `!Send`
/// (a `Cell` would not — `Cell<T>` is `Send` but `!Sync`). Passing this
/// adapter to `ConcreteEinsumPlan::execute_typed_into` only compiles because
/// the caller converts `out` to `TypedTensorWrite` before the
/// engine-holding block, so `O` itself never needs `Send`.
struct NonSendIntoWrite<'a, T> {
    write: TypedTensorWrite<'a, T>,
    _not_send: std::marker::PhantomData<std::rc::Rc<()>>,
}

impl<'a, T> From<NonSendIntoWrite<'a, T>> for TypedTensorWrite<'a, T> {
    fn from(adapter: NonSendIntoWrite<'a, T>) -> Self {
        adapter.write
    }
}

#[test]
fn concrete_einsum_plan_execute_typed_into_accepts_non_send_adapter() {
    let typed_lhs =
        TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let typed_rhs =
        TypedTensor::<f64>::from_vec_col_major(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .unwrap();
    let plan = ConcreteEinsumPlan::prepare_typed([&typed_lhs, &typed_rhs], "ij,jk->ik").unwrap();

    let mut backend = CpuEngine::serial();
    let mut out = TypedTensor::<f64>::from_vec_col_major(vec![2, 2], vec![0.0; 4]).unwrap();
    let write = TypedTensorWrite::from(NonSendIntoWrite {
        write: TypedTensorWrite::from_tensor(&mut out),
        _not_send: std::marker::PhantomData,
    });
    {
        let session = &mut backend;
        {
            plan.execute_typed_into([&typed_lhs, &typed_rhs], session, write)
        }
    }
    .unwrap();
    assert_eq!(out.as_slice().unwrap(), &[22.0, 28.0, 49.0, 64.0]);
}

/// N-ary read-into: the last contraction step writes into the destination
/// (exact and transposed output labels, and a label the tree pre-reduces), and
/// a tree without a contraction step hands the destination back for a copy;
/// every case matches the allocating path.
#[test]
fn prepared_nary_accumulation_reuses_plan_with_runtime_scalars() {
    let a = Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
    let id = Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 0.0, 0.0, 1.0]).unwrap();
    let reads = [
        TensorRead::from_tensor(&a),
        TensorRead::from_tensor(&id),
        TensorRead::from_tensor(&id),
    ];
    let mut plan = ConcreteEinsumPlan::prepare_read(&reads, "ij,jk,kl->il").unwrap();
    let mut backend = CpuEngine::new(tenferro_cpu::Pool::builder().threads(1).build().unwrap());
    let mut out = Tensor::from_vec_col_major(vec![2, 2], vec![10.0_f64; 4]).unwrap();
    for (alpha, beta, expected) in [
        (2.0, 3.0, [32.0, 34.0, 36.0, 38.0]),
        (1.0, 0.0, [1.0, 2.0, 3.0, 4.0]),
    ] {
        {
            let session = &mut backend;
            {
                plan.execute_read_into_accum(
                    &reads,
                    session,
                    DotGeneralAccumulation {
                        lhs_conj: false,
                        rhs_conj: false,
                        alpha: ContractionScalar::F64(alpha),
                        beta: ContractionScalar::F64(beta),
                    },
                    TensorWrite::from_tensor(&mut out),
                )
            }
        }
        .unwrap();
        assert_eq!(out.as_slice::<f64>().unwrap(), expected);
    }
    let stats = plan.cpu_cache_stats().unwrap();
    assert_eq!((stats.entries, stats.hits, stats.misses), (1, 1, 1));
    plan.clear_cpu_cache();
    assert_eq!(plan.cpu_cache_stats().unwrap().entries, 0);
}

#[test]
fn nary_read_into_matches_allocating_execution() {
    let mut backend = CpuEngine::serial();
    let a = Tensor::from_vec_col_major(vec![2, 3], (1..=6).map(f64::from).collect()).unwrap();
    let b = Tensor::from_vec_col_major(vec![3, 4], (1..=12).map(|x| f64::from(x) * 0.5).collect())
        .unwrap();
    let c = Tensor::from_vec_col_major(vec![4, 2], (1..=8).map(|x| f64::from(x) - 3.0).collect())
        .unwrap();
    let reads = [
        TensorRead::from_tensor(&a),
        TensorRead::from_tensor(&b),
        TensorRead::from_tensor(&c),
    ];

    for (subscripts, shape) in [
        // The last step writes the exact output labels.
        ("ij,jk,kl->il", vec![2, 2]),
        // Transposed output order.
        ("ij,jk,kl->li", vec![2, 2]),
        // The tree sums `l` out of `c` before the last step.
        ("ij,jk,kl->i", vec![2]),
    ] {
        let plan = ConcreteEinsumPlan::prepare_read(&reads, subscripts).unwrap();
        let expected = {
            let session = &mut backend;
            plan.execute_read(&reads, session)
        }
        .unwrap();
        let len = shape.iter().product();
        let mut out = Tensor::from_vec_col_major(shape.clone(), vec![-9.0_f64; len]).unwrap();
        {
            let session = &mut backend;
            {
                plan.execute_read_into(&reads, session, TensorWrite::from_tensor(&mut out))
            }
        }
        .unwrap();
        assert_f64_tensor(&out, &shape, expected.as_slice::<f64>().unwrap());
        if shape == [2, 2] {
            for (strides, offset, positions) in
                [([1, 3], 1, [1, 2, 4, 5]), ([-1, 3], 2, [2, 1, 5, 4])]
            {
                // Repeat the same layout to exercise the warm numerical plan.
                for _ in 0..2 {
                    let mut backing = [-99.0_f64; 9];
                    let view = TensorViewMut::F64(
                        TypedTensorViewMut::from_slice([2, 2], strides, offset, &mut backing)
                            .unwrap(),
                    );
                    {
                        let session = &mut backend;
                        {
                            plan.execute_read_into(&reads, session, TensorWrite::from_view(view))
                        }
                    }
                    .unwrap();
                    let mut expected_backing = [-99.0; 9];
                    for (position, value) in positions
                        .into_iter()
                        .zip(expected.as_slice::<f64>().unwrap())
                    {
                        expected_backing[position] = *value;
                    }
                    assert_eq!(backing, expected_backing);
                }
            }
        }
    }

    // A single operand has no contraction step: the result is materialized
    // and copied into the handed-back destination.
    let unary = [TensorRead::from_tensor(&a)];
    let plan = ConcreteEinsumPlan::prepare_read(&unary, "ij->ji").unwrap();
    let mut out = Tensor::from_vec_col_major(vec![3, 2], vec![-9.0_f64; 6]).unwrap();
    {
        let session = &mut backend;
        {
            plan.execute_read_into(&unary, session, TensorWrite::from_tensor(&mut out))
        }
    }
    .unwrap();
    assert_f64_tensor(&out, &[3, 2], &[1.0, 3.0, 5.0, 2.0, 4.0, 6.0]);
}

/// `ij,jk->ki` needs the swapped operand order for an exact-output dot, so the
/// prepared accumulate path must swap the conjugation flags with the operands.
#[test]
fn prepared_swapped_binary_dot_accumulates_with_conjugation_on_the_named_inputs() {
    let lhs_data: Vec<Complex64> = (0..6)
        .map(|index| Complex64::new(index as f64 + 1.0, 0.5 * index as f64 - 1.0))
        .collect();
    let rhs_data: Vec<Complex64> = (0..6)
        .map(|index| Complex64::new(2.0 - index as f64, 0.25 * index as f64))
        .collect();
    let lhs = Tensor::from_vec_col_major(vec![2, 3], lhs_data.clone()).unwrap();
    let rhs = Tensor::from_vec_col_major(vec![3, 2], rhs_data.clone()).unwrap();
    let plan = ConcreteEinsumPlan::prepare([&lhs, &rhs], "ij,jk->ki").unwrap();
    let initial = Complex64::new(0.5, -0.5);
    let (alpha, beta) = (Complex64::new(2.0, 1.0), Complex64::new(1.0, 0.0));

    for (lhs_conj, rhs_conj) in [(true, false), (false, true)] {
        let mut out = Tensor::from_vec_col_major(vec![2, 2], vec![initial; 4]).unwrap();
        let mut accumulation = DotGeneralAccumulation::scaled(
            ContractionScalar::C64(alpha),
            ContractionScalar::C64(beta),
        )
        .unwrap();
        accumulation.lhs_conj = lhs_conj;
        accumulation.rhs_conj = rhs_conj;
        let mut backend = CpuEngine::serial();
        {
            let session = &mut backend;
            {
                plan.execute_read_into_accum(
                    [TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs)],
                    session,
                    accumulation,
                    TensorWrite::from_tensor(&mut out),
                )
            }
        }
        .unwrap();

        let values = out.as_slice::<Complex64>().unwrap();
        for i in 0..2 {
            for k in 0..2 {
                let sum: Complex64 = (0..3)
                    .map(|j| {
                        let a = lhs_data[i + 2 * j];
                        let b = rhs_data[j + 3 * k];
                        (if lhs_conj { a.conj() } else { a })
                            * (if rhs_conj { b.conj() } else { b })
                    })
                    .sum();
                let expected = alpha * sum + beta * initial;
                let got = values[k + 2 * i];
                assert!(
                    (got - expected).norm() < 1e-12,
                    "conj=({lhs_conj}, {rhs_conj}) [{k}, {i}]: {got} vs {expected}"
                );
            }
        }
    }
}