typedef float2 complex_f;

/*
 * Addition of two complex numbers:
 *
 * a + b = (Re(a) + Re(b)) + i(Im(a) + Im(b))
 */
static complex_f add(complex_f a, complex_f b)
{
    return (complex_f)(a.x + b.x, a.y + b.y);
}

/*
 * Negation of a complex numbers:
 *
 * -a = -(Re(a) - i(Im(a))
 */
static complex_f neg(complex_f a)
{
    return (complex_f)(-a.x, -a.y);
}

/*
 * Multiplication of two complex numbers:
 *
 * a * b =
 *   ((Re(a) * Re(b)) - (Im(a) * Im(b)))
 * + ((Im(a) * Re(b)) + (Re(a) * Im(b)))i
 */
static complex_f mul(complex_f a, complex_f b)
{
    return (complex_f)(
        (a.x * b.x) - (a.y * b.y),
        (a.y * b.x) + (a.x * b.y));
}
/**
 * Absolute value of a complex number
 *
 * |a| = √(Re(a)^2 + Im(a)^2)
 */
static float complex_abs(complex_f a)
{
    return sqrt((a.x * a.x) + (a.y * a.y));
}

static complex_f cexp(float a) {
    return (complex_f)(cos(a), sin(a));
}

/*
 * Returns the nth number where a given digit
 * is cleared in the binary representation of the number
 */
static uint nth_cleared(uint n, uint target)
{
    uint mask = (1 << target) - 1;
    uint not_mask = ~mask;

    return (n & mask) | ((n & not_mask) << 1);
}

///////////////////////////////////////////////
// KERNELS
///////////////////////////////////////////////

/*
 * Applies a single qubit gate to the register.
 * The gate matrix must be given in the form:
 *
 *  A B
 *  C D
 */
__kernel void apply_gate(
    __global complex_f *amplitudes,
    uint target,
    complex_f A,
    complex_f B,
    complex_f C,
    complex_f D)
{
    uint const global_id = get_global_id(0);

    uint const zero_state = nth_cleared(global_id, target);

    // uint const zero_state = state & (~(1 << target)); // Could just be state
    uint const one_state = zero_state | (1 << target);

    complex_f const zero_amp = amplitudes[zero_state];
    complex_f const one_amp = amplitudes[one_state];

    amplitudes[zero_state] = add(mul(A, zero_amp), mul(B, one_amp));
    amplitudes[one_state] = add(mul(D, one_amp), mul(C, zero_amp));
}

/*
 * Applies a controlled single qubit gate to the register.
 */
__kernel void apply_controlled_gate(
    __global complex_f *amplitudes,
    uint control,
    uint target,
    complex_f A,
    complex_f B,
    complex_f C,
    complex_f D)
{
    uint const global_id = get_global_id(0);
    uint const zero_state = nth_cleared(global_id, target);
    uint const one_state = zero_state | (1 << target); // Set the target bit

    uint const control_val_zero = (((1 << control) & zero_state) > 0) ? 1 : 0;
    uint const control_val_one = (((1 << control) & one_state) > 0) ? 1 : 0;

    complex_f const zero_amp = amplitudes[zero_state];
    complex_f const one_amp = amplitudes[one_state];

    if (control_val_zero == 1)
    {
        amplitudes[zero_state] = add(mul(A, zero_amp), mul(B, one_amp));
    }

    if (control_val_one == 1)
    {
        amplitudes[one_state] = add(mul(D, one_amp), mul(C, zero_amp));
    }


}

/*
 * Applies a controlled-controlled single qubit gate to the register.
 * NOT MIGRATED
 */
__kernel void apply_controlled_controlled_gate(
    __global complex_f *const amplitudes,
    __global complex_f *amps,
    uint control1,
    uint control2,
    uint target,
    complex_f A,
    complex_f B,
    complex_f C,
    complex_f D)
{
    uint const state = get_global_id(0);
    complex_f const amp = amplitudes[state];

    uint const zero_state = state & (~(1 << target));
    uint const one_state = state | (1 << target);

    uint const bit_val = (((1 << target) & state) > 0) ? 1 : 0;
    uint const control1_val = (((1 << control1) & state) > 0) ? 1 : 0;
    uint const control2_val = (((1 << control2) & state) > 0) ? 1 : 0;

    if (control1_val == 0 || control2_val == 0)
    {
        // Control is 0, don't apply gate
        amps[state] = amp;
    }
    else
    {
        // control is 1, apply gate.
        if (bit_val == 0)
        {
            // Bitval = 0
            amps[state] = add(mul(A, amp), mul(B, amplitudes[one_state]));
        }
        else
        {
            amps[state] = add(mul(D, amp), mul(C, amplitudes[zero_state]));
        }
    }
}

/*
 * Swaps the states of two qubits in the register
 * NOT MIGRATED
 */
__kernel void swap(
    __global complex_f *const amplitudes,
    __global complex_f *amps,
    uint first_qubit,
    uint second_qubit)
{
    uint const state = get_global_id(0);

    uint const first_bit_mask = 1 << first_qubit;
    uint const second_bit_mask = 1 << second_qubit;

    uint const new_second_bit = ((state & first_bit_mask) >> first_qubit) << second_qubit;
    uint const new_first_bit = ((state & second_bit_mask) >> second_qubit) << first_qubit;

    uint const new_state = (state & !first_bit_mask & !second_bit_mask) | new_first_bit | new_second_bit;

    amps[new_state] = amplitudes[state];
}


/**
 * Calculates The Probabilities Of A State Vector
 */
__kernel void calculate_probabilities(
    __global complex_f *const amplitudes,
    __global float *probabilities)
{
    uint const state = get_global_id(0);
    complex_f amp = amplitudes[state];

    probabilities[state] = complex_abs(mul(amp, amp));
}

/**
 * Initializes a register to the value 1|0..100...0>
 *                                          ^ target
 */
__kernel void initialize_register(
    __global complex_f *amplitudes,
    uint const target)
{
    uint const state = get_global_id(0);
    if (state == target)
    {
        amplitudes[state] = (complex_f)(1, 0);
    }
    else
    {
        amplitudes[state] = (complex_f)(0, 0);
    }
}