r"""Demonstration of the generation of IPEA unitaries for statevector simulation."""

# `IterativePhaseEstimation` class may be used for designing
# the user-defined iterative QPE algorithm.
# Here we demonstrate how to use this feature with the Hadamard test as an example.
#
# Although the Hadamard test is not a stochastic QPE algorithm,
# It uses the same type of circuit and thus suitable for the simplest demonstration of QPE circuit generation.
# We calculate the phase by
# $$
#   \phi = \mathrm{arg}[\langle \phi | e^{i\phi} | \phi \rangle]
# $$
# The real and the imaginary parts in the square bracket is evaluated
# by setting $\beta = 0$ and $\beta=-\frac{\pi}{2}$, respectively.

# imports
import numpy as np

from typing import Callable

from pytket.circuit import (
    Circuit,
    OpType,
)

from inquanto.protocols import IterativePhaseEstimationStatevector


# Computational setup.

# Prepare a random phase to be used in Ctrl-U construction
# e0 = 2 * np.random.random()
e0 = 1.198059919365895


### Circuit preparation
# Define simple method to obtain Ctrl-U circuit gadget.
def get_ctrlu(k: int):
    circ = Circuit(2)
    for _ in range(k):
        circ.add_gate(OpType.CU1, e0, (0, 1))
    return circ


ctrlu = get_ctrlu(k=1)

# State preparation circuit.
state = Circuit(1).X(0)

# instance the IQPE statevector protocol
# note that this protocol does not accept backends
# and will perform the algebra internally
protocol = IterativePhaseEstimationStatevector()
# build protocol with the state and function to build ctrl-u circuits
protocol.build_from_circuit(
    state=state,
    get_ctrlu=get_ctrlu,
)


def measure(k: int, beta: float) -> Circuit:
    protocol.update_k_and_beta(k=k, beta=beta)
    # perform statevector algebra to produce probability dists
    protocol.run()
    p0 = protocol.get_distribution()[(0,)]
    return p0


###
# use the measure_function to prepare and run circuits for a range of k and beta
p0s = [measure(k=1, beta=beta) for beta in (0.0, -0.5)]

###
# analyze outputs
pcos, psin = [(2 * m - 1) for m in p0s]
phase = (np.arctan2(psin, pcos) / np.pi) % 2.0
estimate = phase
print(f"P_cos, P_sin     = {pcos:.4f}, {psin:.4f}")
print(f"Energy estimate  = {estimate:.10f} hartree")
print(f"Reference energy = {e0:.10f} hartree")