#pragma once

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

#include "maths/complex.h"
#include "maths/matrix.h"

struct PsiKrausOperator
{
	struct PsiMatrix matrix;
	const char* name;
};

struct PsiKrausOperator psi_new_kraus_operator(const char* name, struct PsiMatrix matrix);
void psi_free_kraus_operator(struct PsiKrausOperator* op);

struct PsiNoiseChannel
{
	const char* name;
	struct PsiKrausOperator* operators;
	size_t operator_count;
	size_t num_qubits;
};

struct PsiNoiseChannel psi_new_noise_channel(const char* name,
                                             const struct PsiKrausOperator* operators, size_t count,
                                             size_t num_qubits);
void psi_free_noise_channel(struct PsiNoiseChannel* channel);

struct PsiNoiseChannel psi_depolarising_channel(double p);
struct PsiNoiseChannel psi_amplitude_damping_channel(double gamma);
struct PsiNoiseChannel psi_phase_damping_channel(double gamma);
struct PsiNoiseChannel psi_bit_flip_channel(double p);
struct PsiNoiseChannel psi_phase_flip_channel(double p);
struct PsiNoiseChannel psi_bit_phase_flip_channel(double p);
struct PsiNoiseChannel psi_generalised_amplitude_damping_channel(double p, double gamma);

struct PsiDensityMatrix
{
	struct PsiComplex* data;
	size_t dim;
	size_t num_qubits;
};

struct PsiDensityMatrix psi_new_density_matrix(size_t num_qubits);
struct PsiDensityMatrix psi_new_density_matrix_from_state(const struct PsiComplex* state,
                                                          size_t len);
void psi_free_density_matrix(struct PsiDensityMatrix* dm);

struct PsiComplex psi_get_density_matrix(struct PsiDensityMatrix dm, size_t row, size_t col);
void psi_set_density_matrix(struct PsiDensityMatrix* dm, size_t row, size_t col,
                            struct PsiComplex value);
struct PsiComplex psi_trace_density_matrix(struct PsiDensityMatrix dm);
double psi_purity_density_matrix(struct PsiDensityMatrix dm);
bool psi_is_pure_density_matrix(struct PsiDensityMatrix dm, double tolerance);
void psi_density_matrix_probabilities(struct PsiDensityMatrix dm, double* out);
void psi_apply_unitary_density_matrix(struct PsiDensityMatrix* dm, struct PsiMatrix gate,
                                      const size_t* targets, size_t target_count);
void psi_apply_noise_channel(struct PsiDensityMatrix* dm, struct PsiNoiseChannel channel,
                             size_t target);
double psi_measure_probability_density_matrix(struct PsiDensityMatrix dm, size_t qubit,
                                              size_t outcome);
double psi_fidelity_density_matrix(struct PsiDensityMatrix dm, const struct PsiComplex* state);