#include "core/noise.h"

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

struct PsiKrausOperator psi_new_kraus_operator(const char* name, struct PsiMatrix matrix)
{
	return (struct PsiKrausOperator){
		matrix,
		name,
	};
}

void psi_free_kraus_operator(struct PsiKrausOperator* op)
{
	psi_free_matrix(&op->matrix);
}

struct PsiNoiseChannel psi_new_noise_channel(const char* name,
                                             const struct PsiKrausOperator* operators, size_t count,
                                             size_t num_qubits)
{
	struct PsiKrausOperator* owned = malloc(count * sizeof(struct PsiKrausOperator));
	assert(owned != NULL || count == 0);

	if (count > 0)
		memcpy(owned, operators, count * sizeof(struct PsiKrausOperator));

	return (struct PsiNoiseChannel){
		name,
		owned,
		count,
		num_qubits,
	};
}

void psi_free_noise_channel(struct PsiNoiseChannel* channel)
{
	for (size_t i = 0; i < channel->operator_count; i++)
		psi_free_matrix(&channel->operators[i].matrix);

	free(channel->operators);
	channel->operators = NULL;
	channel->operator_count = 0;
}

struct PsiNoiseChannel psi_depolarising_channel(double p)
{
	double sqrt_1_p = sqrt(1.0 - p);
	double sqrt_p3 = sqrt(p / 3.0);

	struct PsiKrausOperator ops[] = {
		psi_new_kraus_operator("K0",
		                       psi_matrix(2, 2, psi_new_complex(sqrt_1_p, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_1_p, 0.0))),
		psi_new_kraus_operator(
				"K1(X)",
				psi_matrix(2, 2, psi_new_complex(0.0, 0.0), psi_new_complex(sqrt_p3, 0.0),
				           psi_new_complex(sqrt_p3, 0.0), psi_new_complex(0.0, 0.0))),
		psi_new_kraus_operator(
				"K2(Y)",
				psi_matrix(2, 2, psi_new_complex(0.0, 0.0), psi_new_complex(0.0, -sqrt_p3),
				           psi_new_complex(0.0, sqrt_p3), psi_new_complex(0.0, 0.0))),
		psi_new_kraus_operator("K3(Z)",
		                       psi_matrix(2, 2, psi_new_complex(sqrt_p3, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(-sqrt_p3, 0.0))),
	};

	return psi_new_noise_channel("Depolarising", ops, 4, 1);
}

struct PsiNoiseChannel psi_amplitude_damping_channel(double gamma)
{
	double sqrt_gamma = sqrt(gamma);
	double sqrt_1_gamma = sqrt(1.0 - gamma);

	struct PsiKrausOperator ops[] = {
		psi_new_kraus_operator("K0",
		                       psi_matrix(2, 2, psi_new_complex(1.0, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_1_gamma, 0.0))),
		psi_new_kraus_operator("K1",
		                       psi_matrix(2, 2, psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_gamma, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0))),
	};

	return psi_new_noise_channel("AmplitudeDamping", ops, 2, 1);
}

struct PsiNoiseChannel psi_phase_damping_channel(double gamma)
{
	double sqrt_gamma = sqrt(gamma);
	double sqrt_1_gamma = sqrt(1.0 - gamma);

	struct PsiKrausOperator ops[] = {
		psi_new_kraus_operator("K0",
		                       psi_matrix(2, 2, psi_new_complex(1.0, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_1_gamma, 0.0))),
		psi_new_kraus_operator("K1",
		                       psi_matrix(2, 2, psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_gamma, 0.0))),
	};

	return psi_new_noise_channel("PhaseDamping", ops, 2, 1);
}

struct PsiNoiseChannel psi_bit_flip_channel(double p)
{
	double sqrt_1_p = sqrt(1.0 - p);
	double sqrt_p = sqrt(p);

	struct PsiKrausOperator ops[] = {
		psi_new_kraus_operator("K0(I)",
		                       psi_matrix(2, 2, psi_new_complex(sqrt_1_p, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_1_p, 0.0))),
		psi_new_kraus_operator("K1(X)",
		                       psi_matrix(2, 2, psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_p, 0.0),
		                                  psi_new_complex(sqrt_p, 0.0), psi_new_complex(0.0, 0.0))),
	};

	return psi_new_noise_channel("BitFlip", ops, 2, 1);
}

struct PsiNoiseChannel psi_phase_flip_channel(double p)
{
	double sqrt_1_p = sqrt(1.0 - p);
	double sqrt_p = sqrt(p);

	struct PsiKrausOperator ops[] = {
		psi_new_kraus_operator("K0(I)",
		                       psi_matrix(2, 2, psi_new_complex(sqrt_1_p, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_1_p, 0.0))),
		psi_new_kraus_operator("K1(Z)",
		                       psi_matrix(2, 2, psi_new_complex(sqrt_p, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(-sqrt_p, 0.0))),
	};

	return psi_new_noise_channel("PhaseFlip", ops, 2, 1);
}

struct PsiNoiseChannel psi_bit_phase_flip_channel(double p)
{
	double sqrt_1_p = sqrt(1.0 - p);
	double sqrt_p = sqrt(p);

	struct PsiKrausOperator ops[] = {
		psi_new_kraus_operator("K0(I)",
		                       psi_matrix(2, 2, psi_new_complex(sqrt_1_p, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_1_p, 0.0))),
		psi_new_kraus_operator("K1(Y)",
		                       psi_matrix(2, 2, psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(0.0, -sqrt_p),
		                                  psi_new_complex(0.0, sqrt_p), psi_new_complex(0.0, 0.0))),
	};

	return psi_new_noise_channel("BitPhaseFlip", ops, 2, 1);
}

struct PsiNoiseChannel psi_generalised_amplitude_damping_channel(double p, double gamma)
{
	double sqrt_p = sqrt(p);
	double sqrt_1_p = sqrt(1.0 - p);
	double sqrt_gamma = sqrt(gamma);
	double sqrt_1_gamma = sqrt(1.0 - gamma);

	struct PsiKrausOperator ops[] = {
		psi_new_kraus_operator("K0",
		                       psi_matrix(2, 2, psi_new_complex(sqrt_p, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_p * sqrt_1_gamma, 0.0))),
		psi_new_kraus_operator("K1",
		                       psi_matrix(2, 2, psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_p * sqrt_gamma, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0))),
		psi_new_kraus_operator("K2",
		                       psi_matrix(2, 2, psi_new_complex(sqrt_1_p * sqrt_1_gamma, 0.0),
		                                  psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
		                                  psi_new_complex(sqrt_1_p, 0.0))),
		psi_new_kraus_operator(
				"K3",
				psi_matrix(2, 2, psi_new_complex(0.0, 0.0), psi_new_complex(0.0, 0.0),
				           psi_new_complex(sqrt_1_p * sqrt_gamma, 0.0), psi_new_complex(0.0, 0.0))),
	};

	return psi_new_noise_channel("GeneralisedAmplitudeDamping", ops, 4, 1);
}

struct PsiDensityMatrix psi_new_density_matrix(size_t num_qubits)
{
	size_t dim = (size_t)1 << num_qubits;
	struct PsiComplex* data = calloc(dim * dim, sizeof(struct PsiComplex));
	assert(data != NULL);
	data[0] = psi_new_complex(1.0, 0.0);

	return (struct PsiDensityMatrix){
		data,
		dim,
		num_qubits,
	};
}

struct PsiDensityMatrix psi_new_density_matrix_from_state(const struct PsiComplex* state,
                                                          size_t len)
{
	size_t dim = len;
	size_t num_qubits = 0;
	while (((size_t)1 << num_qubits) < dim)
		num_qubits++;

	struct PsiComplex* data = malloc(dim * dim * sizeof(struct PsiComplex));
	assert(data != NULL);

	for (size_t i = 0; i < dim; i++)
		for (size_t j = 0; j < dim; j++)
			data[i * dim + j] = psi_mul_complex(state[i], psi_conjugate_complex(state[j]));

	return (struct PsiDensityMatrix){
		data,
		dim,
		num_qubits,
	};
}

void psi_free_density_matrix(struct PsiDensityMatrix* dm)
{
	free(dm->data);
	dm->data = NULL;
	dm->dim = 0;
	dm->num_qubits = 0;
}

struct PsiComplex psi_get_density_matrix(struct PsiDensityMatrix dm, size_t row, size_t col)
{
	assert(row < dm.dim && col < dm.dim);
	return dm.data[row * dm.dim + col];
}

void psi_set_density_matrix(struct PsiDensityMatrix* dm, size_t row, size_t col,
                            struct PsiComplex value)
{
	assert(row < dm->dim && col < dm->dim);
	dm->data[row * dm->dim + col] = value;
}

struct PsiComplex psi_trace_density_matrix(struct PsiDensityMatrix dm)
{
	struct PsiComplex sum = psi_new_complex(0.0, 0.0);
	for (size_t i = 0; i < dm.dim; i++)
		sum = psi_add_complex(sum, dm.data[i * dm.dim + i]);

	return sum;
}

double psi_purity_density_matrix(struct PsiDensityMatrix dm)
{
	struct PsiComplex sum = psi_new_complex(0.0, 0.0);
	for (size_t i = 0; i < dm.dim; i++)
		for (size_t j = 0; j < dm.dim; j++)
			sum = psi_add_complex(
					sum, psi_mul_complex(dm.data[i * dm.dim + j], dm.data[j * dm.dim + i]));

	return sum.real;
}

bool psi_is_pure_density_matrix(struct PsiDensityMatrix dm, double tolerance)
{
	return fabs(psi_purity_density_matrix(dm) - 1.0) < tolerance;
}

void psi_density_matrix_probabilities(struct PsiDensityMatrix dm, double* out)
{
	for (size_t i = 0; i < dm.dim; i++)
		out[i] = dm.data[i * dm.dim + i].real;
}

void psi_apply_unitary_density_matrix(struct PsiDensityMatrix* dm, struct PsiMatrix gate,
                                      const size_t* targets, size_t target_count)
{
	size_t g = target_count;
	size_t gate_dim = (size_t)1 << g;
	size_t dim = dm->dim;

	size_t* target_bits = malloc(g * sizeof(size_t));
	assert(target_bits != NULL || g == 0);
	for (size_t t = 0; t < g; t++)
		target_bits[t] = dm->num_qubits - 1 - targets[t];

	size_t non_target_mask = dim - 1;
	for (size_t t = 0; t < g; t++)
		non_target_mask &= ~((size_t)1 << target_bits[t]);

	struct PsiComplex* new_data = calloc(dim * dim, sizeof(struct PsiComplex));
	assert(new_data != NULL);

	for (size_t i = 0; i < dim; i++)
		for (size_t j = 0; j < dim; j++)
		{
			struct PsiComplex sum = psi_new_complex(0.0, 0.0);

			for (size_t k = 0; k < gate_dim; k++)
				for (size_t l = 0; l < gate_dim; l++)
				{
					size_t src_i = i & non_target_mask;
					size_t src_j = j & non_target_mask;

					for (size_t idx = 0; idx < g; idx++)
					{
						if ((k >> (g - 1 - idx)) & 1)
							src_i |= (size_t)1 << target_bits[idx];
						if ((l >> (g - 1 - idx)) & 1)
							src_j |= (size_t)1 << target_bits[idx];
					}

					size_t tgt_i = 0;
					size_t tgt_j = 0;
					for (size_t idx = 0; idx < g; idx++)
					{
						if ((i >> target_bits[idx]) & 1)
							tgt_i |= (size_t)1 << (g - 1 - idx);
						if ((j >> target_bits[idx]) & 1)
							tgt_j |= (size_t)1 << (g - 1 - idx);
					}

					struct PsiComplex u_ik = gate.data[tgt_i * gate_dim + k];
					struct PsiComplex u_jl_dag =
							psi_conjugate_complex(gate.data[tgt_j * gate_dim + l]);
					struct PsiComplex rho_kl = dm->data[src_i * dim + src_j];

					sum = psi_add_complex(sum,
					                      psi_mul_complex(psi_mul_complex(u_ik, rho_kl), u_jl_dag));
				}

			new_data[i * dim + j] = sum;
		}

	free(target_bits);
	free(dm->data);
	dm->data = new_data;
}

void psi_apply_noise_channel(struct PsiDensityMatrix* dm, struct PsiNoiseChannel channel,
                             size_t target)
{
	assert(channel.num_qubits == 1);

	size_t dim = dm->dim;
	size_t target_bit = dm->num_qubits - 1 - target;

	struct PsiComplex* new_data = calloc(dim * dim, sizeof(struct PsiComplex));
	assert(new_data != NULL);

	for (size_t op = 0; op < channel.operator_count; op++)
	{
		struct PsiMatrix k = channel.operators[op].matrix;

		for (size_t i = 0; i < dim; i++)
			for (size_t j = 0; j < dim; j++)
			{
				size_t i_target = (i >> target_bit) & 1;
				size_t j_target = (j >> target_bit) & 1;

				for (size_t ki = 0; ki < 2; ki++)
					for (size_t kj = 0; kj < 2; kj++)
					{
						size_t src_i = (i & ~((size_t)1 << target_bit)) | (ki << target_bit);
						size_t src_j = (j & ~((size_t)1 << target_bit)) | (kj << target_bit);

						struct PsiComplex k_elem = k.data[i_target * 2 + ki];
						struct PsiComplex k_dag_elem =
								psi_conjugate_complex(k.data[j_target * 2 + kj]);
						struct PsiComplex rho_elem = dm->data[src_i * dim + src_j];

						struct PsiComplex term =
								psi_mul_complex(psi_mul_complex(k_elem, rho_elem), k_dag_elem);
						new_data[i * dim + j] = psi_add_complex(new_data[i * dim + j], term);
					}
			}
	}

	free(dm->data);
	dm->data = new_data;
}

double psi_measure_probability_density_matrix(struct PsiDensityMatrix dm, size_t qubit,
                                              size_t outcome)
{
	size_t target_bit = dm.num_qubits - 1 - qubit;
	double prob = 0.0;

	for (size_t i = 0; i < dm.dim; i++)
		if (((i >> target_bit) & 1) == outcome)
			prob += dm.data[i * dm.dim + i].real;

	return prob;
}

double psi_fidelity_density_matrix(struct PsiDensityMatrix dm, const struct PsiComplex* state)
{
	struct PsiComplex sum = psi_new_complex(0.0, 0.0);

	for (size_t i = 0; i < dm.dim; i++)
		for (size_t j = 0; j < dm.dim; j++)
			sum = psi_add_complex(sum,
			                      psi_mul_complex(psi_mul_complex(psi_conjugate_complex(state[i]),
			                                                      dm.data[i * dm.dim + j]),
			                                      state[j]));

	return sum.real;
}