#include "core/quantum_components.h"

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

struct PsiVector psi_new_state_0(void)
{
	return psi_column_vector(psi_new_complex(1.0, 0.0), psi_new_complex(0.0, 0.0));
}

struct PsiVector psi_new_state_1(void)
{
	return psi_column_vector(psi_new_complex(0.0, 0.0), psi_new_complex(1.0, 0.0));
}

struct PsiQuantumGate psi_new_quantum_gate(const char* name, struct PsiMatrix matrix,
                                           size_t num_qubits)
{
	size_t expected_dim = (size_t)1 << num_qubits;
	assert(matrix.rows == expected_dim);
	assert(matrix.cols == expected_dim);

	return (struct PsiQuantumGate){
		name,
		matrix,
		num_qubits,
	};
}

struct PsiQuantumGate psi_new_quantum_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++;

	return (struct PsiQuantumGate){
		name,
		matrix,
		num_qubits,
	};
}

void psi_free_quantum_gate(struct PsiQuantumGate* gate)
{
	psi_free_matrix(&gate->matrix);
}

struct PsiQuantumBit psi_new_quantum_bit(const char* name, struct PsiVector state)
{
	return (struct PsiQuantumBit){
		name,
		state,
	};
}

void psi_free_quantum_bit(struct PsiQuantumBit* bit)
{
	psi_free_vector(&bit->state);
}

static void update_register(struct PsiQuantumRegister* reg)
{
	psi_free_vector(&reg->state_vector);

	if (reg->num_qubits == 0)
	{
		reg->state_vector = psi_new_vector(0, PSI_COLUMN_VECTOR);
		return;
	}

	struct PsiMatrix result = psi_matrix_from_vector(reg->qubits[0].state);
	for (size_t i = 1; i < reg->num_qubits; i++)
	{
		struct PsiMatrix part = psi_matrix_from_vector(reg->qubits[i].state);
		struct PsiMatrix next = psi_kronecker_matrix(result, part);
		psi_free_matrix(&result);
		psi_free_matrix(&part);
		result = next;
	}

	reg->state_vector = psi_vector_from_matrix(result, PSI_COLUMN_VECTOR);
	psi_free_matrix(&result);
}

struct PsiQuantumRegister psi_new_quantum_register(const char* name, const char** names,
                                                   size_t count)
{
	struct PsiQuantumBit* qubits = malloc(count * sizeof(struct PsiQuantumBit));
	assert(qubits != NULL || count == 0);

	for (size_t i = 0; i < count; i++)
		qubits[i] = psi_new_quantum_bit(names[i], psi_new_state_0());

	struct PsiQuantumRegister reg = {
		name,
		psi_new_vector(0, PSI_COLUMN_VECTOR),
		qubits,
		count,
	};

	update_register(&reg);
	return reg;
}

struct PsiQuantumRegister
psi_new_quantum_register_from(const char* name, const struct PsiQuantumBit* bits, size_t count)
{
	struct PsiQuantumBit* qubits = malloc(count * sizeof(struct PsiQuantumBit));
	assert(qubits != NULL || count == 0);

	for (size_t i = 0; i < count; i++)
		qubits[i] = psi_new_quantum_bit(bits[i].name, psi_clone_vector(bits[i].state));

	struct PsiQuantumRegister reg = {
		name,
		psi_new_vector(0, PSI_COLUMN_VECTOR),
		qubits,
		count,
	};

	update_register(&reg);
	return reg;
}

void psi_free_quantum_register(struct PsiQuantumRegister* reg)
{
	for (size_t i = 0; i < reg->num_qubits; i++)
		psi_free_quantum_bit(&reg->qubits[i]);

	free(reg->qubits);
	reg->qubits = NULL;
	reg->num_qubits = 0;
	psi_free_vector(&reg->state_vector);
}

static bool targets_contain(const size_t* targets, size_t count, size_t value)
{
	for (size_t i = 0; i < count; i++)
		if (targets[i] == value)
			return true;

	return false;
}

static struct PsiMatrix build_contiguous_operator(struct PsiQuantumRegister reg,
                                                  struct PsiQuantumGate gate, size_t start_idx)
{
	size_t n = reg.num_qubits;
	size_t g = gate.num_qubits;

	bool has_result = false;
	struct PsiMatrix result = { 0 };

	for (size_t i = 0; i < n; i++)
	{
		if (i > start_idx && i < start_idx + g)
			continue;

		struct PsiMatrix part =
				i == start_idx ? psi_clone_matrix(gate.matrix) : psi_identity_matrix(2);

		if (!has_result)
		{
			result = part;
			has_result = true;
			continue;
		}

		struct PsiMatrix next = psi_kronecker_matrix(result, part);
		psi_free_matrix(&result);
		psi_free_matrix(&part);
		result = next;
	}

	if (!has_result)
		return psi_identity_matrix((size_t)1 << n);

	return result;
}

static struct PsiMatrix build_full_operator(struct PsiQuantumRegister reg,
                                            struct PsiQuantumGate gate, const size_t* targets,
                                            size_t target_count)
{
	size_t n = reg.num_qubits;
	size_t g = gate.num_qubits;
	size_t dim = (size_t)1 << n;

	bool contiguous = true;
	for (size_t i = 1; i < target_count; i++)
		if (targets[i] != targets[i - 1] + 1)
		{
			contiguous = false;
			break;
		}

	if (contiguous && g == n)
		return psi_clone_matrix(gate.matrix);

	if (contiguous)
		return build_contiguous_operator(reg, gate, targets[0]);

	struct PsiMatrix result = psi_new_matrix(dim, dim);

	for (size_t col = 0; col < dim; col++)
		for (size_t row = 0; row < dim; row++)
		{
			size_t target_row_bits = 0;
			size_t target_col_bits = 0;

			for (size_t i = 0; i < target_count; i++)
			{
				size_t qubit_pos = n - 1 - targets[i];
				if ((row >> qubit_pos) & 1)
					target_row_bits |= (size_t)1 << (g - 1 - i);
				if ((col >> qubit_pos) & 1)
					target_col_bits |= (size_t)1 << (g - 1 - i);
			}

			bool non_target_match = true;
			for (size_t q = 0; q < n; q++)
			{
				if (targets_contain(targets, target_count, q))
					continue;

				size_t qubit_pos = n - 1 - q;
				if (((row >> qubit_pos) & 1) != ((col >> qubit_pos) & 1))
				{
					non_target_match = false;
					break;
				}
			}

			if (non_target_match)
				result.data[row * result.cols + col] =
						psi_get_matrix(gate.matrix, target_row_bits, target_col_bits);
		}

	return result;
}

void psi_apply_gate(struct PsiQuantumRegister* reg, struct PsiQuantumGate gate,
                    const size_t* targets, size_t target_count)
{
	size_t n = reg->num_qubits;

	assert(gate.num_qubits == target_count);
	for (size_t i = 0; i < target_count; i++)
		assert(targets[i] < n);

	for (size_t i = 0; i < target_count; i++)
		for (size_t j = i + 1; j < target_count; j++)
			assert(targets[i] != targets[j]);

	struct PsiMatrix full_operator = build_full_operator(*reg, gate, targets, target_count);
	struct PsiVector new_state = psi_mul_vector_matrix(reg->state_vector, full_operator);

	psi_free_vector(&reg->state_vector);
	psi_free_matrix(&full_operator);
	reg->state_vector = new_state;
}