r"""Use HadamardTest with a chemistry ansatz and a circuit kernel."""

import numpy as np
from pytket import Circuit
from pytket.extensions.qiskit import AerBackend

from inquanto.ansatzes import FermionSpaceAnsatzUCCSD
from inquanto.operators import QubitOperator
from inquanto.protocols import HadamardTest
from inquanto.spaces import FermionSpace
from inquanto.states import FermionState

backend = AerBackend()

space = FermionSpace(4)
reference_state = FermionState([1, 1, 0, 0], 1)
ansatz = FermionSpaceAnsatzUCCSD(space, reference_state)
parameters = ansatz.state_symbols.construct_random(seed=0, mu=0.0, sigma=0.1)

state_circuit = ansatz.get_circuit(parameters)
statevector = state_circuit.get_statevector()

# H(0) Z(0) H(0) is equivalent to X0.
kernel_circuit = Circuit(state_circuit.n_qubits).H(0).Z(0).H(0)
kernel_operator = QubitOperator("X0")
circuit_reference = np.vdot(statevector, kernel_circuit.get_unitary() @ statevector)
operator_reference = kernel_operator.state_expectation(statevector)
print(
    "Exact circuit/operator agreement: "
    f"{np.isclose(circuit_reference, operator_reference, atol=1e-12, rtol=0.0)}"
)

protocol_circuit = HadamardTest(backend, shots_per_circuit=50000)
protocol_circuit.build(parameters, ansatz, kernel_circuit, component="real")
# Disable preoptimization because this chemistry ansatz expands to CircBox blocks,
# and the default GuidedPauliSimp pass cannot process them.
protocol_circuit.compile_circuits(preoptimize_passes=False)
protocol_circuit.run(seed=0)
circuit_value = protocol_circuit.evaluate_expectation_value()[0]

protocol_operator = HadamardTest(backend, shots_per_circuit=50000)
protocol_operator.build(parameters, ansatz, kernel_operator)
# Same reason as above.
protocol_operator.compile_circuits(preoptimize_passes=False)
protocol_operator.run(seed=0)
operator_value = protocol_operator.evaluate_expectation_value(ansatz, kernel_operator)

print(f"HadamardTest from circuit kernel: {circuit_value}")
print(f"HadamardTest from qubit operator: {operator_value}")
print(
    "Agreement within shot noise: "
    f"{np.isclose(circuit_value, operator_value, atol=0.03, rtol=0.0)}"
)  # Loose comparison due to shot noise making a tighter comparison unreliable here.