import pytest

import numpy as np

import piquasso as pq
import strawberryfields as sf


pytestmark = pytest.mark.benchmark(
    group="pure-fock",
)


@pytest.fixture
def theta():
    return np.pi / 5


@pytest.fixture
def d():
    return 5


@pytest.mark.parametrize("cutoff", range(3, 14))
def piquasso_benchmark(benchmark, d, cutoff, theta):
    @benchmark
    def func():
        state_vector = [cutoff // d] * d
        state_vector[0] += cutoff % d - 1

        with pq.Program() as program:
            pq.Q(all) | pq.StateVector(state_vector)
            for i in range(d - 1):
                pq.Q(i, i + 1) | pq.Beamsplitter(theta)

        simulator_fock = pq.PureFockSimulator(d=d, config=pq.Config(cutoff=cutoff))

        simulator_fock.execute(program)


@pytest.mark.parametrize("cutoff", (3, 4, 5))
def strawberryfields_benchmark(benchmark, d, cutoff, theta):
    @benchmark
    def func():
        eng = sf.Engine(backend="fock", backend_options={"cutoff_dim": cutoff})

        circuit = sf.Program(d)

        state_vector = [cutoff // d] * d
        state_vector[0] += cutoff % d - 1

        with circuit.context as q:
            for i, n in enumerate(state_vector):
                sf.ops.Fock(n) | q[i]

            for w in range(d - 1):
                sf.ops.BSgate(theta) | (q[w], q[w + 1])

        eng.run(circuit).state
