from qiskit import QuantumCircuit
from qiskit.quantum_info import SparsePauliOp

from circuit_knitting.cutting import (
    partition_problem,
    generate_cutting_experiments,
)

circuit = QuantumCircuit(2)
circuit.h(0)
circuit.cx(0, 1)
observable = SparsePauliOp(["ZZ"])
partitioned_problem = partition_problem(
    circuit=circuit, partition_labels="AB", observables=observable.paulis
)
subcircuits = partitioned_problem.subcircuits
subobservables = partitioned_problem.subobservables

subexperiments, coefficients = generate_cutting_experiments(
    circuits=subcircuits,
    observables=subobservables,
    num_samples=1000,
)

from circuit_knitting.utils.simulation import ExactSampler

exact_sampler = ExactSampler()

results = {
    label: exact_sampler.run(subexperiment).result()
    for label, subexperiment in subexperiments.items()
}