import cirq

# Кубиты: два входных (q0, q1) и целевой вспомогательный (q2)
q0, q1, q2 = cirq.LineQubit.range(3)

# Оракулы для константных функций: f(x)=0 и f(x)=1
constant = (
    [],  # ничего не делаем с целевым
    [cirq.X(q2)],  # инвертируем целевой
)

# Оракулы для сбалансированных функций
balanced = (
    [cirq.CNOT(q0, q2)],  # f=x0
    [cirq.CNOT(q1, q2)],  # f=x1
    [cirq.CNOT(q0, q2), cirq.CNOT(q1, q2)],
    [cirq.CNOT(q0, q2), cirq.X(q2)],
    [cirq.CNOT(q1, q2), cirq.X(q2)],
    [cirq.CNOT(q0, q2), cirq.CNOT(q1, q2), cirq.X(q2)],
)


def deutsch_jozsa_circuit(oracle):
    """Схема DJ через yield: измеряем только q0 и q1."""
    # 1) Входы в равномерную суперпозицию
    yield cirq.H(q0), cirq.H(q1)

    # 2) Подготовка целевого в |-> для фазового трюка
    yield cirq.X(q2), cirq.H(q2)

    # 3) Вызов оракула U_f
    yield oracle

    # 4) Интерференция на входах
    yield cirq.H(q0), cirq.H(q1)

    # 5) Измеряем только входные кубиты
    yield cirq.measure(q0, q1, key="in")


def interpret(bits):
    # bits — это массив вида [q0, q1]
    return "constant" if (bits[0] == 0 and bits[1] == 0) else "balanced"


# Симулятор
sim = cirq.Simulator()

print("Результаты для константных функций:")
for i, oracle in enumerate(constant):
    result = sim.run(cirq.Circuit(deutsch_jozsa_circuit(oracle)), repetitions=1)
    bits = result.measurements["in"][0]
    print(f"  Константный оракул {i}: meas(q0,q1)={bits.tolist()} -> {interpret(bits)}")

print("\nРезультаты для сбалансированных функций:")
for i, oracle in enumerate(balanced):
    result = sim.run(cirq.Circuit(deutsch_jozsa_circuit(oracle)), repetitions=1)
    bits = result.measurements["in"][0]
    print(
        f"  Сбалансированный оракул {i}: meas(q0,q1)={bits.tolist()} -> {interpret(bits)}"
    )