from math import acos, cos, pi, sqrt

import cirq
import numpy as np


def naive_kitaev_single(phi_real: float, repetitions: int = 2000):
    """
    Наивный вариант алгоритма Китаева
    ---------------------------------------------------
    Реализует схему:
         H --●-- H -- измерение
               |
               U
    где U|1> = e^{i2πφ}|1>.

    После схемы:
        P(0) = cos²(πφ)
        P(1) = sin²(πφ)

    Измеряя P(0), можно получить множество кандидатов для φ из уравнения:
        φ = (±arccos(√P(0)) + mπ) / π
        где m ∈ ℤ, φ ∈ [0,1).
    """

    # Симулятор и кубиты
    sim = cirq.Simulator()
    a, u = cirq.LineQubit.range(2)  # a — фазовый кубит, u — системный

    # Создаём унитар U с фазой φ
    U = cirq.MatrixGate(np.diag([1.0, np.exp(1j * 2 * np.pi * phi_real)]))

    # Собираем схему
    circuit = cirq.Circuit(
        cirq.X(u),  # подготавливаем |ψ⟩ = |1⟩
        cirq.H(a),  # суперпозиция на фазовом кубите
        cirq.ControlledGate(U).on(a, u),
        cirq.H(a),  # интерференция
        cirq.measure(a, key="m"),  # измеряем фазовый кубит
    )

    # Запускаем симуляцию
    result = sim.run(circuit, repetitions=repetitions)
    m = result.measurements["m"][:, 0]

    # Эмпирическая вероятность
    p0_emp = np.mean(m == 0)
    p0_theory = cos(pi * phi_real) ** 2

    # Решаем уравнение cos²(πφ) = p0_emp
    c = sqrt(max(0.0, min(1.0, p0_emp)))  # защита от ошибок округления
    alpha = acos(c)  # угол α ∈ [0, π/2]

    # Все кандидаты φ = (±α + mπ) / π в [0,1)
    candidates = []
    for m in range(2):  # m=0 и m=1 дают все уникальные решения в [0,1)
        for sign in (+1, -1):
            phi_c = (sign * alpha + m * pi) / pi
            phi_c_mod = phi_c % 1.0
            candidates.append((m, sign, phi_c_mod))

    # Удалим дубли и отсортируем
    uniq = {}
    for m, sign, val in candidates:
        key = round(val, 12)
        if key not in uniq:
            uniq[key] = (m, sign, val)
    candidates_unique = sorted(uniq.values(), key=lambda t: t[2])

    print(f"φ(real)      = {phi_real:.6f}")
    print(f"P(0)_emp     = {p0_emp:.6f}")
    print(f"P(0)_theory  = {p0_theory:.6f}")
    print(f"c = sqrt(P0) = {c:.6f},  α = arccos(c) = {alpha:.6f} рад")
    print()
    print("Все кандидаты φ (m, sign, φ ∈ [0,1)):")
    for m, sign, val in candidates_unique:
        s = "+" if sign > 0 else "-"
        print(f"  m={m:1d}, sign={s}  →  φ = {val:.12f}")
    print()
    print(f"repetitions = {repetitions}")

    return candidates_unique


# === Пример запуска ===
if __name__ == "__main__":
    phi_real = 0.228  # истинная фаза
    repetitions = 2000
    naive_kitaev_single(phi_real, repetitions)