#pragma once

#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>

#include "core/custom_gate.h"
#include "maths/vector.h"

enum PsiGateKind
{
    PSI_GATE_H,
    PSI_GATE_X,
    PSI_GATE_Y,
    PSI_GATE_Z,
    PSI_GATE_S,
    PSI_GATE_T,
    PSI_GATE_SDG,
    PSI_GATE_TDG,
    PSI_GATE_SX,
    PSI_GATE_SXDG,
    PSI_GATE_RX,
    PSI_GATE_RY,
    PSI_GATE_RZ,
    PSI_GATE_P,
    PSI_GATE_U1,
    PSI_GATE_U2,
    PSI_GATE_U3,
    PSI_GATE_CNOT,
    PSI_GATE_CZ,
    PSI_GATE_SWAP,
    PSI_GATE_CRX,
    PSI_GATE_CRY,
    PSI_GATE_CRZ,
    PSI_GATE_CP,
    PSI_GATE_CCNOT,
    PSI_GATE_CSWAP,
    PSI_GATE_MEASURE,
    PSI_GATE_CUSTOM,
};

struct PsiGateOp
{
    enum PsiGateKind kind;
    size_t qubits[3];
    size_t qubit_count;
    double params[3];
    size_t param_count;
    size_t classical;
    struct PsiCustomGate* custom;
    size_t* custom_targets;
    size_t custom_target_count;
};

const char* psi_gate_op_name(struct PsiGateOp op);
const size_t* psi_gate_op_quantum_targets(const struct PsiGateOp* op, size_t* out_count);
const size_t* psi_gate_op_classical_targets(const struct PsiGateOp* op, size_t* out_count);
bool psi_gate_op_is_measurement(struct PsiGateOp op);
bool psi_gate_op_is_custom(struct PsiGateOp op);
bool psi_gate_op_is_non_clifford(struct PsiGateOp op);

struct PsiQuantumCircuit
{
    size_t num_qubits;
    size_t num_classical;
    struct PsiGateOp* operations;
    size_t operation_count;
    size_t operation_capacity;
    struct PsiVector computed_state;
    bool is_computed;
};

struct PsiQuantumCircuit psi_new_quantum_circuit(size_t num_qubits);
struct PsiQuantumCircuit psi_new_quantum_circuit_with_classical(size_t num_qubits,
                                                                size_t num_classical);
void psi_free_quantum_circuit(struct PsiQuantumCircuit* c);
void psi_reset_circuit(struct PsiQuantumCircuit* c);

void psi_apply_h(struct PsiQuantumCircuit* c, size_t target);
void psi_apply_x(struct PsiQuantumCircuit* c, size_t target);
void psi_apply_y(struct PsiQuantumCircuit* c, size_t target);
void psi_apply_z(struct PsiQuantumCircuit* c, size_t target);
void psi_apply_s(struct PsiQuantumCircuit* c, size_t target);
void psi_apply_t(struct PsiQuantumCircuit* c, size_t target);
void psi_apply_sdg(struct PsiQuantumCircuit* c, size_t target);
void psi_apply_tdg(struct PsiQuantumCircuit* c, size_t target);
void psi_apply_sx(struct PsiQuantumCircuit* c, size_t target);
void psi_apply_sxdg(struct PsiQuantumCircuit* c, size_t target);

void psi_apply_rx(struct PsiQuantumCircuit* c, size_t target, double theta);
void psi_apply_ry(struct PsiQuantumCircuit* c, size_t target, double theta);
void psi_apply_rz(struct PsiQuantumCircuit* c, size_t target, double theta);
void psi_apply_p(struct PsiQuantumCircuit* c, size_t target, double theta);
void psi_apply_u1(struct PsiQuantumCircuit* c, size_t target, double lambda);
void psi_apply_u2(struct PsiQuantumCircuit* c, size_t target, double phi, double lambda);
void psi_apply_u3(struct PsiQuantumCircuit* c, size_t target, double theta, double phi,
                  double lambda);

void psi_apply_cnot(struct PsiQuantumCircuit* c, size_t control, size_t target);
void psi_apply_cz(struct PsiQuantumCircuit* c, size_t control, size_t target);
void psi_apply_swap(struct PsiQuantumCircuit* c, size_t qubit1, size_t qubit2);
void psi_apply_crx(struct PsiQuantumCircuit* c, size_t control, size_t target, double theta);
void psi_apply_cry(struct PsiQuantumCircuit* c, size_t control, size_t target, double theta);
void psi_apply_crz(struct PsiQuantumCircuit* c, size_t control, size_t target, double theta);
void psi_apply_cp(struct PsiQuantumCircuit* c, size_t control, size_t target, double theta);

void psi_apply_ccnot(struct PsiQuantumCircuit* c, size_t control1, size_t control2, size_t target);
void psi_apply_cswap(struct PsiQuantumCircuit* c, size_t control, size_t target1, size_t target2);

void psi_measure(struct PsiQuantumCircuit* c, size_t qubit, size_t classical);
void psi_measure_all(struct PsiQuantumCircuit* c);

void psi_apply_custom(struct PsiQuantumCircuit* c, struct PsiCustomGate gate, const size_t* targets,
                      size_t count);

void psi_print_circuit(const struct PsiQuantumCircuit* circuit, FILE* out);