use num_complex::{Complex32, Complex64};
use tenferro_cpu::CpuBackend;
use tenferro_runtime::extension::{ExecInstruction, ExecOp, ExecProgram};
use tenferro_runtime::{DType, ExtensionCacheStore, ExtensionExecutionContext, GraphExecutor};
use tenferro_tensor::{DotGeneralConfig, Tensor, TensorValue, TypedTensor};

#[cfg(feature = "cuda")]
use tenferro_gpu::{download_tensor, upload_tensor, CudaBackend};

fn f64_tensor(shape: Vec<usize>, data: Vec<f64>) -> Tensor {
    Tensor::F64(TypedTensor::from_vec_col_major(shape, data).unwrap())
}

fn scalar_extents(
    output_count: usize,
) -> Vec<Vec<tenferro_ops::ShapeExtent<tenferro_ops::dim_expr::DimExpr>>> {
    vec![Vec::new(); output_count]
}

fn const_shape(shape: &[usize]) -> Vec<tenferro_ops::dim_expr::DimExpr> {
    shape
        .iter()
        .map(|&dim| tenferro_ops::dim_expr::DimExpr::Const(dim))
        .collect()
}

fn exact_extents(
    shape: &[usize],
) -> Vec<tenferro_ops::ShapeExtent<tenferro_ops::dim_expr::DimExpr>> {
    const_shape(shape)
        .into_iter()
        .map(tenferro_ops::ShapeExtent::exact)
        .collect()
}

fn matmul_config() -> DotGeneralConfig {
    DotGeneralConfig {
        lhs_contracting_dims: vec![1],
        rhs_contracting_dims: vec![0],
        lhs_batch_dims: vec![],
        rhs_batch_dims: vec![],
    }
}

fn cpu_parity_program() -> ExecProgram {
    ExecProgram {
        instructions: vec![
            ExecInstruction {
                op: ExecOp::Add,
                input_slots: vec![0, 1],
                output_slots: vec![4],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![false, false],
            },
            ExecInstruction {
                op: ExecOp::Exp,
                input_slots: vec![4],
                output_slots: vec![5],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![false],
            },
            ExecInstruction {
                op: ExecOp::ShapeOf { axis: 0 },
                input_slots: vec![5],
                output_slots: vec![6],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![false],
            },
            ExecInstruction {
                op: ExecOp::DynamicTruncate { axis: 0 },
                input_slots: vec![5, 6],
                output_slots: vec![7],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![true, true],
            },
            ExecInstruction {
                op: ExecOp::ReduceSum { axes: vec![1] },
                input_slots: vec![7],
                output_slots: vec![8],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![true],
            },
            ExecInstruction {
                op: ExecOp::DotGeneral(matmul_config()),
                input_slots: vec![0, 1],
                output_slots: vec![9],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![true, true],
            },
            ExecInstruction {
                op: ExecOp::PadToMatch { axis: 0 },
                input_slots: vec![2, 3],
                output_slots: vec![11],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![true, true],
            },
            ExecInstruction {
                op: ExecOp::Constant {
                    dtype: DType::F64,
                    bytes: 1.5_f64.to_le_bytes().to_vec(),
                },
                input_slots: vec![],
                output_slots: vec![12],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![],
            },
        ],
        input_slots: vec![0, 1, 2, 3],
        output_slots: vec![8, 9, 11, 12],
        n_slots: 13,
    }
}

fn cpu_parity_inputs() -> Vec<Tensor> {
    vec![
        f64_tensor(vec![2, 2], vec![4.0, 1.0, 1.0, 3.0]),
        f64_tensor(vec![2, 2], vec![0.5, 0.0, 0.0, 0.5]),
        f64_tensor(vec![2], vec![1.0, 2.0]),
        f64_tensor(vec![4], vec![0.0, 0.0, 0.0, 0.0]),
    ]
}

fn assert_tensor_eq(lhs: &Tensor, rhs: &Tensor) {
    assert_eq!(lhs.shape(), rhs.shape());
    match (lhs, rhs) {
        (Tensor::F32(lhs), Tensor::F32(rhs)) => {
            assert_eq!(lhs.host_data().unwrap(), rhs.host_data().unwrap())
        }
        (Tensor::F64(lhs), Tensor::F64(rhs)) => {
            assert_eq!(lhs.host_data().unwrap(), rhs.host_data().unwrap())
        }
        (Tensor::C32(lhs), Tensor::C32(rhs)) => {
            assert_eq!(lhs.host_data().unwrap(), rhs.host_data().unwrap())
        }
        (Tensor::C64(lhs), Tensor::C64(rhs)) => {
            assert_eq!(lhs.host_data().unwrap(), rhs.host_data().unwrap())
        }
        _ => panic!("dtype mismatch: lhs={lhs:?} rhs={rhs:?}"),
    }
}

fn assert_tensor_vec_eq(lhs: &[Tensor], rhs: &[Tensor]) {
    assert_eq!(lhs.len(), rhs.len());
    for (lhs, rhs) in lhs.iter().zip(rhs.iter()) {
        assert_tensor_eq(lhs, rhs);
    }
}

#[test]
fn segmented_dispatch_matches_unsegmented_dispatch_on_cpu() {
    let program = cpu_parity_program();
    let inputs = cpu_parity_inputs();

    let mut cpu_unsegmented = CpuBackend::new();
    let mut extension_caches = ExtensionCacheStore::new();
    let mut unsegmented_context =
        ExtensionExecutionContext::new(&mut cpu_unsegmented, &mut extension_caches);
    let unsegmented = unsegmented_context
        .execute_core_exec_program_unsegmented(&program, inputs.clone())
        .unwrap();

    let mut cpu_segmented = GraphExecutor::new(CpuBackend::new());
    let segmented = cpu_segmented
        .eval_exec_ir(&program, inputs.clone())
        .unwrap();

    let mut cpu_non_consuming = GraphExecutor::new(CpuBackend::new());
    let segmented_non_consuming = cpu_non_consuming
        .eval_exec_ir_non_consuming(&program, &inputs)
        .unwrap();

    assert_tensor_vec_eq(&unsegmented, &segmented);
    assert_tensor_vec_eq(&segmented, &segmented_non_consuming);
}

fn terminal_noncompact_broadcast_multiply_program() -> ExecProgram {
    let output_shape = vec![2, 2, 4, 3];
    let output_exprs = const_shape(&output_shape);
    let output_extents = exact_extents(&output_shape);

    ExecProgram {
        instructions: vec![
            ExecInstruction {
                op: ExecOp::BroadcastInDim {
                    shape: output_exprs.clone(),
                    dims: vec![1, 0, 3],
                },
                input_slots: vec![0],
                output_slots: vec![2],
                dtype: DType::F64,
                output_shapes: vec![output_exprs.clone()].into(),
                output_extents: vec![output_extents.clone()].into(),
                last_use: vec![false],
            },
            ExecInstruction {
                op: ExecOp::BroadcastInDim {
                    shape: output_exprs.clone(),
                    dims: vec![2, 3],
                },
                input_slots: vec![1],
                output_slots: vec![3],
                dtype: DType::F64,
                output_shapes: vec![output_exprs.clone()].into(),
                output_extents: vec![output_extents.clone()].into(),
                last_use: vec![false],
            },
            ExecInstruction {
                op: ExecOp::Multiply,
                input_slots: vec![2, 3],
                output_slots: vec![4],
                dtype: DType::F64,
                output_shapes: vec![output_exprs].into(),
                output_extents: vec![output_extents].into(),
                last_use: vec![true, true],
            },
        ],
        input_slots: vec![0, 1],
        output_slots: vec![4],
        n_slots: 5,
    }
}

fn terminal_noncompact_broadcast_multiply_inputs() -> Vec<Tensor> {
    vec![
        f64_tensor(vec![2, 2, 3], (0..12).map(|idx| idx as f64 + 1.0).collect()),
        f64_tensor(vec![4, 3], (0..12).map(|idx| idx as f64 + 101.0).collect()),
    ]
}

#[test]
fn segmented_value_dispatch_preserves_terminal_lazy_broadcast_multiply_view() {
    let program = terminal_noncompact_broadcast_multiply_program();
    let inputs = terminal_noncompact_broadcast_multiply_inputs();

    let mut executor = GraphExecutor::new(CpuBackend::new());
    let mut outputs = executor
        .eval_exec_ir_non_consuming_values(&program, &inputs)
        .unwrap();
    assert_eq!(outputs.len(), 1);
    assert!(
        matches!(outputs[0], TensorValue::View(_)),
        "terminal broadcast-multiply segment should preserve a lazy output view"
    );

    let output = outputs.pop().unwrap().to_tensor().unwrap();
    assert_eq!(output.shape(), &[2, 2, 4, 3]);
    let Tensor::F64(output) = output else {
        panic!("expected f64 output")
    };
    let lhs = inputs[0].as_slice::<f64>().unwrap();
    let rhs = inputs[1].as_slice::<f64>().unwrap();
    let actual = output.host_data().unwrap();
    for t in 0..3 {
        for o in 0..4 {
            for k in 0..2 {
                for j in 0..2 {
                    let out_idx = j + 2 * k + 4 * o + 16 * t;
                    let lhs_idx = k + 2 * j + 4 * t;
                    let rhs_idx = o + 4 * t;
                    assert_eq!(actual[out_idx], lhs[lhs_idx] * rhs[rhs_idx]);
                }
            }
        }
    }
}

#[cfg(feature = "cuda")]
fn gpu_host_boundary_program() -> ExecProgram {
    ExecProgram {
        instructions: vec![
            ExecInstruction {
                op: ExecOp::DotGeneral(matmul_config()),
                input_slots: vec![0, 0],
                output_slots: vec![3],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![true],
            },
            ExecInstruction {
                op: ExecOp::ShapeOf { axis: 0 },
                input_slots: vec![3],
                output_slots: vec![4],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![true],
            },
            ExecInstruction {
                op: ExecOp::PadToMatch { axis: 0 },
                input_slots: vec![1, 2],
                output_slots: vec![5],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![true, true],
            },
            ExecInstruction {
                op: ExecOp::Exp,
                input_slots: vec![5],
                output_slots: vec![6],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![false],
            },
            ExecInstruction {
                op: ExecOp::ShapeOf { axis: 0 },
                input_slots: vec![6],
                output_slots: vec![7],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![false],
            },
            ExecInstruction {
                op: ExecOp::DynamicTruncate { axis: 0 },
                input_slots: vec![6, 7],
                output_slots: vec![8],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![true, true],
            },
            ExecInstruction {
                op: ExecOp::Negate,
                input_slots: vec![8],
                output_slots: vec![9],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![true],
            },
            ExecInstruction {
                op: ExecOp::Constant {
                    dtype: DType::F64,
                    bytes: 2.0_f64.to_le_bytes().to_vec(),
                },
                input_slots: vec![],
                output_slots: vec![10],
                dtype: DType::F64,
                output_shapes: vec![Vec::new()].into(),
                output_extents: scalar_extents(1).into(),
                last_use: vec![],
            },
        ],
        input_slots: vec![0, 1, 2],
        output_slots: vec![4, 9, 10],
        n_slots: 11,
    }
}

#[cfg(feature = "cuda")]
fn gpu_host_boundary_inputs() -> Vec<Tensor> {
    vec![
        f64_tensor(vec![2, 2], vec![5.0, 1.0, 1.0, 4.0]),
        f64_tensor(vec![2], vec![1.0, 2.0]),
        f64_tensor(vec![4], vec![0.0, 0.0, 0.0, 0.0]),
    ]
}

#[cfg(feature = "cuda")]
fn upload_all(backend: &CudaBackend, tensors: &[Tensor]) -> Vec<Tensor> {
    tensors
        .iter()
        .map(|tensor| upload_tensor(backend.runtime(), tensor).unwrap())
        .collect()
}

#[cfg(feature = "cuda")]
fn download_all(backend: &CudaBackend, tensors: &[Tensor]) -> Vec<Tensor> {
    tensors
        .iter()
        .map(|tensor| download_tensor(backend.runtime(), tensor).unwrap())
        .collect()
}

#[cfg(feature = "cuda")]
#[test]
fn segmented_dispatch_matches_unsegmented_dispatch_on_cubecl_host_boundaries() {
    let program = gpu_host_boundary_program();
    let host_inputs = gpu_host_boundary_inputs();

    let mut gpu_unsegmented = CudaBackend::new(0).unwrap();
    let unsegmented_inputs = upload_all(&gpu_unsegmented, &host_inputs);
    let mut extension_caches = ExtensionCacheStore::new();
    let mut unsegmented_context =
        ExtensionExecutionContext::new(&mut gpu_unsegmented, &mut extension_caches);
    let unsegmented = unsegmented_context
        .execute_core_exec_program_unsegmented(&program, unsegmented_inputs)
        .unwrap();
    drop(unsegmented_context);
    let unsegmented_host = download_all(&gpu_unsegmented, &unsegmented);

    let gpu_segmented_backend = CudaBackend::new(0).unwrap();
    let segmented_inputs = upload_all(&gpu_segmented_backend, &host_inputs);
    let mut gpu_segmented = GraphExecutor::new(gpu_segmented_backend);
    let segmented = gpu_segmented
        .eval_exec_ir(&program, segmented_inputs)
        .unwrap();
    let segmented_host = download_all(gpu_segmented.backend(), &segmented);

    assert_tensor_vec_eq(&unsegmented_host, &segmented_host);
}

#[allow(dead_code)]
fn _keep_complex_imports(_: Complex32, _: Complex64) {}
