#include "core/circuit.h"

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

#include "maths/format.h"

const char* psi_gate_op_name(struct PsiGateOp op)
{
    switch (op.kind)
    {
        case PSI_GATE_H: return "H";
        case PSI_GATE_X: return "X";
        case PSI_GATE_Y: return "Y";
        case PSI_GATE_Z: return "Z";
        case PSI_GATE_S: return "S";
        case PSI_GATE_T: return "T";
        case PSI_GATE_SDG: return "S†";
        case PSI_GATE_TDG: return "T†";
        case PSI_GATE_SX: return "√X";
        case PSI_GATE_SXDG: return "√X†";
        case PSI_GATE_RX: return "Rx";
        case PSI_GATE_RY: return "Ry";
        case PSI_GATE_RZ: return "Rz";
        case PSI_GATE_P: return "P";
        case PSI_GATE_U1: return "U1";
        case PSI_GATE_U2: return "U2";
        case PSI_GATE_U3: return "U3";
        case PSI_GATE_CNOT: return "CNOT";
        case PSI_GATE_CZ: return "CZ";
        case PSI_GATE_SWAP: return "SWAP";
        case PSI_GATE_CRX: return "CRx";
        case PSI_GATE_CRY: return "CRy";
        case PSI_GATE_CRZ: return "CRz";
        case PSI_GATE_CP: return "CP";
        case PSI_GATE_CCNOT: return "CCNOT";
        case PSI_GATE_CSWAP: return "CSWAP";
        case PSI_GATE_MEASURE: return "M";
        case PSI_GATE_CUSTOM: return op.custom->name;
    }

    return "?";
}

const size_t* psi_gate_op_quantum_targets(const struct PsiGateOp* op, size_t* out_count)
{
    if (op->kind == PSI_GATE_CUSTOM)
    {
        *out_count = op->custom_target_count;
        return op->custom_targets;
    }

    *out_count = op->qubit_count;
    return op->qubits;
}

const size_t* psi_gate_op_classical_targets(const struct PsiGateOp* op, size_t* out_count)
{
    if (op->kind == PSI_GATE_MEASURE)
    {
        *out_count = 1;
        return &op->classical;
    }

    *out_count = 0;
    return NULL;
}

bool psi_gate_op_is_measurement(struct PsiGateOp op)
{
    return op.kind == PSI_GATE_MEASURE;
}

bool psi_gate_op_is_custom(struct PsiGateOp op)
{
    return op.kind == PSI_GATE_CUSTOM;
}

bool psi_gate_op_is_non_clifford(struct PsiGateOp op)
{
    switch (op.kind)
    {
        case PSI_GATE_T:
        case PSI_GATE_TDG:
        case PSI_GATE_SX:
        case PSI_GATE_SXDG:
        case PSI_GATE_RX:
        case PSI_GATE_RY:
        case PSI_GATE_RZ:
        case PSI_GATE_P:
        case PSI_GATE_U1:
        case PSI_GATE_U2:
        case PSI_GATE_U3:
        case PSI_GATE_CRX:
        case PSI_GATE_CRY:
        case PSI_GATE_CRZ:
        case PSI_GATE_CP: return true;
        default: return false;
    }
}

struct PsiQuantumCircuit psi_new_quantum_circuit(size_t num_qubits)
{
    return psi_new_quantum_circuit_with_classical(num_qubits, 0);
}

struct PsiQuantumCircuit psi_new_quantum_circuit_with_classical(size_t num_qubits,
                                                                size_t num_classical)
{
    struct PsiQuantumCircuit c;
    c.num_qubits = num_qubits;
    c.num_classical = num_classical;
    c.operations = NULL;
    c.operation_count = 0;
    c.operation_capacity = 0;
    c.computed_state = psi_new_vector(0, PSI_COLUMN_VECTOR);
    c.is_computed = false;

    return c;
}

static void free_operations(struct PsiQuantumCircuit* c)
{
    for (size_t i = 0; i < c->operation_count; i++)
    {
        struct PsiGateOp* op = &c->operations[i];
        if (op->kind != PSI_GATE_CUSTOM)
            continue;

        psi_free_custom_gate(op->custom);
        free(op->custom);
        free(op->custom_targets);
    }
}

void psi_free_quantum_circuit(struct PsiQuantumCircuit* c)
{
    free_operations(c);
    free(c->operations);
    c->operations = NULL;
    c->operation_count = 0;
    c->operation_capacity = 0;
    psi_free_vector(&c->computed_state);
    c->is_computed = false;
}

void psi_reset_circuit(struct PsiQuantumCircuit* c)
{
    free_operations(c);
    c->operation_count = 0;
    psi_free_vector(&c->computed_state);
    c->computed_state = psi_new_vector(0, PSI_COLUMN_VECTOR);
    c->is_computed = false;
}

static struct PsiGateOp* append_op(struct PsiQuantumCircuit* c)
{
    if (c->operation_count == c->operation_capacity)
    {
        size_t new_capacity = c->operation_capacity == 0 ? 8 : c->operation_capacity * 2;
        c->operations = realloc(c->operations, new_capacity * sizeof(struct PsiGateOp));
        assert(c->operations != NULL);
        c->operation_capacity = new_capacity;
    }

    struct PsiGateOp* op = &c->operations[c->operation_count++];
    memset(op, 0, sizeof(*op));
    c->is_computed = false;

    return op;
}

static void push_1q(struct PsiQuantumCircuit* c, enum PsiGateKind kind, size_t target)
{
    struct PsiGateOp* op = append_op(c);
    op->kind = kind;
    op->qubits[0] = target;
    op->qubit_count = 1;
}

static void push_1q_1p(struct PsiQuantumCircuit* c, enum PsiGateKind kind, size_t target,
                       double theta)
{
    struct PsiGateOp* op = append_op(c);
    op->kind = kind;
    op->qubits[0] = target;
    op->qubit_count = 1;
    op->params[0] = theta;
    op->param_count = 1;
}

static void push_2q(struct PsiQuantumCircuit* c, enum PsiGateKind kind, size_t a, size_t b)
{
    struct PsiGateOp* op = append_op(c);
    op->kind = kind;
    op->qubits[0] = a;
    op->qubits[1] = b;
    op->qubit_count = 2;
}

static void push_2q_1p(struct PsiQuantumCircuit* c, enum PsiGateKind kind, size_t control,
                       size_t target, double theta)
{
    struct PsiGateOp* op = append_op(c);
    op->kind = kind;
    op->qubits[0] = control;
    op->qubits[1] = target;
    op->qubit_count = 2;
    op->params[0] = theta;
    op->param_count = 1;
}

static void push_3q(struct PsiQuantumCircuit* c, enum PsiGateKind kind, size_t a, size_t b,
                    size_t d)
{
    struct PsiGateOp* op = append_op(c);
    op->kind = kind;
    op->qubits[0] = a;
    op->qubits[1] = b;
    op->qubits[2] = d;
    op->qubit_count = 3;
}

void psi_apply_h(struct PsiQuantumCircuit* c, size_t target)
{
    push_1q(c, PSI_GATE_H, target);
}
void psi_apply_x(struct PsiQuantumCircuit* c, size_t target)
{
    push_1q(c, PSI_GATE_X, target);
}
void psi_apply_y(struct PsiQuantumCircuit* c, size_t target)
{
    push_1q(c, PSI_GATE_Y, target);
}
void psi_apply_z(struct PsiQuantumCircuit* c, size_t target)
{
    push_1q(c, PSI_GATE_Z, target);
}
void psi_apply_s(struct PsiQuantumCircuit* c, size_t target)
{
    push_1q(c, PSI_GATE_S, target);
}
void psi_apply_t(struct PsiQuantumCircuit* c, size_t target)
{
    push_1q(c, PSI_GATE_T, target);
}
void psi_apply_sdg(struct PsiQuantumCircuit* c, size_t target)
{
    push_1q(c, PSI_GATE_SDG, target);
}
void psi_apply_tdg(struct PsiQuantumCircuit* c, size_t target)
{
    push_1q(c, PSI_GATE_TDG, target);
}
void psi_apply_sx(struct PsiQuantumCircuit* c, size_t target)
{
    push_1q(c, PSI_GATE_SX, target);
}
void psi_apply_sxdg(struct PsiQuantumCircuit* c, size_t target)
{
    push_1q(c, PSI_GATE_SXDG, target);
}

void psi_apply_rx(struct PsiQuantumCircuit* c, size_t target, double theta)
{
    push_1q_1p(c, PSI_GATE_RX, target, theta);
}
void psi_apply_ry(struct PsiQuantumCircuit* c, size_t target, double theta)
{
    push_1q_1p(c, PSI_GATE_RY, target, theta);
}
void psi_apply_rz(struct PsiQuantumCircuit* c, size_t target, double theta)
{
    push_1q_1p(c, PSI_GATE_RZ, target, theta);
}
void psi_apply_p(struct PsiQuantumCircuit* c, size_t target, double theta)
{
    push_1q_1p(c, PSI_GATE_P, target, theta);
}
void psi_apply_u1(struct PsiQuantumCircuit* c, size_t target, double lambda)
{
    push_1q_1p(c, PSI_GATE_U1, target, lambda);
}

void psi_apply_u2(struct PsiQuantumCircuit* c, size_t target, double phi, double lambda)
{
    struct PsiGateOp* op = append_op(c);
    op->kind = PSI_GATE_U2;
    op->qubits[0] = target;
    op->qubit_count = 1;
    op->params[0] = phi;
    op->params[1] = lambda;
    op->param_count = 2;
}

void psi_apply_u3(struct PsiQuantumCircuit* c, size_t target, double theta, double phi,
                  double lambda)
{
    struct PsiGateOp* op = append_op(c);
    op->kind = PSI_GATE_U3;
    op->qubits[0] = target;
    op->qubit_count = 1;
    op->params[0] = theta;
    op->params[1] = phi;
    op->params[2] = lambda;
    op->param_count = 3;
}

void psi_apply_cnot(struct PsiQuantumCircuit* c, size_t control, size_t target)
{
    push_2q(c, PSI_GATE_CNOT, control, target);
}
void psi_apply_cz(struct PsiQuantumCircuit* c, size_t control, size_t target)
{
    push_2q(c, PSI_GATE_CZ, control, target);
}
void psi_apply_swap(struct PsiQuantumCircuit* c, size_t qubit1, size_t qubit2)
{
    push_2q(c, PSI_GATE_SWAP, qubit1, qubit2);
}
void psi_apply_crx(struct PsiQuantumCircuit* c, size_t control, size_t target, double theta)
{
    push_2q_1p(c, PSI_GATE_CRX, control, target, theta);
}
void psi_apply_cry(struct PsiQuantumCircuit* c, size_t control, size_t target, double theta)
{
    push_2q_1p(c, PSI_GATE_CRY, control, target, theta);
}
void psi_apply_crz(struct PsiQuantumCircuit* c, size_t control, size_t target, double theta)
{
    push_2q_1p(c, PSI_GATE_CRZ, control, target, theta);
}
void psi_apply_cp(struct PsiQuantumCircuit* c, size_t control, size_t target, double theta)
{
    push_2q_1p(c, PSI_GATE_CP, control, target, theta);
}

void psi_apply_ccnot(struct PsiQuantumCircuit* c, size_t control1, size_t control2, size_t target)
{
    push_3q(c, PSI_GATE_CCNOT, control1, control2, target);
}
void psi_apply_cswap(struct PsiQuantumCircuit* c, size_t control, size_t target1, size_t target2)
{
    push_3q(c, PSI_GATE_CSWAP, control, target1, target2);
}

void psi_measure(struct PsiQuantumCircuit* c, size_t qubit, size_t classical)
{
    if (classical >= c->num_classical)
        c->num_classical = classical + 1;

    struct PsiGateOp* op = append_op(c);
    op->kind = PSI_GATE_MEASURE;
    op->qubits[0] = qubit;
    op->qubit_count = 1;
    op->classical = classical;
}

void psi_measure_all(struct PsiQuantumCircuit* c)
{
    for (size_t i = 0; i < c->num_qubits; i++)
        psi_measure(c, i, i);
}

void psi_apply_custom(struct PsiQuantumCircuit* c, struct PsiCustomGate gate, const size_t* targets,
                      size_t count)
{
    struct PsiCustomGate* owned = malloc(sizeof(struct PsiCustomGate));
    assert(owned != NULL);
    *owned = gate;

    size_t* owned_targets = malloc(count * sizeof(size_t));
    assert(owned_targets != NULL || count == 0);
    if (count > 0)
        memcpy(owned_targets, targets, count * sizeof(size_t));

    struct PsiGateOp* op = append_op(c);
    op->kind = PSI_GATE_CUSTOM;
    op->custom = owned;
    op->custom_targets = owned_targets;
    op->custom_target_count = count;
}

void psi_print_circuit(const struct PsiQuantumCircuit* circuit, FILE* out)
{
    size_t nq = circuit->num_qubits;
    fprintf(out, "QuantumCircuit (%zu qubits, %zu classical)\n", nq, circuit->num_classical);
    fprintf(out, "Operations:\n");

    for (size_t i = 0; i < circuit->operation_count; i++)
    {
        struct PsiGateOp op = circuit->operations[i];

        if (op.kind == PSI_GATE_MEASURE)
        {
            fprintf(out, "  %zu: %s q%zu → c%zu\n", i, psi_gate_op_name(op), op.qubits[0],
                    op.classical);
            continue;
        }

        if (op.kind == PSI_GATE_CUSTOM)
            fprintf(out, "  %zu: [%s] on [", i, op.custom->name);
        else
            fprintf(out, "  %zu: %s on [", i, psi_gate_op_name(op));

        size_t target_count;
        const size_t* targets = psi_gate_op_quantum_targets(&op, &target_count);
        for (size_t j = 0; j < target_count; j++)
        {
            fprintf(out, "%zu", targets[j]);
            if (j != target_count - 1)
                fputs(", ", out);
        }
        fputs("]\n", out);
    }

    if (!circuit->is_computed)
    {
        fputs("State: (not computed)\n", out);
        return;
    }

    fputs("State:\n", out);
    size_t dim = (size_t)1 << nq;
    for (size_t i = 0; i < dim; i++)
    {
        struct PsiComplex amp = circuit->computed_state.data[i];
        if (fabs(amp.real) < 1e-10 && fabs(amp.imaginary) < 1e-10)
            continue;

        char basis[65];
        for (size_t b = 0; b < nq; b++)
            basis[b] = ((i >> (nq - 1 - b)) & 1) ? '1' : '0';
        basis[nq] = '\0';

        char amp_buf[64];
        fprintf(out, "  |%s⟩: %s\n", basis, psi_format_amplitude(amp, amp_buf, sizeof amp_buf));
    }
}