#include "tests.h"

static const double R2 = 0.7071067811865476;

void run_custom_tests(void)
{
    psi_test_section("Custom gates");

    struct PsiCompositeGateOp bell_ops[] = {
        { PSI_OP_H, { 0 }, 1 },
        { PSI_OP_CNOT, { 0, 1 }, 2 },
    };
    struct PsiCustomGate bell = psi_new_custom_gate_from_composite("BELL", 2, bell_ops, 2);
    struct PsiQuantumCircuit c = psi_new_quantum_circuit(2);
    size_t bell_targets[] = { 0, 1 };
    psi_apply_custom(&c, bell, bell_targets, 2);
    struct PsiComplex bell_exp[] = {
        psi_new_complex(R2, 0.0),
        psi_new_complex(0.0, 0.0),
        psi_new_complex(0.0, 0.0),
        psi_new_complex(R2, 0.0),
    };
    psi_check_circuit("composite BELL gate -> Bell state", &c, bell_exp, 4);
    psi_free_quantum_circuit(&c);

    struct PsiCustomGate custom_x = psi_new_custom_gate_from_matrix(
            "MYX",
            psi_matrix(2, 2, psi_new_complex(0.0, 0.0), psi_new_complex(1.0, 0.0),
                       psi_new_complex(1.0, 0.0), psi_new_complex(0.0, 0.0)));
    struct PsiQuantumCircuit c2 = psi_new_quantum_circuit(1);
    size_t x_targets[] = { 0 };
    psi_apply_custom(&c2, custom_x, x_targets, 1);
    struct PsiComplex x_exp[] = { psi_new_complex(0.0, 0.0), psi_new_complex(1.0, 0.0) };
    psi_check_circuit("from_matrix X custom gate: |0> -> |1>", &c2, x_exp, 2);
    psi_free_quantum_circuit(&c2);
}