#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;
	}
}

void psi_apply_kernel(struct PsiVector *state, struct PsiKernel kernel, size_t num_qubits)
{
	struct PsiComplex *next = apply_kernel(state->data, kernel, num_qubits);
	free(state->data);
	state->data = next;
}

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

	(*kernels)[(*count)++] = kernel;
}

static bool layer_can_add(struct PsiExecutionLayer layer, struct PsiKernel kernel)
{
	for (size_t i = 0; i < layer.count; i++)
		if (psi_kernels_share_qubits(layer.kernels[i], kernel))
			return false;

	return true;
}

static void free_layer(struct PsiExecutionLayer *layer)
{
	for (size_t i = 0; i < layer->count; i++)
		psi_free_kernel(&layer->kernels[i]);

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

struct PsiStructureAwareBatch psi_new_structure_aware_batch(size_t num_qubits)
{
	struct PsiStructureAwareBatch batch;
	batch.kernels = NULL;
	batch.count = 0;
	batch.capacity = 0;
	batch.layers = NULL;
	batch.layer_count = 0;
	batch.layer_capacity = 0;
	batch.num_qubits = num_qubits;
	batch.optimised = false;

	return batch;
}

static void clear_layers(struct PsiStructureAwareBatch *batch)
{
	for (size_t i = 0; i < batch->layer_count; i++)
		free_layer(&batch->layers[i]);

	free(batch->layers);
	batch->layers = NULL;
	batch->layer_count = 0;
	batch->layer_capacity = 0;
}

void psi_free_structure_aware_batch(struct PsiStructureAwareBatch *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;
	clear_layers(batch);
}

void psi_add_structure_aware_kernel(struct PsiStructureAwareBatch *batch, struct PsiKernel kernel)
{
	push_kernel(&batch->kernels, &batch->count, &batch->capacity, kernel);
	batch->optimised = false;
}

static struct PsiKernel remove_kernel_at(struct PsiStructureAwareBatch *batch, size_t index)
{
	struct PsiKernel removed = batch->kernels[index];
	for (size_t i = index; i + 1 < batch->count; i++)
		batch->kernels[i] = batch->kernels[i + 1];

	batch->count--;
	return removed;
}

static void insert_kernel_at(struct PsiStructureAwareBatch *batch, size_t index, 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;
	}

	for (size_t i = batch->count; i > index; i--)
		batch->kernels[i] = batch->kernels[i - 1];

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

static void reorder_commuting_gates(struct PsiStructureAwareBatch *batch)
{
	bool changed = true;
	size_t iterations = 0;
	const size_t MAX_ITERATIONS = 100;

	while (changed && iterations < MAX_ITERATIONS)
	{
		changed = false;
		iterations++;

		for (size_t i = 0; i + 1 < batch->count; i++)
		{
			struct PsiKernel current = batch->kernels[i];
			struct PsiKernel next = batch->kernels[i + 1];

			if (current.target_count != 1 || next.target_count != 1
				|| current.targets[0] == next.targets[0]
				|| !psi_kernels_commute(current, next))
				continue;

			for (size_t j = i + 2; j < batch->count; j++)
			{
				struct PsiKernel candidate = batch->kernels[j];
				if (candidate.target_count != 1 || candidate.targets[0] != current.targets[0])
					continue;

				bool can_move = true;
				for (size_t k = i + 1; k < j; k++)
				{
					struct PsiKernel between = batch->kernels[k];
					if (psi_kernels_share_qubits(between, current) && !psi_kernels_commute(current, between))
					{
						can_move = false;
						break;
					}
				}

				if (can_move && psi_kernels_can_fuse(current, candidate))
				{
					struct PsiKernel moved = remove_kernel_at(batch, j);
					insert_kernel_at(batch, i + 1, moved);
					changed = true;
					break;
				}
			}
		}
	}
}

static void multi_pass_fusion(struct PsiStructureAwareBatch *batch)
{
	bool changed = true;
	size_t iterations = 0;
	const size_t MAX_ITERATIONS = 50;

	while (changed && iterations < MAX_ITERATIONS)
	{
		changed = false;
		iterations++;

		struct PsiKernel *new_kernels = NULL;
		size_t new_count = 0;
		size_t new_capacity = 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]);
					push_kernel(&new_kernels, &new_count, &new_capacity, fused);
					i += 2;
					changed = true;
					continue;
				}
			}

			push_kernel(&new_kernels, &new_count, &new_capacity, batch->kernels[i]);
			i++;
		}

		free(batch->kernels);
		batch->kernels = new_kernels;
		batch->count = new_count;
		batch->capacity = new_capacity;
	}
}

static void build_execution_layers(struct PsiStructureAwareBatch *batch)
{
	clear_layers(batch);

	for (size_t i = 0; i < batch->count; i++)
	{
		struct PsiKernel kernel = batch->kernels[i];
		bool placed = false;

		for (size_t l = 0; l < batch->layer_count; l++)
			if (layer_can_add(batch->layers[l], kernel))
			{
				struct PsiExecutionLayer *layer = &batch->layers[l];
				push_kernel(&layer->kernels, &layer->count, &layer->capacity, psi_clone_kernel(kernel));
				placed = true;
				break;
			}

		if (placed)
			continue;

		if (batch->layer_count == batch->layer_capacity)
		{
			size_t new_capacity = batch->layer_capacity == 0 ? 4 : batch->layer_capacity * 2;
			batch->layers = realloc(batch->layers, new_capacity * sizeof(struct PsiExecutionLayer));
			assert(batch->layers != NULL);
			batch->layer_capacity = new_capacity;
		}

		struct PsiExecutionLayer layer;
		layer.kernels = NULL;
		layer.count = 0;
		layer.capacity = 0;
		push_kernel(&layer.kernels, &layer.count, &layer.capacity, psi_clone_kernel(kernel));
		batch->layers[batch->layer_count++] = layer;
	}
}

void psi_optimize_structure_aware_batch(struct PsiStructureAwareBatch *batch)
{
	if (batch->optimised || batch->count < 2)
		return;

	reorder_commuting_gates(batch);
	multi_pass_fusion(batch);
	build_execution_layers(batch);
	batch->optimised = true;
}

void psi_execute_structure_aware_batch(struct PsiStructureAwareBatch 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;
	}
}

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

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

struct PsiKernelStats psi_structure_aware_batch_stats(struct PsiStructureAwareBatch batch)
{
	struct PsiKernelStats stats;
	stats.total_kernels = batch.count;
	stats.single_qubit = 0;
	stats.two_qubit = 0;
	stats.multi_qubit = 0;
	stats.diagonal = 0;
	stats.execution_layers = batch.layer_count;

	for (size_t i = 0; i < batch.count; i++)
	{
		struct PsiKernel kernel = batch.kernels[i];

		if (kernel.target_count == 1)
			stats.single_qubit++;
		else if (kernel.target_count == 2)
			stats.two_qubit++;
		else if (kernel.target_count > 2)
			stats.multi_qubit++;

		if (kernel.gate_type == PSI_GATE_TYPE_DIAGONAL)
			stats.diagonal++;
	}

	return stats;
}