%load_ext autoreload
%autoreload 2

import numpy as np

from qiskit import QuantumCircuit, BasicAer
from qiskit.circuit.library import RealAmplitudes
from qiskit.opflow import PauliSumOp

observable = PauliSumOp.from_list(
    [
        ("IIXZ", -1.052373245772859),
        ("IZXZ", 0.39793742484318045),
        ("ZIZZ", -0.39793742484318045),
        ("ZZXZ", -0.01128010425623538),
        ("XXZZ", 0.18093119978423156),
    ]
)

U = RealAmplitudes(num_qubits=2, reps=2)

bitsrings = [[0, 1], [1, 0]]
schmidts = [1/np.sqrt(2**2), 1/np.sqrt(2**2)]

print("Schmidts:", schmidts)
print("Observable:")
print(observable)
print("U:")
U.decompose().draw()

from qiskit.evaluators import EntanglementForgingExpectationValue

expval = EntanglementForgingExpectationValue(U, observable,
                                             bitstrings=bitsrings,
                                             schmidts=schmidts,
                                             backend=BasicAer.get_backend("qasm_simulator"))

expectation = expval.evaluate([0, 1, 1, 2, 3, 5], shots=1000).value
print("EF exp value: {}".format(expectation))

