#include "maths/simd.h"

#include <assert.h>
#include <stdlib.h>

#include "parallel.h"

#if defined(__x86_64__) || defined(__i386__)
#include <immintrin.h>
#elif defined(__aarch64__)
#include <arm_neon.h>
#endif

enum PsiSimdCapability psi_detect_simd(void)
{
#if defined(__x86_64__) || defined(__i386__)
    if (__builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512dq"))
        return PSI_SIMD_AVX512;
    if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma"))
        return PSI_SIMD_AVX2;

    return PSI_SIMD_NONE;
#elif defined(__aarch64__)
    return PSI_SIMD_NEON;
#else
    return PSI_SIMD_NONE;
#endif
}

const char* psi_simd_name(enum PsiSimdCapability cap)
{
    switch (cap)
    {
        case PSI_SIMD_NONE: return "Scalar";
        case PSI_SIMD_AVX2: return "AVX2+FMA";
        case PSI_SIMD_AVX512: return "AVX-512";
        case PSI_SIMD_NEON: return "NEON";
    }

    return "Scalar";
}

static void apply_pair(struct PsiComplex* state, size_t i, size_t j, struct PsiComplex g00,
                       struct PsiComplex g01, struct PsiComplex g10, struct PsiComplex g11)
{
    struct PsiComplex s0 = state[i];
    struct PsiComplex s1 = state[j];

    state[i] = psi_new_complex(s0.real * g00.real - s0.imaginary * g00.imaginary +
                                       s1.real * g01.real - s1.imaginary * g01.imaginary,
                               s0.real * g00.imaginary + s0.imaginary * g00.real +
                                       s1.real * g01.imaginary + s1.imaginary * g01.real);
    state[j] = psi_new_complex(s0.real * g10.real - s0.imaginary * g10.imaginary +
                                       s1.real * g11.real - s1.imaginary * g11.imaginary,
                               s0.real * g10.imaginary + s0.imaginary * g10.real +
                                       s1.real * g11.imaginary + s1.imaginary * g11.real);
}

static void apply_scalar(struct PsiComplex* state, const struct PsiComplex gate[2][2],
                         size_t target, size_t num_qubits)
{
    size_t target_bit = num_qubits - 1 - target;
    size_t step = (size_t)1 << target_bit;
    size_t dim = (size_t)1 << num_qubits;

    for (size_t i = 0; i < dim; i++)
    {
        if ((i >> target_bit) & 1)
            continue;

        apply_pair(state, i, i | step, gate[0][0], gate[0][1], gate[1][0], gate[1][1]);
    }
}

static size_t (*build_pairs(size_t dim, size_t target_bit, size_t step, size_t* out_count))[2]
{
    size_t (*pairs)[2] = malloc((dim / 2) * sizeof(*pairs));
    assert(pairs != NULL || dim == 0);

    size_t n = 0;
    for (size_t i = 0; i < dim; i++)
        if (((i >> target_bit) & 1) == 0)
        {
            pairs[n][0] = i;
            pairs[n][1] = i | step;
            n++;
        }

    *out_count = n;
    return pairs;
}

#if defined(__aarch64__)
static void apply_neon(struct PsiComplex* state, const struct PsiComplex gate[2][2], size_t target,
                       size_t num_qubits)
{
    size_t target_bit = num_qubits - 1 - target;
    size_t step = (size_t)1 << target_bit;
    size_t dim = (size_t)1 << num_qubits;

    struct PsiComplex g00 = gate[0][0];
    struct PsiComplex g01 = gate[0][1];
    struct PsiComplex g10 = gate[1][0];
    struct PsiComplex g11 = gate[1][1];

    size_t np;
    size_t (*pairs)[2] = build_pairs(dim, target_bit, step, &np);
    size_t chunks = np / 2;

    float64x2_t g00_re = vdupq_n_f64(g00.real);
    float64x2_t g00_im = vdupq_n_f64(g00.imaginary);
    float64x2_t g01_re = vdupq_n_f64(g01.real);
    float64x2_t g01_im = vdupq_n_f64(g01.imaginary);
    float64x2_t g10_re = vdupq_n_f64(g10.real);
    float64x2_t g10_im = vdupq_n_f64(g10.imaginary);
    float64x2_t g11_re = vdupq_n_f64(g11.real);
    float64x2_t g11_im = vdupq_n_f64(g11.imaginary);

    for (size_t chunk = 0; chunk < chunks; chunk++)
    {
        size_t i0 = pairs[chunk * 2][0];
        size_t j0 = pairs[chunk * 2][1];
        size_t i1 = pairs[chunk * 2 + 1][0];
        size_t j1 = pairs[chunk * 2 + 1][1];

        double s0r[2] = { state[i0].real, state[i1].real };
        double s0i[2] = { state[i0].imaginary, state[i1].imaginary };
        double s1r[2] = { state[j0].real, state[j1].real };
        double s1i[2] = { state[j0].imaginary, state[j1].imaginary };

        float64x2_t s0_re = vld1q_f64(s0r);
        float64x2_t s0_im = vld1q_f64(s0i);
        float64x2_t s1_re = vld1q_f64(s1r);
        float64x2_t s1_im = vld1q_f64(s1i);

        float64x2_t new0_re = vaddq_f64(vfmsq_f64(vmulq_f64(s0_re, g00_re), s0_im, g00_im),
                                        vfmsq_f64(vmulq_f64(s1_re, g01_re), s1_im, g01_im));
        float64x2_t new0_im = vaddq_f64(vfmaq_f64(vmulq_f64(s0_re, g00_im), s0_im, g00_re),
                                        vfmaq_f64(vmulq_f64(s1_re, g01_im), s1_im, g01_re));
        float64x2_t new1_re = vaddq_f64(vfmsq_f64(vmulq_f64(s0_re, g10_re), s0_im, g10_im),
                                        vfmsq_f64(vmulq_f64(s1_re, g11_re), s1_im, g11_im));
        float64x2_t new1_im = vaddq_f64(vfmaq_f64(vmulq_f64(s0_re, g10_im), s0_im, g10_re),
                                        vfmaq_f64(vmulq_f64(s1_re, g11_im), s1_im, g11_re));

        state[i0] = psi_new_complex(vgetq_lane_f64(new0_re, 0), vgetq_lane_f64(new0_im, 0));
        state[j0] = psi_new_complex(vgetq_lane_f64(new1_re, 0), vgetq_lane_f64(new1_im, 0));
        state[i1] = psi_new_complex(vgetq_lane_f64(new0_re, 1), vgetq_lane_f64(new0_im, 1));
        state[j1] = psi_new_complex(vgetq_lane_f64(new1_re, 1), vgetq_lane_f64(new1_im, 1));
    }

    for (size_t p = chunks * 2; p < np; p++)
        apply_pair(state, pairs[p][0], pairs[p][1], g00, g01, g10, g11);

    free(pairs);
}
#endif

#if defined(__x86_64__) || defined(__i386__)
__attribute__((target("avx2,fma"))) static void apply_avx2(struct PsiComplex* state,
                                                           const struct PsiComplex gate[2][2],
                                                           size_t target, size_t num_qubits)
{
    size_t target_bit = num_qubits - 1 - target;
    size_t step = (size_t)1 << target_bit;
    size_t dim = (size_t)1 << num_qubits;

    struct PsiComplex g00 = gate[0][0];
    struct PsiComplex g01 = gate[0][1];
    struct PsiComplex g10 = gate[1][0];
    struct PsiComplex g11 = gate[1][1];

    size_t np;
    size_t (*pairs)[2] = build_pairs(dim, target_bit, step, &np);
    size_t chunks = np / 2;

    for (size_t chunk = 0; chunk < chunks; chunk++)
    {
        size_t i0 = pairs[chunk * 2][0];
        size_t j0 = pairs[chunk * 2][1];
        size_t i1 = pairs[chunk * 2 + 1][0];
        size_t j1 = pairs[chunk * 2 + 1][1];

        __m256d s0_re =
                _mm256_set_pd(state[j1].real, state[i1].real, state[j0].real, state[i0].real);
        __m256d s0_im = _mm256_set_pd(state[j1].imaginary, state[i1].imaginary, state[j0].imaginary,
                                      state[i0].imaginary);

        __m256d g_re_0 = _mm256_set_pd(g01.real, g00.real, g01.real, g00.real);
        __m256d g_im_0 = _mm256_set_pd(g01.imaginary, g00.imaginary, g01.imaginary, g00.imaginary);
        __m256d g_re_1 = _mm256_set_pd(g11.real, g10.real, g11.real, g10.real);
        __m256d g_im_1 = _mm256_set_pd(g11.imaginary, g10.imaginary, g11.imaginary, g10.imaginary);

        __m256d prod0_re = _mm256_fmsub_pd(s0_re, g_re_0, _mm256_mul_pd(s0_im, g_im_0));
        __m256d prod0_im = _mm256_fmadd_pd(s0_re, g_im_0, _mm256_mul_pd(s0_im, g_re_0));
        __m256d prod1_re = _mm256_fmsub_pd(s0_re, g_re_1, _mm256_mul_pd(s0_im, g_im_1));
        __m256d prod1_im = _mm256_fmadd_pd(s0_re, g_im_1, _mm256_mul_pd(s0_im, g_re_1));

        double res0_re[4];
        double res0_im[4];
        double res1_re[4];
        double res1_im[4];

        _mm256_storeu_pd(res0_re, prod0_re);
        _mm256_storeu_pd(res0_im, prod0_im);
        _mm256_storeu_pd(res1_re, prod1_re);
        _mm256_storeu_pd(res1_im, prod1_im);

        state[i0] = psi_new_complex(res0_re[0] + res0_re[1], res0_im[0] + res0_im[1]);
        state[j0] = psi_new_complex(res1_re[0] + res1_re[1], res1_im[0] + res1_im[1]);
        state[i1] = psi_new_complex(res0_re[2] + res0_re[3], res0_im[2] + res0_im[3]);
        state[j1] = psi_new_complex(res1_re[2] + res1_re[3], res1_im[2] + res1_im[3]);
    }

    for (size_t p = chunks * 2; p < np; p++)
        apply_pair(state, pairs[p][0], pairs[p][1], g00, g01, g10, g11);

    free(pairs);
}

__attribute__((target("avx512f,avx512dq"))) static void
apply_avx512(struct PsiComplex* state, const struct PsiComplex gate[2][2], size_t target,
             size_t num_qubits)
{
    size_t target_bit = num_qubits - 1 - target;
    size_t step = (size_t)1 << target_bit;
    size_t dim = (size_t)1 << num_qubits;

    struct PsiComplex g00 = gate[0][0];
    struct PsiComplex g01 = gate[0][1];
    struct PsiComplex g10 = gate[1][0];
    struct PsiComplex g11 = gate[1][1];

    size_t np;
    size_t (*pairs)[2] = build_pairs(dim, target_bit, step, &np);
    size_t chunks = np / 4;

    for (size_t chunk = 0; chunk < chunks; chunk++)
    {
        size_t base = chunk * 4;
        size_t i0 = pairs[base][0], j0 = pairs[base][1];
        size_t i1 = pairs[base + 1][0], j1 = pairs[base + 1][1];
        size_t i2 = pairs[base + 2][0], j2 = pairs[base + 2][1];
        size_t i3 = pairs[base + 3][0], j3 = pairs[base + 3][1];

        __m512d s0_re =
                _mm512_set_pd(state[j3].real, state[i3].real, state[j2].real, state[i2].real,
                              state[j1].real, state[i1].real, state[j0].real, state[i0].real);
        __m512d s0_im = _mm512_set_pd(state[j3].imaginary, state[i3].imaginary, state[j2].imaginary,
                                      state[i2].imaginary, state[j1].imaginary, state[i1].imaginary,
                                      state[j0].imaginary, state[i0].imaginary);

        __m512d g_re_0 = _mm512_set_pd(g01.real, g00.real, g01.real, g00.real, g01.real, g00.real,
                                       g01.real, g00.real);
        __m512d g_im_0 = _mm512_set_pd(g01.imaginary, g00.imaginary, g01.imaginary, g00.imaginary,
                                       g01.imaginary, g00.imaginary, g01.imaginary, g00.imaginary);
        __m512d g_re_1 = _mm512_set_pd(g11.real, g10.real, g11.real, g10.real, g11.real, g10.real,
                                       g11.real, g10.real);
        __m512d g_im_1 = _mm512_set_pd(g11.imaginary, g10.imaginary, g11.imaginary, g10.imaginary,
                                       g11.imaginary, g10.imaginary, g11.imaginary, g10.imaginary);

        __m512d prod0_re = _mm512_fmsub_pd(s0_re, g_re_0, _mm512_mul_pd(s0_im, g_im_0));
        __m512d prod0_im = _mm512_fmadd_pd(s0_re, g_im_0, _mm512_mul_pd(s0_im, g_re_0));
        __m512d prod1_re = _mm512_fmsub_pd(s0_re, g_re_1, _mm512_mul_pd(s0_im, g_im_1));
        __m512d prod1_im = _mm512_fmadd_pd(s0_re, g_im_1, _mm512_mul_pd(s0_im, g_re_1));

        double res0_re[8];
        double res0_im[8];
        double res1_re[8];
        double res1_im[8];

        _mm512_storeu_pd(res0_re, prod0_re);
        _mm512_storeu_pd(res0_im, prod0_im);
        _mm512_storeu_pd(res1_re, prod1_re);
        _mm512_storeu_pd(res1_im, prod1_im);

        state[i0] = psi_new_complex(res0_re[0] + res0_re[1], res0_im[0] + res0_im[1]);
        state[j0] = psi_new_complex(res1_re[0] + res1_re[1], res1_im[0] + res1_im[1]);
        state[i1] = psi_new_complex(res0_re[2] + res0_re[3], res0_im[2] + res0_im[3]);
        state[j1] = psi_new_complex(res1_re[2] + res1_re[3], res1_im[2] + res1_im[3]);
        state[i2] = psi_new_complex(res0_re[4] + res0_re[5], res0_im[4] + res0_im[5]);
        state[j2] = psi_new_complex(res1_re[4] + res1_re[5], res1_im[4] + res1_im[5]);
        state[i3] = psi_new_complex(res0_re[6] + res0_re[7], res0_im[6] + res0_im[7]);
        state[j3] = psi_new_complex(res1_re[6] + res1_re[7], res1_im[6] + res1_im[7]);
    }

    for (size_t p = chunks * 4; p < np; p++)
        apply_pair(state, pairs[p][0], pairs[p][1], g00, g01, g10, g11);

    free(pairs);
}
#endif

void psi_apply_single_qubit_gate_simd(struct PsiComplex* state, const struct PsiComplex gate[2][2],
                                      size_t target, size_t num_qubits)
{
    enum PsiSimdCapability cap = psi_detect_simd();

#if defined(__x86_64__) || defined(__i386__)
    if (cap == PSI_SIMD_AVX512)
    {
        apply_avx512(state, gate, target, num_qubits);
        return;
    }
    if (cap == PSI_SIMD_AVX2)
    {
        apply_avx2(state, gate, target, num_qubits);
        return;
    }
#elif defined(__aarch64__)
    if (cap == PSI_SIMD_NEON)
    {
        apply_neon(state, gate, target, num_qubits);
        return;
    }
#endif

    (void)cap;
    apply_scalar(state, gate, target, num_qubits);
}

struct PsiSimdPairs
{
    struct PsiComplex* state;
    size_t (*pairs)[2];
    struct PsiComplex g00;
    struct PsiComplex g01;
    struct PsiComplex g10;
    struct PsiComplex g11;
};

static void simd_pair_range(size_t start, size_t end, void* vctx)
{
    struct PsiSimdPairs* c = vctx;

    for (size_t p = start; p < end; p++)
        apply_pair(c->state, c->pairs[p][0], c->pairs[p][1], c->g00, c->g01, c->g10, c->g11);
}

void psi_apply_single_qubit_gate_simd_parallel(struct PsiComplex* state,
                                               const struct PsiComplex gate[2][2], size_t target,
                                               size_t num_qubits)
{
    size_t target_bit = num_qubits - 1 - target;
    size_t step = (size_t)1 << target_bit;
    size_t dim = (size_t)1 << num_qubits;

    size_t np;
    size_t (*pairs)[2] = build_pairs(dim, target_bit, step, &np);

    struct PsiSimdPairs ctx = {
        state, pairs, gate[0][0], gate[0][1], gate[1][0], gate[1][1],
    };
    psi_parallel_for(np, simd_pair_range, &ctx);

    free(pairs);
}