r"""Demonstration of the generation and use of sin-state phase estimation circuits, and plotting of results."""

from inquanto.operators import FermionOperatorString, FermionOperator
from inquanto.mappings import QubitMappingJordanWigner
from inquanto.ansatzes import MultiConfigurationState
from inquanto.states import QubitState, QubitStateString
from inquanto.algorithms import AlgorithmDeterministicQPE
from inquanto.protocols import SinStatePhaseEstimation, SinStatePhaseEstimator
from pytket.utils import probs_from_counts
from pytket.extensions.qiskit import AerBackend
import matplotlib.pyplot as plt
import scipy.linalg as la
import numpy as np


def plot_counts(probs):
    bins = [bi for bi, count in probs.items()]

    def bin2fraction(tup):
        bits = "".join([str(x) for x in tup])
        frac = int(bits, 2) / 2 ** len(bits)
        return frac

    fractions = [bin2fraction(elem) for elem in bins]
    outcome = [p for bi, p in probs.items()]
    fin_ene = -2.0 * (np.array(fractions)) / time
    eigs = ["{:.3f}".format(x) for x in fin_ene]

    plt.figure()
    plt.bar(eigs, outcome)
    plt.xlabel("eigenvalues")
    plt.ylabel("probability")
    plt.xticks(rotation=90)
    plt.title("probs")
    plt.show()


# define Hamiltonian: 2x2 matrix given by a,b,c,d
a, b, c, d = -0.16, 0.1, 0.1, 0.16
m = np.array([[a, b], [c, d]])
e, v = la.eig(m)
fos0 = FermionOperatorString(((0, 0), (0, 1)))
fos1 = FermionOperatorString(((0, 0), (1, 1)))
fos2 = FermionOperatorString(((1, 0), (0, 1)))
fos3 = FermionOperatorString(((1, 0), (1, 1)))
fop = FermionOperator.from_list([(a, fos0), (b, fos1), (c, fos2), (d, fos3)])
time = 1 / abs(np.real(e[1] - e[0]))

# define a 2-state ansatz
n_trot = 1
qop = QubitMappingJordanWigner().operator_map(fop)
qop_list = qop.trotterize(trotter_number=n_trot, constant=time)
c_1, c_2 = 0.0, 1.0
qubit_states = QubitState(
    {QubitStateString([1, 0]): c_1, QubitStateString([0, 1]): c_2}
)
ansatz = MultiConfigurationState(qubit_states)

# define sin-state QPE
algorithm = AlgorithmDeterministicQPE(
    ansatz,
    qop_list,
)

n_round = 3
n_shots = 10000

backend = AerBackend()
protocol = SinStatePhaseEstimation(
    backend=backend,
    n_rounds=n_round,
    n_shots=n_shots,
)
estimator = SinStatePhaseEstimator()
algorithm.build(
    protocol=protocol,
    phase_calculation_protocol=estimator,
)

# run & plot
algorithm.run()
energy = algorithm.final_energy(time=time, phase_estimator_protocol=estimator)
print("scipy eigenvalues:", e)
print(f"sin-state QPE energy estimate  = {energy:8.4f} hartree")

# Sin-state QPE uses an extra ancilla qubit for post-selecting the state of the rest of the ancilla
# For the purpose of plotting the results, we extract the experiment results and perform post-selection
post_counts = estimator.generate_report()["post_selected_counts"]
probs = probs_from_counts(post_counts)
plot_counts(probs)