r"""BQPE calculation using H1 device and the iceberg detection code."""
import numpy as np
import phayes
from pytket.circuit import (
    Circuit,
    Pauli,
    PauliExpBox,
    Qubit,
)
from pytket.pauli import QubitPauliString
from inquanto.operators import QubitOperator
from inquanto.ansatzes import CircuitAnsatz
from pytket.circuit.display import render_circuit_jupyter
from inquanto.protocols import (
    IterativePhaseEstimationQuantinuum,
    CircuitEncoderQuantinuum,
    CompilationLevelQuantinuum,
    IcebergOptions,
)
from inquanto.express import get_noisy_backend
from inquanto.extensions.phayes import AlgorithmBayesianQPE

# 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")

# Hamiltonian terms mapping.
terms_map = {
    QubitPauliString([Qubit(0)], [Pauli.Z]): QubitPauliString([Qubit(1)], [Pauli.Z]),
    QubitPauliString([Qubit(1)], [Pauli.Z]): QubitPauliString([Qubit(2)], [Pauli.Z]),
    QubitPauliString([Qubit(0), Qubit(1)], [Pauli.Y, Pauli.Y]): QubitPauliString(
        [Qubit(1), Qubit(2)], [Pauli.X, Pauli.X]
    ),
    QubitPauliString([Qubit(0), Qubit(1)], [Pauli.Z, Pauli.Z]): QubitPauliString(
        [Qubit(1), Qubit(2)], [Pauli.Z, Pauli.Z]
    ),
    QubitPauliString(): QubitPauliString(),
}
print("Terms mapping")
print(terms_map, end="\n\n")

# Pauli operator mapping primary for the iceberg code for the circuit optimization.
qubits = [Qubit(i) for i in range(4)]
paulis_map = {
    QubitPauliString(qubits, [Pauli.Z, Pauli.I, Pauli.X, Pauli.X]): QubitPauliString(
        qubits, [Pauli.Z, Pauli.X, Pauli.X, Pauli.X]
    )
}
print("Pauli term mappings")
print(paulis_map)

# 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(3)
state.add_pauliexpbox(
    PauliExpBox([Pauli.I, Pauli.Y, Pauli.X], -0.07113),
    state.qubits,
)
state.Sdg(1)
state.Sdg(2)
render_circuit_jupyter(state)
ansatz = CircuitAnsatz(state)


backend = get_noisy_backend(
    cx_err=2.5e-4,
    ro_err=0.0,
    n_qubits=8,
)


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

options = IcebergOptions(
    n_plus_states=2,
    syndrome_interval=syndrome_interval,
    sx_insertion=True,
)
protocol = IterativePhaseEstimationQuantinuum(
    backend=backend,
    optimisation_level=0,
    compilation_level=compilation_level,
).build(
    state=ansatz,
    evolution_operator_exponents=evo_ope_exp,
    eoe_totally_commuting=eoe_tot_com,
    encoding_method=CircuitEncoderQuantinuum.ICEBERG,
    encoding_options=options,
    terms_map=terms_map,
    paulis_map=paulis_map,
    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


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,
).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")