#include <stdio.h>

#include "psi.h"

int main(void)
{
	printf("psi %s\n", psi_version());

	struct PsiComplex a = psi_new_complex(1.0, 2.0);
	struct PsiComplex b = psi_new_complex(3.0, -1.0);

	struct PsiComplex sum = psi_add_complex(a, b);
	struct PsiComplex product = psi_mul_complex(a, b);

	printf("a       = %g + %gi\n", a.real, a.imaginary);
	printf("b       = %g + %gi\n", b.real, b.imaginary);
	printf("a + b   = %g + %gi\n", sum.real, sum.imaginary);
	printf("a * b   = %g + %gi\n", product.real, product.imaginary);
	printf("|a|     = %g\n", psi_abs_complex(a));

	struct PsiVector u = psi_column_vector(
		psi_new_complex(1.0, 0.0),
		psi_new_complex(2.0, 0.0),
		psi_new_complex(3.0, 0.0));
	struct PsiVector v = psi_column_vector(
		psi_new_complex(4.0, 0.0),
		psi_new_complex(5.0, 0.0),
		psi_new_complex(6.0, 0.0));

	struct PsiComplex dot = psi_dot_vector(u, v);
	struct PsiComplex norm = psi_norm_vector(u);

	printf("u . v   = %g + %gi\n", dot.real, dot.imaginary);
	printf("|u|     = %g + %gi\n", norm.real, norm.imaginary);

	psi_free_vector(&u);
	psi_free_vector(&v);

	struct PsiMatrix m = psi_matrix(2, 2,
		psi_new_complex(1.0, 0.0), psi_new_complex(2.0, 0.0),
		psi_new_complex(3.0, 0.0), psi_new_complex(4.0, 0.0));
	struct PsiMatrix mm = psi_dot_matrix(m, m);

	printf("m*m     = [%g %g; %g %g]\n",
		psi_get_matrix(mm, 0, 0).real, psi_get_matrix(mm, 0, 1).real,
		psi_get_matrix(mm, 1, 0).real, psi_get_matrix(mm, 1, 1).real);

	struct PsiVector x = psi_column_vector(
		psi_new_complex(1.0, 0.0),
		psi_new_complex(1.0, 0.0));
	struct PsiVector mx = psi_mul_vector_matrix(x, m);

	printf("m*x     = [%g, %g]\n", mx.data[0].real, mx.data[1].real);

	psi_free_vector(&x);
	psi_free_vector(&mx);
	psi_free_matrix(&m);
	psi_free_matrix(&mm);

	double inv_sqrt2 = 0.7071067811865475;
	struct PsiMatrix h = psi_matrix(2, 2,
		psi_new_complex(inv_sqrt2, 0.0), psi_new_complex(inv_sqrt2, 0.0),
		psi_new_complex(inv_sqrt2, 0.0), psi_new_complex(-inv_sqrt2, 0.0));
	struct PsiQuantumGate gate = psi_new_quantum_gate_from_matrix("H", h);

	struct PsiVector s0 = psi_new_state_0();
	struct PsiVector hs0 = psi_mul_vector_matrix(s0, gate.matrix);

	printf("gate    = %s (%zu qubit)\n", gate.name, gate.num_qubits);
	printf("H|0>    = [%g, %g]\n", hs0.data[0].real, hs0.data[1].real);

	psi_free_vector(&s0);
	psi_free_vector(&hs0);
	psi_free_quantum_gate(&gate);

	const char *names[] = { "q0", "q1" };
	struct PsiQuantumRegister reg = psi_new_quantum_register("bell", names, 2);

	struct PsiQuantumGate hadamard = psi_hadamard_gate();
	struct PsiQuantumGate cnot = psi_cnot_gate();

	size_t t0[] = { 0 };
	size_t t01[] = { 0, 1 };
	psi_apply_gate(&reg, hadamard, t0, 1);
	psi_apply_gate(&reg, cnot, t01, 2);

	printf("bell    = [%g, %g, %g, %g]\n",
		reg.state_vector.data[0].real, reg.state_vector.data[1].real,
		reg.state_vector.data[2].real, reg.state_vector.data[3].real);

	psi_free_quantum_gate(&hadamard);
	psi_free_quantum_gate(&cnot);
	psi_free_quantum_register(&reg);

	struct PsiCompositeGateOp bell_ops[] = {
		{ PSI_OP_H, { 0 }, 1 },
		{ PSI_OP_CNOT, { 0, 1 }, 2 },
	};
	struct PsiCustomGate bell = psi_new_custom_gate_from_composite("BELL", 2, bell_ops, 2);
	struct PsiQuantumGate bell_gate = psi_to_quantum_gate(bell);

	struct PsiQuantumRegister reg2 = psi_new_quantum_register("bell2", names, 2);
	size_t t01b[] = { 0, 1 };
	psi_apply_gate(&reg2, bell_gate, t01b, 2);

	printf("custom  = %s -> [%g, %g, %g, %g]\n", bell_gate.name,
		reg2.state_vector.data[0].real, reg2.state_vector.data[1].real,
		reg2.state_vector.data[2].real, reg2.state_vector.data[3].real);

	psi_free_custom_gate(&bell);
	psi_free_quantum_gate(&bell_gate);
	psi_free_quantum_register(&reg2);

	const char *cnames[] = { "c0", "c1" };
	struct PsiClassicalRegister creg = psi_new_classical_register("c", cnames, 2);
	creg.bits[1].state = true;

	printf("creg    = %s=%d %s=%d\n",
		creg.bits[0].name, creg.bits[0].state,
		creg.bits[1].name, creg.bits[1].state);

	psi_free_classical_register(&creg);

	struct PsiQuantumCircuit circuit = psi_new_quantum_circuit(3);
	psi_apply_h(&circuit, 0);
	psi_apply_cnot(&circuit, 0, 1);
	psi_apply_rx(&circuit, 2, 1.5707963267948966);
	psi_apply_ccnot(&circuit, 0, 1, 2);

	struct PsiCompositeGateOp cops[] = { { PSI_OP_X, { 0 }, 1 } };
	struct PsiCustomGate xg = psi_new_custom_gate_from_composite("MYX", 1, cops, 1);
	size_t ct[] = { 2 };
	psi_apply_custom(&circuit, xg, ct, 1);

	psi_measure_all(&circuit);

	printf("circuit = %zu qubits, %zu classical, %zu ops\n",
		circuit.num_qubits, circuit.num_classical, circuit.operation_count);
	for (size_t i = 0; i < circuit.operation_count; i++)
	{
		size_t tc;
		const size_t *targets = psi_gate_op_quantum_targets(&circuit.operations[i], &tc);
		printf("  %zu: %s on", i, psi_gate_op_name(circuit.operations[i]));
		for (size_t j = 0; j < tc; j++)
			printf(" q%zu", targets[j]);
		printf("\n");
	}

	psi_free_quantum_circuit(&circuit);
}