#include "core/gates.h"

#include <math.h>

static const double INV_SQRT_2 = 0.7071067811865476;

struct PsiMatrix psi_rx_matrix(double theta)
{
	double cos_h = cos(theta / 2.0);
	double sin_h = sin(theta / 2.0);

	return psi_matrix(2, 2,
		psi_new_complex(cos_h, 0.0), psi_new_complex(0.0, -sin_h),
		psi_new_complex(0.0, -sin_h), psi_new_complex(cos_h, 0.0));
}

struct PsiMatrix psi_ry_matrix(double theta)
{
	double cos_h = cos(theta / 2.0);
	double sin_h = sin(theta / 2.0);

	return psi_matrix(2, 2,
		psi_new_complex(cos_h, 0.0), psi_new_complex(-sin_h, 0.0),
		psi_new_complex(sin_h, 0.0), psi_new_complex(cos_h, 0.0));
}

struct PsiMatrix psi_rz_matrix(double theta)
{
	double half = theta / 2.0;

	return psi_matrix(2, 2,
		psi_new_complex(cos(half), -sin(half)), psi_new_complex(0.0, 0.0),
		psi_new_complex(0.0, 0.0), psi_new_complex(cos(half), sin(half)));
}

struct PsiMatrix psi_p_matrix(double theta)
{
	return 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(cos(theta), sin(theta)));
}

struct PsiMatrix psi_u1_matrix(double lambda)
{
	return psi_p_matrix(lambda);
}

struct PsiMatrix psi_u2_matrix(double phi, double lambda)
{
	return psi_matrix(2, 2,
		psi_new_complex(INV_SQRT_2, 0.0),
		psi_new_complex(-INV_SQRT_2 * cos(lambda), -INV_SQRT_2 * sin(lambda)),
		psi_new_complex(INV_SQRT_2 * cos(phi), INV_SQRT_2 * sin(phi)),
		psi_new_complex(cos(phi + lambda) * INV_SQRT_2, sin(phi + lambda) * INV_SQRT_2));
}

struct PsiMatrix psi_u3_matrix(double theta, double phi, double lambda)
{
	double cos_h = cos(theta / 2.0);
	double sin_h = sin(theta / 2.0);

	return psi_matrix(2, 2,
		psi_new_complex(cos_h, 0.0),
		psi_new_complex(-sin_h * cos(lambda), -sin_h * sin(lambda)),
		psi_new_complex(sin_h * cos(phi), sin_h * sin(phi)),
		psi_new_complex(cos_h * cos(phi + lambda), cos_h * sin(phi + lambda)));
}

struct PsiMatrix psi_crx_matrix(double theta)
{
	double cos_h = cos(theta / 2.0);
	double sin_h = sin(theta / 2.0);

	struct PsiMatrix m = psi_identity_matrix(4);
	psi_set_matrix(&m, 2, 2, psi_new_complex(cos_h, 0.0));
	psi_set_matrix(&m, 2, 3, psi_new_complex(0.0, -sin_h));
	psi_set_matrix(&m, 3, 2, psi_new_complex(0.0, -sin_h));
	psi_set_matrix(&m, 3, 3, psi_new_complex(cos_h, 0.0));

	return m;
}

struct PsiMatrix psi_cry_matrix(double theta)
{
	double cos_h = cos(theta / 2.0);
	double sin_h = sin(theta / 2.0);

	struct PsiMatrix m = psi_identity_matrix(4);
	psi_set_matrix(&m, 2, 2, psi_new_complex(cos_h, 0.0));
	psi_set_matrix(&m, 2, 3, psi_new_complex(-sin_h, 0.0));
	psi_set_matrix(&m, 3, 2, psi_new_complex(sin_h, 0.0));
	psi_set_matrix(&m, 3, 3, psi_new_complex(cos_h, 0.0));

	return m;
}

struct PsiMatrix psi_crz_matrix(double theta)
{
	double half = theta / 2.0;

	struct PsiMatrix m = psi_identity_matrix(4);
	psi_set_matrix(&m, 2, 2, psi_new_complex(cos(half), -sin(half)));
	psi_set_matrix(&m, 3, 3, psi_new_complex(cos(half), sin(half)));

	return m;
}

struct PsiMatrix psi_cp_matrix(double theta)
{
	struct PsiMatrix m = psi_identity_matrix(4);
	psi_set_matrix(&m, 3, 3, psi_new_complex(cos(theta), sin(theta)));

	return m;
}

struct PsiQuantumGate psi_hadamard_gate(void)
{
	struct PsiMatrix m = psi_matrix(2, 2,
		psi_new_complex(INV_SQRT_2, 0.0), psi_new_complex(INV_SQRT_2, 0.0),
		psi_new_complex(INV_SQRT_2, 0.0), psi_new_complex(-INV_SQRT_2, 0.0));

	return psi_new_quantum_gate_from_matrix("H", m);
}

struct PsiQuantumGate psi_pauli_x_gate(void)
{
	struct PsiMatrix m = psi_matrix(2, 2,
		psi_new_complex(0.0, 0.0), psi_new_complex(1.0, 0.0),
		psi_new_complex(1.0, 0.0), psi_new_complex(0.0, 0.0));

	return psi_new_quantum_gate_from_matrix("X", m);
}

struct PsiQuantumGate psi_pauli_y_gate(void)
{
	struct PsiMatrix m = psi_matrix(2, 2,
		psi_new_complex(0.0, 0.0), psi_new_complex(0.0, -1.0),
		psi_new_complex(0.0, 1.0), psi_new_complex(0.0, 0.0));

	return psi_new_quantum_gate_from_matrix("Y", m);
}

struct PsiQuantumGate psi_pauli_z_gate(void)
{
	struct PsiMatrix m = 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(-1.0, 0.0));

	return psi_new_quantum_gate_from_matrix("Z", m);
}

struct PsiQuantumGate psi_s_gate(void)
{
	struct PsiMatrix m = 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(0.0, 1.0));

	return psi_new_quantum_gate_from_matrix("S", m);
}

struct PsiQuantumGate psi_t_gate(void)
{
	struct PsiMatrix m = 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(INV_SQRT_2, INV_SQRT_2));

	return psi_new_quantum_gate_from_matrix("T", m);
}

struct PsiQuantumGate psi_sdg_gate(void)
{
	struct PsiMatrix m = 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(0.0, -1.0));

	return psi_new_quantum_gate_from_matrix("S†", m);
}

struct PsiQuantumGate psi_tdg_gate(void)
{
	struct PsiMatrix m = 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(INV_SQRT_2, -INV_SQRT_2));

	return psi_new_quantum_gate_from_matrix("T†", m);
}

struct PsiQuantumGate psi_sx_gate(void)
{
	struct PsiMatrix m = psi_matrix(2, 2,
		psi_new_complex(0.5, 0.5), psi_new_complex(0.5, -0.5),
		psi_new_complex(0.5, -0.5), psi_new_complex(0.5, 0.5));

	return psi_new_quantum_gate_from_matrix("√X", m);
}

struct PsiQuantumGate psi_sxdg_gate(void)
{
	struct PsiMatrix m = psi_matrix(2, 2,
		psi_new_complex(0.5, -0.5), psi_new_complex(0.5, 0.5),
		psi_new_complex(0.5, 0.5), psi_new_complex(0.5, -0.5));

	return psi_new_quantum_gate_from_matrix("√X†", m);
}

struct PsiQuantumGate psi_identity_gate(void)
{
	return psi_new_quantum_gate_from_matrix("I", psi_identity_matrix(2));
}

struct PsiQuantumGate psi_cnot_gate(void)
{
	struct PsiMatrix m = psi_identity_matrix(4);
	psi_set_matrix(&m, 2, 2, psi_new_complex(0.0, 0.0));
	psi_set_matrix(&m, 3, 3, psi_new_complex(0.0, 0.0));
	psi_set_matrix(&m, 2, 3, psi_new_complex(1.0, 0.0));
	psi_set_matrix(&m, 3, 2, psi_new_complex(1.0, 0.0));

	return psi_new_quantum_gate_from_matrix("CNOT", m);
}

struct PsiQuantumGate psi_cz_gate(void)
{
	struct PsiMatrix m = psi_identity_matrix(4);
	psi_set_matrix(&m, 3, 3, psi_new_complex(-1.0, 0.0));

	return psi_new_quantum_gate_from_matrix("CZ", m);
}

struct PsiQuantumGate psi_swap_gate(void)
{
	struct PsiMatrix m = psi_identity_matrix(4);
	psi_set_matrix(&m, 1, 1, psi_new_complex(0.0, 0.0));
	psi_set_matrix(&m, 2, 2, psi_new_complex(0.0, 0.0));
	psi_set_matrix(&m, 1, 2, psi_new_complex(1.0, 0.0));
	psi_set_matrix(&m, 2, 1, psi_new_complex(1.0, 0.0));

	return psi_new_quantum_gate_from_matrix("SWAP", m);
}

struct PsiQuantumGate psi_iswap_gate(void)
{
	struct PsiMatrix m = psi_identity_matrix(4);
	psi_set_matrix(&m, 1, 1, psi_new_complex(0.0, 0.0));
	psi_set_matrix(&m, 2, 2, psi_new_complex(0.0, 0.0));
	psi_set_matrix(&m, 1, 2, psi_new_complex(0.0, 1.0));
	psi_set_matrix(&m, 2, 1, psi_new_complex(0.0, 1.0));

	return psi_new_quantum_gate_from_matrix("iSWAP", m);
}

struct PsiQuantumGate psi_sqrt_swap_gate(void)
{
	struct PsiMatrix m = psi_identity_matrix(4);
	psi_set_matrix(&m, 1, 1, psi_new_complex(0.5, 0.5));
	psi_set_matrix(&m, 1, 2, psi_new_complex(0.5, -0.5));
	psi_set_matrix(&m, 2, 1, psi_new_complex(0.5, -0.5));
	psi_set_matrix(&m, 2, 2, psi_new_complex(0.5, 0.5));

	return psi_new_quantum_gate_from_matrix("√SWAP", m);
}

struct PsiQuantumGate psi_toffoli_gate(void)
{
	struct PsiMatrix m = psi_identity_matrix(8);
	psi_set_matrix(&m, 6, 6, psi_new_complex(0.0, 0.0));
	psi_set_matrix(&m, 7, 7, psi_new_complex(0.0, 0.0));
	psi_set_matrix(&m, 6, 7, psi_new_complex(1.0, 0.0));
	psi_set_matrix(&m, 7, 6, psi_new_complex(1.0, 0.0));

	return psi_new_quantum_gate_from_matrix("CCNOT", m);
}

struct PsiQuantumGate psi_fredkin_gate(void)
{
	struct PsiMatrix m = psi_identity_matrix(8);
	psi_set_matrix(&m, 5, 5, psi_new_complex(0.0, 0.0));
	psi_set_matrix(&m, 6, 6, psi_new_complex(0.0, 0.0));
	psi_set_matrix(&m, 5, 6, psi_new_complex(1.0, 0.0));
	psi_set_matrix(&m, 6, 5, psi_new_complex(1.0, 0.0));

	return psi_new_quantum_gate_from_matrix("CSWAP", m);
}