#include "tests.h"

static const double R2 = 0.7071067811865476;

static void build_fusion(struct PsiQuantumCircuit* c)
{
    psi_apply_h(c, 0);
    psi_apply_z(c, 0);
    psi_apply_h(c, 0);
    psi_apply_x(c, 1);
    psi_apply_x(c, 1);
}

void run_kernel_tests(void)
{
    psi_test_section("Kernel batching and fusion");

    struct PsiQuantumCircuit hh = psi_new_quantum_circuit(1);
    psi_apply_h(&hh, 0);
    psi_apply_h(&hh, 0);
    struct PsiComplex hh_exp[] = { psi_new_complex(1.0, 0.0), psi_new_complex(0.0, 0.0) };
    psi_check_circuit("H then H = identity", &hh, hh_exp, 2);
    psi_free_quantum_circuit(&hh);

    struct PsiKernel a =
            psi_new_kernel("H",
                           psi_matrix(2, 2, psi_new_complex(R2, 0.0), psi_new_complex(R2, 0.0),
                                      psi_new_complex(R2, 0.0), psi_new_complex(-R2, 0.0)),
                           (size_t[]){ 0 }, 1);
    struct PsiKernel b = psi_clone_kernel(a);
    psi_test_check(psi_kernels_can_fuse(a, b), "adjacent single-qubit kernels fuse");
    psi_free_kernel(&a);
    psi_free_kernel(&b);

    struct PsiKernelBatch batch = psi_new_kernel_batch(1);
    psi_add_kernel(
            &batch,
            psi_new_kernel("H",
                           psi_matrix(2, 2, psi_new_complex(R2, 0.0), psi_new_complex(R2, 0.0),
                                      psi_new_complex(R2, 0.0), psi_new_complex(-R2, 0.0)),
                           (size_t[]){ 0 }, 1));
    psi_add_kernel(
            &batch,
            psi_new_kernel("H",
                           psi_matrix(2, 2, psi_new_complex(R2, 0.0), psi_new_complex(R2, 0.0),
                                      psi_new_complex(R2, 0.0), psi_new_complex(-R2, 0.0)),
                           (size_t[]){ 0 }, 1));
    psi_optimize_kernel_batch(&batch);
    psi_test_check(batch.count == 1, "batch fuses two H kernels into one");
    psi_free_kernel_batch(&batch);

    psi_check_runtimes_agree("Runtimes agree on fusion circuit", 2, build_fusion);
}