
#include <gpusim/memory_ops.h>
#include <gpusim/update_ops_cuda.h>

#include "test_util.hpp"

void test_single_qubit_named_gate(UINT n, std::string name,
    std::function<void(UINT, void*, ITYPE, void*, UINT)> func,
    Eigen::MatrixXcd mat) {
    const ITYPE dim = 1ULL << n;
    const UINT max_repeat = 2;
    int ngpus = get_num_device();
    for (int idx = 0; idx < ngpus; ++idx) {
        set_device(idx);
        auto stream_ptr = allocate_cuda_stream_host(1, idx);
        auto state = allocate_quantum_state_host(dim, idx);
        initialize_Haar_random_state_with_seed_host(
            state, dim, 0, stream_ptr, idx);

        Eigen::VectorXcd test_state =
            copy_cpu_from_gpu(state, dim, stream_ptr, idx);
        std::vector<UINT> indices;
        for (UINT i = 0; i < n; ++i) indices.push_back(i);

        std::random_device seed_gen;
        std::mt19937 engine(seed_gen());
        for (UINT rep = 0; rep < max_repeat; ++rep) {
            for (UINT i = 0; i < n; ++i) {
                UINT target = indices[i];
                func(target, state, dim, stream_ptr, idx);
                test_state =
                    get_expanded_eigen_matrix_with_identity(target, mat, n) *
                    test_state;
                state_equal_gpu(state, test_state, dim, name, stream_ptr, idx);
            }
            std::shuffle(indices.begin(), indices.end(), engine);
        }
        release_quantum_state_host(state, idx);
        release_cuda_stream_host(stream_ptr, 1, idx);
    }
}

TEST(UpdateTest, XGate) {
    Eigen::MatrixXcd mat(2, 2);
    mat << 0, 1, 1, 0;
    void (*func)(UINT, void*, ITYPE, void*, UINT) = &X_gate_host;
    test_single_qubit_named_gate(6, "XGate", func, mat);
}
TEST(UpdateTest, YGate) {
    Eigen::MatrixXcd mat(2, 2);
    mat << 0, -1.i, 1.i, 0;
    void (*func)(UINT, void*, ITYPE, void*, UINT) = &Y_gate_host;
    test_single_qubit_named_gate(6, "YGate", func, mat);
}
TEST(UpdateTest, ZGate) {
    const UINT n = 3;
    Eigen::MatrixXcd mat(2, 2);
    mat << 1, 0, 0, -1;
    void (*func)(UINT, void*, ITYPE, void*, UINT) = &Z_gate_host;
    test_single_qubit_named_gate(6, "ZGate", func, mat);
}
TEST(UpdateTest, HGate) {
    const UINT n = 3;
    Eigen::MatrixXcd mat(2, 2);
    mat << 1, 1, 1, -1;
    mat /= sqrt(2.);
    void (*func)(UINT, void*, ITYPE, void*, UINT) = &H_gate_host;
    test_single_qubit_named_gate(n, "HGate", func, mat);
}

TEST(UpdateTest, SGate) {
    const UINT n = 3;
    Eigen::MatrixXcd mat(2, 2);
    mat << 1, 0, 0, 1.i;
    void (*func)(UINT, void*, ITYPE, void*, UINT) = &S_gate_host;
    test_single_qubit_named_gate(n, "SGate", func, mat);
    void (*func2)(UINT, void*, ITYPE, void*, UINT) = &Sdag_gate_host;
    test_single_qubit_named_gate(n, "SGate", func2, mat.adjoint());
}

TEST(UpdateTest, TGate) {
    const UINT n = 3;
    Eigen::MatrixXcd mat(2, 2);
    mat << 1, 0, 0, (1. + 1.i) / sqrt(2.);
    void (*func)(UINT, void*, ITYPE, void*, UINT) = &T_gate_host;
    test_single_qubit_named_gate(n, "TGate", func, mat);
    void (*func2)(UINT, void*, ITYPE, void*, UINT) = &Tdag_gate_host;
    test_single_qubit_named_gate(n, "TGate", func2, mat.adjoint());
}

TEST(UpdateTest, sqrtXGate) {
    const UINT n = 3;
    Eigen::MatrixXcd mat(2, 2);
    mat << 0.5 + 0.5i, 0.5 - 0.5i, 0.5 - 0.5i, 0.5 + 0.5i;
    void (*func)(UINT, void*, ITYPE, void*, UINT) = &sqrtX_gate_host;
    test_single_qubit_named_gate(n, "SqrtXGate", func, mat);
    void (*func2)(UINT, void*, ITYPE, void*, UINT) = &sqrtXdag_gate_host;
    test_single_qubit_named_gate(n, "SqrtXGate", func2, mat.adjoint());
}

TEST(UpdateTest, sqrtYGate) {
    const UINT n = 3;
    Eigen::MatrixXcd mat(2, 2);
    mat << 0.5 + 0.5i, -0.5 - 0.5i, 0.5 + 0.5i, 0.5 + 0.5i;
    void (*func)(UINT, void*, ITYPE, void*, UINT) = &sqrtY_gate_host;
    test_single_qubit_named_gate(n, "SqrtYGate", func, mat);
    void (*func2)(UINT, void*, ITYPE, void*, UINT) = &sqrtYdag_gate_host;
    test_single_qubit_named_gate(n, "SqrtYGate", func2, mat.adjoint());
}

void test_projection_gate(
    std::function<void(UINT, void*, ITYPE, void*, UINT)> func,
    std::function<double(UINT, void*, ITYPE, void*, UINT)> prob_func,
    Eigen::MatrixXcd mat) {
    const UINT n = 6;
    const ITYPE dim = 1ULL << n;
    const UINT max_repeat = 10;

    UINT target;
    double prob;
    int ngpus = get_num_device();
    for (int idx = 0; idx < ngpus; ++idx) {
        set_device(idx);
        auto stream_ptr = allocate_cuda_stream_host(1, idx);
        auto state = allocate_quantum_state_host(dim, idx);

        for (UINT rep = 0; rep < max_repeat; ++rep) {
            initialize_Haar_random_state_host(state, dim, stream_ptr, idx);
            Eigen::VectorXcd test_state =
                copy_cpu_from_gpu(state, dim, stream_ptr, idx);

            target = rand_int(n);
            // Z-projection operators
            prob = prob_func(target, state, dim, stream_ptr, idx);
            EXPECT_GT(prob, 1e-10);
            func(target, state, dim, stream_ptr, idx);
            ASSERT_NEAR(state_norm_squared_host(state, dim, stream_ptr, idx),
                prob, eps);
            normalize_host(prob, state, dim, stream_ptr, idx);

            test_state =
                get_expanded_eigen_matrix_with_identity(target, mat, n) *
                test_state;
            ASSERT_NEAR(test_state.squaredNorm(), prob, eps);
            test_state.normalize();
            state_equal_gpu(
                state, test_state, dim, "Projection gate", stream_ptr, idx);
        }
        release_quantum_state_host(state, idx);
        release_cuda_stream_host(stream_ptr, 1, idx);
    }
}

TEST(UpdateTest, ProjectionAndNormalizeTest) {
    Eigen::MatrixXcd P0(2, 2), P1(2, 2);
    P0 << 1, 0, 0, 0;
    P1 << 0, 0, 0, 1;
    void (*func1)(UINT, void*, ITYPE, void*, UINT) = &P0_gate_host;
    double (*func2)(UINT, void*, ITYPE, void*, UINT) = &M0_prob_host;
    test_projection_gate(func1, func2, P0);
    func1 = &P1_gate_host;
    func2 = &M1_prob_host;
    test_projection_gate(func1, func2, P1);
}

TEST(UpdateTest, SingleQubitRotationGateTest) {
    const UINT n = 6;
    const ITYPE dim = 1ULL << n;
    const UINT max_repeat = 10;

    Eigen::MatrixXcd Identity(2, 2), X(2, 2), Y(2, 2), Z(2, 2);
    Identity << 1, 0, 0, 1;
    X << 0, 1, 1, 0;
    Y << 0, -1.i, 1.i, 0;
    Z << 1, 0, 0, -1;
    std::complex<double> imag_unit(0, 1);

    UINT target;
    double angle;

    int ngpus = get_num_device();
    for (int idx = 0; idx < ngpus; ++idx) {
        set_device(idx);
        auto stream_ptr = allocate_cuda_stream_host(1, idx);
        auto state = allocate_quantum_state_host(dim, idx);
        initialize_Haar_random_state_host(state, dim, stream_ptr, idx);
        Eigen::VectorXcd test_state =
            copy_cpu_from_gpu(state, dim, stream_ptr, idx);
        using testset = std::tuple<
            std::function<void(UINT, double, void*, ITYPE, void*, UINT)>,
            Eigen::MatrixXcd, std::string>;
        std::vector<testset> test_list;
        void (*func)(UINT, double, void*, ITYPE, void*, UINT) = &RX_gate_host;
        test_list.push_back(std::make_tuple(func, X, "Xrot"));
        func = &RY_gate_host;
        test_list.push_back(std::make_tuple(func, Y, "Yrot"));
        func = &RZ_gate_host;
        test_list.push_back(std::make_tuple(func, Z, "Zrot"));

        for (UINT rep = 0; rep < max_repeat; ++rep) {
            for (auto tup : test_list) {
                target = rand_int(n);
                angle = rand_real();
                auto func = std::get<0>(tup);
                auto mat = std::get<1>(tup);
                auto name = std::get<2>(tup);
                func(target, angle, state, dim, stream_ptr, idx);
                test_state = get_expanded_eigen_matrix_with_identity(target,
                                 cos(angle / 2) * Identity +
                                     imag_unit * sin(angle / 2) * mat,
                                 n) *
                             test_state;
                state_equal_gpu(state, test_state, dim, name, stream_ptr, idx);
            }
        }
        release_quantum_state_host(state, idx);
        release_cuda_stream_host(stream_ptr, 1, idx);
    }
}

void test_two_qubit_named_gate(UINT n, std::string name,
    std::function<void(UINT, UINT, void*, ITYPE, void*, UINT)> func,
    std::function<Eigen::MatrixXcd(UINT, UINT, UINT)> matfunc) {
    const ITYPE dim = 1ULL << n;
    const UINT max_repeat = 2;

    int ngpus = get_num_device();
    std::random_device seed_gen;
    std::mt19937 engine(seed_gen());
    for (int idx = 0; idx < ngpus; ++idx) {
        set_device(idx);
        auto stream_ptr = allocate_cuda_stream_host(1, idx);
        auto state = allocate_quantum_state_host(dim, idx);
        initialize_Haar_random_state_with_seed_host(
            state, dim, 0, stream_ptr, idx);

        Eigen::VectorXcd test_state =
            copy_cpu_from_gpu(state, dim, stream_ptr, idx);
        std::vector<UINT> indices;
        for (UINT i = 0; i < n; ++i) indices.push_back(i);

        for (UINT rep = 0; rep < max_repeat; ++rep) {
            for (UINT i = 0; i + 1 < n; i += 2) {
                UINT target = indices[i];
                UINT control = indices[i + 1];
                func(control, target, state, dim, stream_ptr, idx);
                Eigen::MatrixXcd mat = matfunc(control, target, n);
                test_state = mat * test_state;
                state_equal_gpu(state, test_state, dim, name, stream_ptr, idx);
            }
            std::shuffle(indices.begin(), indices.end(), engine);
        }
        release_quantum_state_host(state, idx);
        release_cuda_stream_host(stream_ptr, 1, idx);
    }
}

TEST(UpdateTest, CNOTGate) {
    const UINT n = 4;
    void (*func)(UINT, UINT, void*, ITYPE, void*, UINT) = &CNOT_gate_host;
    test_two_qubit_named_gate(
        n, "CNOT", func, get_eigen_matrix_full_qubit_CNOT);
}

TEST(UpdateTest, CZGate) {
    const UINT n = 4;
    void (*func)(UINT, UINT, void*, ITYPE, void*, UINT) = &CZ_gate_host;
    test_two_qubit_named_gate(n, "CZ", func, get_eigen_matrix_full_qubit_CZ);
}

TEST(UpdateTest, SWAPGate) {
    const UINT n = 4;
    void (*func)(UINT, UINT, void*, ITYPE, void*, UINT) = &SWAP_gate_host;
    test_two_qubit_named_gate(
        n, "SWAP", func, get_eigen_matrix_full_qubit_SWAP);
}
