'''
CALCULATING NUMBER OF MEASUREMENTS NEEDED AGAINST AN ERROR RATE
Example: H2 molecule in sto3g basis
'''

import numpy as np
from qiskit_nature.units import DistanceUnit
from qiskit_nature.second_q.drivers import PySCFDriver
from qiskit_nature.second_q.mappers import BravyiKitaevMapper




# specifying the geometry of the molecule in question
driver = PySCFDriver(
    atom="H 0 0 0; H 0 0 0.735",
    basis="sto3g",
    charge=0,
    spin=0,
    unit=DistanceUnit.ANGSTROM,
)


problem = driver.run()
hamiltonian = problem.hamiltonian

# electronic Hamiltonian of the system
second_q_op = hamiltonian.second_q_op()

# The Bravyi-Kitaev representation of the Fermionic Hamiltonian
mapper = BravyiKitaevMapper()
bkencoded_hamiltonian = mapper.map(second_q_op)

print(bkencoded_hamiltonian)



def shadow_bound(error, observables, failure_rate=0.01):
    """
    Method taken from: https://pennylane.ai/qml/demos/tutorial_classical_shadows
    
    Calculate the shadow bound for the Pauli measurement scheme.

    Implements Eq. (S13) from https://arxiv.org/pdf/2002.08953.pdf

    Args:
        error (float): The error on the estimator.
        observables (list) : List of matrices corresponding to the observables we intend to
            measure.
        failure_rate (float): Rate of failure for the bound to hold.

    Returns:
        An integer that gives the number of samples required to satisfy the shadow bound and
        the chunk size required attaining the specified failure rate.
    """
    M = len(observables)
    K = 2 * np.log(2 * M / failure_rate)
    shadow_norm = (
        lambda op: np.linalg.norm(
            op - np.trace(op) / 2 ** int(np.log2(op.shape[0])), ord=np.inf
        )
        ** 2
    )
    N = 34 * max(shadow_norm(o) for o in observables) / error ** 2
    return int(np.ceil(N * K)), int(K)


matrix_observables = [obs.to_matrix().real for obs in bkencoded_hamiltonian]
num_measurements, chunk_size = shadow_bound(error=2e-1, observables=matrix_observables, failure_rate=0.01)
print('Number of measurements (error=2e-1):', num_measurements)


