#include "core/custom_gate.h"

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

#include "core/gates.h"

struct PsiCustomGate psi_new_custom_gate_from_matrix(const char* name, struct PsiMatrix matrix)
{
    assert(matrix.rows == matrix.cols);

    size_t dim = matrix.rows;
    assert(dim > 0 && (dim & (dim - 1)) == 0);

    size_t num_qubits = 0;
    while (((size_t)1 << num_qubits) < dim)
        num_qubits++;

    struct PsiCustomGate gate;
    gate.name = name;
    gate.num_qubits = num_qubits;
    gate.kind = PSI_CUSTOM_GATE_MATRIX;
    gate.definition.matrix = matrix;

    return gate;
}

struct PsiCustomGate psi_new_custom_gate_from_composite(const char* name, size_t num_qubits,
                                                        const struct PsiCompositeGateOp* ops,
                                                        size_t op_count)
{
    struct PsiCompositeGateOp* owned = malloc(op_count * sizeof(struct PsiCompositeGateOp));
    assert(owned != NULL || op_count == 0);

    if (op_count > 0)
        memcpy(owned, ops, op_count * sizeof(struct PsiCompositeGateOp));

    struct PsiCustomGate gate;
    gate.name = name;
    gate.num_qubits = num_qubits;
    gate.kind = PSI_CUSTOM_GATE_COMPOSITE;
    gate.definition.composite.ops = owned;
    gate.definition.composite.count = op_count;

    return gate;
}

void psi_free_custom_gate(struct PsiCustomGate* gate)
{
    if (gate->kind == PSI_CUSTOM_GATE_MATRIX)
    {
        psi_free_matrix(&gate->definition.matrix);
        return;
    }

    free(gate->definition.composite.ops);
    gate->definition.composite.ops = NULL;
    gate->definition.composite.count = 0;
}

static struct PsiQuantumGate op_gate(enum PsiCompositeOp op)
{
    switch (op)
    {
        case PSI_OP_H: return psi_hadamard_gate();
        case PSI_OP_X: return psi_pauli_x_gate();
        case PSI_OP_Y: return psi_pauli_y_gate();
        case PSI_OP_Z: return psi_pauli_z_gate();
        case PSI_OP_S: return psi_s_gate();
        case PSI_OP_T: return psi_t_gate();
        case PSI_OP_CNOT: return psi_cnot_gate();
        case PSI_OP_CZ: return psi_cz_gate();
        case PSI_OP_SWAP: return psi_swap_gate();
        case PSI_OP_CCNOT: return psi_toffoli_gate();
        case PSI_OP_CSWAP: return psi_fredkin_gate();
    }

    return psi_identity_gate();
}

static bool find_target(const size_t* targets, size_t count, size_t q, size_t* pos)
{
    for (size_t i = 0; i < count; i++)
        if (targets[i] == q)
        {
            *pos = i;
            return true;
        }

    return false;
}

static struct PsiMatrix build_full_operator(struct PsiMatrix gate_matrix, const size_t* targets,
                                            size_t num_gate_qubits, size_t total_qubits)
{
    size_t dim = (size_t)1 << total_qubits;
    size_t gate_dim = gate_matrix.rows;

    struct PsiMatrix result = psi_new_matrix(dim, dim);

    for (size_t i = 0; i < dim; i++)
        for (size_t j = 0; j < dim; j++)
        {
            size_t gate_i = 0;
            size_t gate_j = 0;
            bool match_non_targets = true;

            for (size_t q = 0; q < total_qubits; q++)
            {
                size_t bit_i = (i >> (total_qubits - 1 - q)) & 1;
                size_t bit_j = (j >> (total_qubits - 1 - q)) & 1;

                size_t pos;
                if (find_target(targets, num_gate_qubits, q, &pos))
                {
                    gate_i |= bit_i << (num_gate_qubits - 1 - pos);
                    gate_j |= bit_j << (num_gate_qubits - 1 - pos);
                }
                else if (bit_i != bit_j)
                {
                    match_non_targets = false;
                    break;
                }
            }

            if (match_non_targets)
                result.data[i * dim + j] = gate_matrix.data[gate_i * gate_dim + gate_j];
        }

    return result;
}

static struct PsiMatrix compute_composite_matrix(struct PsiCustomGate gate)
{
    size_t dim = (size_t)1 << gate.num_qubits;
    struct PsiMatrix result = psi_identity_matrix(dim);

    for (size_t i = 0; i < gate.definition.composite.count; i++)
    {
        struct PsiCompositeGateOp step = gate.definition.composite.ops[i];
        struct PsiQuantumGate g = op_gate(step.op);

        struct PsiMatrix full =
                build_full_operator(g.matrix, step.targets, step.target_count, gate.num_qubits);
        struct PsiMatrix next = psi_dot_matrix(full, result);

        psi_free_matrix(&full);
        psi_free_matrix(&result);
        psi_free_quantum_gate(&g);
        result = next;
    }

    return result;
}

struct PsiQuantumGate psi_to_quantum_gate(struct PsiCustomGate gate)
{
    if (gate.kind == PSI_CUSTOM_GATE_MATRIX)
        return psi_new_quantum_gate(gate.name, psi_clone_matrix(gate.definition.matrix),
                                    gate.num_qubits);

    return psi_new_quantum_gate(gate.name, compute_composite_matrix(gate), gate.num_qubits);
}