r"""BQPE calculation using a noise aware prior."""
import numpy as np
import phayes
from pytket.circuit import (
    Circuit,
    Pauli,
    PauliExpBox,
)
from pytket.circuit.display import render_circuit_jupyter
from inquanto.operators import QubitOperator
from inquanto.ansatzes import CircuitAnsatz
from inquanto.protocols import (
    IterativePhaseEstimationQuantinuum,
    CircuitEncoderQuantinuum,
    CompilationLevelQuantinuum,
)
from inquanto.extensions.phayes import AlgorithmBayesianQPE


# ### Input data


# Qubit operator.
# Two-qubit H2 with the equilibrium geometry.
qop = QubitOperator.from_string(
    "(-0.398, Z0), (-0.398, Z1), (-0.1809, Y0 Y1)",
)
qop_totally_commuting = QubitOperator.from_string(
    "(0.0112, Z0 Z1), (-0.3322, )",
)
fci_energy = np.linalg.eigh(
    (qop + qop_totally_commuting).to_sparse_matrix(qubits=2).todense()
)[0][0]
print(qop.df())
print(qop_totally_commuting.df())
print(f"FCI energy = {fci_energy:.5f} hartree")


# Parameters for constructing a function to return controlled unitary.
time = 0.1
n_trotter = 1
evo_ope_exp = qop.trotterize(trotter_number=n_trotter) * time
eoe_tot_com = qop_totally_commuting.trotterize(trotter_number=n_trotter) * time
time_split = [-0.5, 2.5]
syndrome_interval = 6


# State preparation circuit.
state = Circuit(2)
state.add_pauliexpbox(
    PauliExpBox([Pauli.Y, Pauli.X], -0.07113),
    state.qubits,
)
render_circuit_jupyter(state)
ansatz = CircuitAnsatz(state)


def noisy_backend(cx_err, ro_err, n_qubits):
    """Return noisy backend."""
    from pytket.extensions.qiskit import AerBackend
    from qiskit_aer.noise import NoiseModel, ReadoutError
    import qiskit_aer.noise as noise

    noise_model = NoiseModel()
    # Add CX error.
    error_1 = noise.depolarizing_error(cx_err, 2)
    for i in range(n_qubits):
        for j in range(n_qubits):
            if i != j:
                noise_model.add_quantum_error(error_1, ["cx"], [i, j])
    # Add RO error.
    probabilities = [[1 - ro_err, ro_err], [ro_err, 1 - ro_err]]
    readout_error = ReadoutError(probabilities)
    for i in range(n_qubits):
        noise_model.add_readout_error(readout_error, qubits=[i])
    # Construct noisy backend.
    backend = AerBackend(noise_model=noise_model)
    return backend


p2 = 2.5e-4

backend = noisy_backend(
    cx_err=p2,
    ro_err=0.0,
    n_qubits=8,
)


# Prepare the protocol.
# compilation_level=CompilationLevelQuantinuum.LOGICAL
# compilation_level=CompilationLevelQuantinuum.ENCODED
compilation_level = CompilationLevelQuantinuum.COMPILED

protocol = IterativePhaseEstimationQuantinuum(
    backend=backend,
    optimisation_level=1,
    compilation_level=compilation_level,
).build(
    state=ansatz,
    evolution_operator_exponents=evo_ope_exp,
    eoe_totally_commuting=eoe_tot_com,
    encoding_method=CircuitEncoderQuantinuum.PLAIN,
    time_split=time_split,
)


# Show the quantum circuit.
if compilation_level < CompilationLevelQuantinuum.COMPILED:
    protocol.update_k_and_beta(k=2 * syndrome_interval, beta=0.25)
    circ = protocol.get_circuits()[0]
    render_circuit_jupyter(circ)


# ### Run the Stochastic QPE algorithm


def error_rate(k: int) -> float:
    q = 1 - (1 - p2) ** (10 * k)
    return q


print(error_rate(120))


phase_state = phayes.init(J=2000)

# Execute Bayesian QPE (modest setting to prioritize the speed).
verbose = 1
k_max = 256
n_updates = 120
conv = 3e-4


# Prepare the algorithm.
algorithm = AlgorithmBayesianQPE(
    phayes_state=phase_state,
    k_max=k_max,
    verbose=verbose,
    error_rate=error_rate,
).build(
    protocol=protocol,
)


# Run the algorithm.
ls_updated = []
for i in range(n_updates):
    handles_mapping = algorithm.run_async()
    algorithm.join(handles_mapping)
    mu, sigma = algorithm.final_value()
    ls_updated.append(algorithm.has_updated)
    print(f"i={i} mu={mu:10.6f} sigma={sigma:10.6f}")
    if sigma < conv:
        break


# Show the number of successful results.
print(f"Success: {ls_updated.count(True)} / {len(ls_updated)}")


# ### Analyze the results


# Show the result.
mu, sigma = algorithm.final_value()
energy_mu = -mu / time
energy_sigma = sigma / time
print(f"Energy(mu)    = {energy_mu:10.6f} hartree")
print(f"Energy(sigma) = {energy_sigma:10.6f} hartree")
print(f"FCI energy =    {fci_energy:10.6f} hartree")