#include "core/runtime.h"

#include "core/custom_gate.h"
#include "core/gates.h"
#include "core/kernel.h"
#include "maths/complex.h"
#include "maths/simd.h"

struct PsiRuntimeConfig psi_new_runtime_config(void)
{
	struct PsiRuntimeConfig config;
	config.parallel = false;
	config.simd = false;
	config.batched = false;
	config.structure_aware = false;
	config.parallel_threshold = PSI_PARALLEL_THRESHOLD;

	return config;
}

struct PsiRuntimeConfig psi_optimal_runtime_config(void)
{
	struct PsiRuntimeConfig config = psi_new_runtime_config();
	config.structure_aware = true;
	config.simd = true;
	config.parallel = true;

	return config;
}

struct PsiRuntimeConfig psi_runtime_to_config(enum PsiRuntime runtime)
{
	struct PsiRuntimeConfig config = psi_new_runtime_config();

	switch (runtime)
	{
		case PSI_RUNTIME_BASIC: break;
		case PSI_RUNTIME_BASIC_MT: config.parallel = true; break;
		case PSI_RUNTIME_BATCHED: config.batched = true; break;
		case PSI_RUNTIME_BATCHED_MT:
			config.batched = true;
			config.parallel = true;
			break;
		case PSI_RUNTIME_SIMD:
			config.batched = true;
			config.simd = true;
			break;
		case PSI_RUNTIME_SIMD_MT:
			config.batched = true;
			config.simd = true;
			config.parallel = true;
			break;
		case PSI_RUNTIME_STRUCTURE_AWARE:
			config.structure_aware = true;
			config.simd = true;
			break;
		case PSI_RUNTIME_STRUCTURE_AWARE_MT:
			config.structure_aware = true;
			config.simd = true;
			config.parallel = true;
			break;
	}

	return config;
}

static bool op_to_kernel(struct PsiGateOp op, struct PsiKernel* out)
{
	struct PsiMatrix matrix;
	const char* name;

	switch (op.kind)
	{
		case PSI_GATE_H:
			matrix = psi_hadamard_gate().matrix;
			name = "H";
			break;
		case PSI_GATE_X:
			matrix = psi_pauli_x_gate().matrix;
			name = "X";
			break;
		case PSI_GATE_Y:
			matrix = psi_pauli_y_gate().matrix;
			name = "Y";
			break;
		case PSI_GATE_Z:
			matrix = psi_pauli_z_gate().matrix;
			name = "Z";
			break;
		case PSI_GATE_S:
			matrix = psi_s_gate().matrix;
			name = "S";
			break;
		case PSI_GATE_T:
			matrix = psi_t_gate().matrix;
			name = "T";
			break;
		case PSI_GATE_SDG:
			matrix = psi_sdg_gate().matrix;
			name = "Sdg";
			break;
		case PSI_GATE_TDG:
			matrix = psi_tdg_gate().matrix;
			name = "Tdg";
			break;
		case PSI_GATE_SX:
			matrix = psi_sx_gate().matrix;
			name = "Sx";
			break;
		case PSI_GATE_SXDG:
			matrix = psi_sxdg_gate().matrix;
			name = "Sxdg";
			break;
		case PSI_GATE_RX:
			matrix = psi_rx_matrix(op.params[0]);
			name = "Rx";
			break;
		case PSI_GATE_RY:
			matrix = psi_ry_matrix(op.params[0]);
			name = "Ry";
			break;
		case PSI_GATE_RZ:
			matrix = psi_rz_matrix(op.params[0]);
			name = "Rz";
			break;
		case PSI_GATE_P:
			matrix = psi_p_matrix(op.params[0]);
			name = "P";
			break;
		case PSI_GATE_U1:
			matrix = psi_u1_matrix(op.params[0]);
			name = "U1";
			break;
		case PSI_GATE_U2:
			matrix = psi_u2_matrix(op.params[0], op.params[1]);
			name = "U2";
			break;
		case PSI_GATE_U3:
			matrix = psi_u3_matrix(op.params[0], op.params[1], op.params[2]);
			name = "U3";
			break;
		case PSI_GATE_CNOT:
			matrix = psi_cnot_gate().matrix;
			name = "CNOT";
			break;
		case PSI_GATE_CZ:
			matrix = psi_cz_gate().matrix;
			name = "CZ";
			break;
		case PSI_GATE_SWAP:
			matrix = psi_swap_gate().matrix;
			name = "SWAP";
			break;
		case PSI_GATE_CRX:
			matrix = psi_crx_matrix(op.params[0]);
			name = "CRx";
			break;
		case PSI_GATE_CRY:
			matrix = psi_cry_matrix(op.params[0]);
			name = "CRy";
			break;
		case PSI_GATE_CRZ:
			matrix = psi_crz_matrix(op.params[0]);
			name = "CRz";
			break;
		case PSI_GATE_CP:
			matrix = psi_cp_matrix(op.params[0]);
			name = "CP";
			break;
		case PSI_GATE_CCNOT:
			matrix = psi_toffoli_gate().matrix;
			name = "CCNOT";
			break;
		case PSI_GATE_CSWAP:
			matrix = psi_fredkin_gate().matrix;
			name = "CSWAP";
			break;
		case PSI_GATE_MEASURE: return false;
		case PSI_GATE_CUSTOM:
			matrix = psi_to_quantum_gate(*op.custom).matrix;
			name = "Custom";
			break;
	}

	size_t target_count;
	const size_t* targets = psi_gate_op_quantum_targets(&op, &target_count);
	*out = psi_new_kernel(name, matrix, targets, target_count);

	return true;
}

static struct PsiVector new_zero_state(size_t num_qubits)
{
	size_t dim = (size_t)1 << num_qubits;
	struct PsiVector state = psi_new_vector(dim, PSI_COLUMN_VECTOR);
	state.data[0] = psi_new_complex(1.0, 0.0);

	return state;
}

static void execute_kernels(struct PsiVector* state, const struct PsiKernel* kernels, size_t count,
                            size_t num_qubits, struct PsiRuntimeConfig config)
{
	bool use_parallel = config.parallel && num_qubits >= config.parallel_threshold;

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

		if (config.simd && kernel.target_count == 1)
		{
			struct PsiComplex gate[2][2] = {
				{ kernel.matrix.data[0], kernel.matrix.data[1] },
				{ kernel.matrix.data[2], kernel.matrix.data[3] },
			};

			if (use_parallel)
				psi_apply_single_qubit_gate_simd_parallel(state->data, gate, kernel.targets[0],
				                                          num_qubits);
			else
				psi_apply_single_qubit_gate_simd(state->data, gate, kernel.targets[0], num_qubits);
		}
		else
			psi_apply_kernel(state, kernel, num_qubits);
	}
}

struct PsiVector psi_compute_runtime_config(struct PsiRuntimeConfig config, size_t num_qubits,
                                            const struct PsiGateOp* operations, size_t op_count)
{
	struct PsiVector state = new_zero_state(num_qubits);

	if (config.structure_aware)
	{
		struct PsiStructureAwareBatch batch = psi_new_structure_aware_batch(num_qubits);
		for (size_t i = 0; i < op_count; i++)
		{
			struct PsiKernel kernel;
			if (op_to_kernel(operations[i], &kernel))
				psi_add_structure_aware_kernel(&batch, kernel);
		}

		psi_optimize_structure_aware_batch(&batch);
		execute_kernels(&state, batch.kernels, batch.count, num_qubits, config);
		psi_free_structure_aware_batch(&batch);

		return state;
	}

	struct PsiKernelBatch batch = psi_new_kernel_batch(num_qubits);
	for (size_t i = 0; i < op_count; i++)
	{
		struct PsiKernel kernel;
		if (op_to_kernel(operations[i], &kernel))
			psi_add_kernel(&batch, kernel);
	}

	if (config.batched)
		psi_optimize_kernel_batch(&batch);

	execute_kernels(&state, batch.kernels, batch.count, num_qubits, config);
	psi_free_kernel_batch(&batch);

	return state;
}

struct PsiVector psi_compute_runtime(enum PsiRuntime runtime, size_t num_qubits,
                                     const struct PsiGateOp* operations, size_t op_count)
{
	return psi_compute_runtime_config(psi_runtime_to_config(runtime), num_qubits, operations,
	                                  op_count);
}

const struct PsiVector* psi_compute_circuit_with_config(struct PsiQuantumCircuit* circuit,
                                                        struct PsiRuntimeConfig config)
{
	if (!circuit->is_computed)
	{
		psi_free_vector(&circuit->computed_state);
		circuit->computed_state = psi_compute_runtime_config(
				config, circuit->num_qubits, circuit->operations, circuit->operation_count);
		circuit->is_computed = true;
	}

	return &circuit->computed_state;
}

const struct PsiVector* psi_compute_circuit_with(struct PsiQuantumCircuit* circuit,
                                                 enum PsiRuntime runtime)
{
	return psi_compute_circuit_with_config(circuit, psi_runtime_to_config(runtime));
}

const struct PsiVector* psi_compute_circuit(struct PsiQuantumCircuit* circuit)
{
	return psi_compute_circuit_with(circuit, PSI_RUNTIME_BASIC);
}

double psi_circuit_probability(struct PsiQuantumCircuit* circuit, size_t state_index)
{
	const struct PsiVector* state = psi_compute_circuit(circuit);
	return psi_norm2_complex(state->data[state_index]);
}

void psi_circuit_probabilities(struct PsiQuantumCircuit* circuit, double* out)
{
	const struct PsiVector* state = psi_compute_circuit(circuit);
	size_t dim = (size_t)1 << circuit->num_qubits;
	for (size_t i = 0; i < dim; i++)
		out[i] = psi_norm2_complex(state->data[i]);
}