#include "core/kernel.h"

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

static char *dup_string(const char *s)
{
	size_t n = strlen(s) + 1;
	char *p = malloc(n);
	assert(p != NULL);
	memcpy(p, s, n);

	return p;
}

static size_t *dup_targets(const size_t *targets, size_t count)
{
	size_t *p = malloc(count * sizeof(size_t));
	assert(p != NULL || count == 0);

	if (count > 0)
		memcpy(p, targets, count * sizeof(size_t));

	return p;
}

static bool starts_with(const char *s, const char *prefix)
{
	return strncmp(s, prefix, strlen(prefix)) == 0;
}

static enum PsiGateType detect_gate_type(const char *name, struct PsiMatrix matrix)
{
	static const char *diagonal[] = { "Z", "S", "T", "Sdg", "Tdg", "Rz", "P", "U1", "CZ", "CP", "CRz" };
	for (size_t i = 0; i < sizeof(diagonal) / sizeof(diagonal[0]); i++)
		if (starts_with(name, diagonal[i]))
			return PSI_GATE_TYPE_DIAGONAL;

	static const char *controlled[] = { "CNOT", "CZ", "SWAP", "CRx", "CRy", "CRz", "CP", "CCNOT", "CSWAP" };
	for (size_t i = 0; i < sizeof(controlled) / sizeof(controlled[0]); i++)
		if (starts_with(name, controlled[i]))
			return PSI_GATE_TYPE_CONTROLLED;

	if (matrix.rows == 2 && matrix.cols == 2)
	{
		bool is_diag = fabs(matrix.data[1].real) < 1e-10
			&& fabs(matrix.data[1].imaginary) < 1e-10
			&& fabs(matrix.data[2].real) < 1e-10
			&& fabs(matrix.data[2].imaginary) < 1e-10;
		if (is_diag)
			return PSI_GATE_TYPE_DIAGONAL;
	}

	return PSI_GATE_TYPE_NON_DIAGONAL;
}

struct PsiKernel psi_new_kernel(const char *name, struct PsiMatrix matrix, const size_t *targets, size_t target_count)
{
	struct PsiKernel kernel;
	kernel.matrix = matrix;
	kernel.targets = dup_targets(targets, target_count);
	kernel.target_count = target_count;
	kernel.name = dup_string(name);
	kernel.gate_type = detect_gate_type(name, matrix);

	return kernel;
}

struct PsiKernel psi_clone_kernel(struct PsiKernel kernel)
{
	struct PsiKernel copy;
	copy.matrix = psi_clone_matrix(kernel.matrix);
	copy.targets = dup_targets(kernel.targets, kernel.target_count);
	copy.target_count = kernel.target_count;
	copy.name = dup_string(kernel.name);
	copy.gate_type = kernel.gate_type;

	return copy;
}

void psi_free_kernel(struct PsiKernel *kernel)
{
	psi_free_matrix(&kernel->matrix);
	free(kernel->targets);
	free(kernel->name);
	kernel->targets = NULL;
	kernel->name = NULL;
	kernel->target_count = 0;
}

static bool targets_equal(struct PsiKernel a, struct PsiKernel b)
{
	if (a.target_count != b.target_count)
		return false;

	for (size_t i = 0; i < a.target_count; i++)
		if (a.targets[i] != b.targets[i])
			return false;

	return true;
}

bool psi_kernels_share_qubits(struct PsiKernel a, struct PsiKernel b)
{
	for (size_t i = 0; i < a.target_count; i++)
		for (size_t j = 0; j < b.target_count; j++)
			if (a.targets[i] == b.targets[j])
				return true;

	return false;
}

bool psi_kernels_commute(struct PsiKernel a, struct PsiKernel b)
{
	if (!psi_kernels_share_qubits(a, b))
		return true;

	if (a.gate_type == PSI_GATE_TYPE_DIAGONAL && b.gate_type == PSI_GATE_TYPE_DIAGONAL && targets_equal(a, b))
		return true;

	return false;
}

bool psi_kernels_can_fuse(struct PsiKernel a, struct PsiKernel b)
{
	if (a.target_count != 1 || b.target_count != 1)
		return false;

	return a.targets[0] == b.targets[0];
}

bool psi_fuse_kernels(struct PsiKernel a, struct PsiKernel b, struct PsiKernel *out)
{
	if (!psi_kernels_can_fuse(a, b))
		return false;

	enum PsiGateType new_type = a.gate_type == PSI_GATE_TYPE_DIAGONAL && b.gate_type == PSI_GATE_TYPE_DIAGONAL
		? PSI_GATE_TYPE_DIAGONAL
		: PSI_GATE_TYPE_NON_DIAGONAL;

	size_t name_len = strlen(a.name) + strlen(b.name) + 2;
	char *fused_name = malloc(name_len);
	assert(fused_name != NULL);
	snprintf(fused_name, name_len, "%s+%s", a.name, b.name);

	out->matrix = psi_dot_matrix(b.matrix, a.matrix);
	out->targets = dup_targets(a.targets, a.target_count);
	out->target_count = a.target_count;
	out->name = fused_name;
	out->gate_type = new_type;

	return true;
}

static struct PsiComplex *apply_kernel(const struct PsiComplex *state, struct PsiKernel kernel, size_t num_qubits)
{
	size_t dim = (size_t)1 << num_qubits;
	size_t g = kernel.target_count;
	size_t gate_dim = (size_t)1 << g;

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

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

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

	for (size_t i = 0; i < dim; i++)
	{
		size_t target_idx = 0;
		for (size_t k = 0; k < g; k++)
			if ((i >> target_bits[k]) & 1)
				target_idx |= (size_t)1 << (g - 1 - k);

		struct PsiComplex sum = psi_new_complex(0.0, 0.0);
		for (size_t j = 0; j < gate_dim; j++)
		{
			struct PsiComplex gate_elem = kernel.matrix.data[target_idx * gate_dim + j];
			if (fabs(gate_elem.real) < 1e-15 && fabs(gate_elem.imaginary) < 1e-15)
				continue;

			size_t source_idx = i & non_target_mask;
			for (size_t k = 0; k < g; k++)
				if ((j >> (g - 1 - k)) & 1)
					source_idx |= (size_t)1 << target_bits[k];

			sum = psi_add_complex(sum, psi_mul_complex(gate_elem, state[source_idx]));
		}

		new_state[i] = sum;
	}

	free(target_bits);
	return new_state;
}

struct PsiKernelBatch psi_new_kernel_batch(size_t num_qubits)
{
	struct PsiKernelBatch batch;
	batch.kernels = NULL;
	batch.count = 0;
	batch.capacity = 0;
	batch.num_qubits = num_qubits;

	return batch;
}

void psi_free_kernel_batch(struct PsiKernelBatch *batch)
{
	for (size_t i = 0; i < batch->count; i++)
		psi_free_kernel(&batch->kernels[i]);

	free(batch->kernels);
	batch->kernels = NULL;
	batch->count = 0;
	batch->capacity = 0;
}

void psi_add_kernel(struct PsiKernelBatch *batch, struct PsiKernel kernel)
{
	if (batch->count == batch->capacity)
	{
		size_t new_capacity = batch->capacity == 0 ? 8 : batch->capacity * 2;
		batch->kernels = realloc(batch->kernels, new_capacity * sizeof(struct PsiKernel));
		assert(batch->kernels != NULL);
		batch->capacity = new_capacity;
	}

	batch->kernels[batch->count++] = kernel;
}

void psi_optimize_kernel_batch(struct PsiKernelBatch *batch)
{
	if (batch->count < 2)
		return;

	size_t original = batch->count;
	struct PsiKernel *out = malloc(original * sizeof(struct PsiKernel));
	assert(out != NULL);
	size_t out_count = 0;

	size_t i = 0;
	while (i < batch->count)
	{
		if (i + 1 < batch->count)
		{
			struct PsiKernel fused;
			if (psi_fuse_kernels(batch->kernels[i], batch->kernels[i + 1], &fused))
			{
				psi_free_kernel(&batch->kernels[i]);
				psi_free_kernel(&batch->kernels[i + 1]);
				out[out_count++] = fused;
				i += 2;
				continue;
			}
		}

		out[out_count++] = batch->kernels[i];
		i += 1;
	}

	free(batch->kernels);
	batch->kernels = out;
	batch->count = out_count;
	batch->capacity = original;
}

void psi_execute_kernel_batch(struct PsiKernelBatch batch, struct PsiVector *state)
{
	size_t dim = (size_t)1 << batch.num_qubits;
	assert(state->size == dim);

	for (size_t i = 0; i < batch.count; i++)
	{
		struct PsiComplex *next = apply_kernel(state->data, batch.kernels[i], batch.num_qubits);
		free(state->data);
		state->data = next;
	}
}