import numpy as np
from qiskit.circuit.library import EfficientSU2
from qiskit.quantum_info import SparsePauliOp

from circuit_knitting.cutting import (
    partition_problem,
    generate_cutting_experiments,
)

circuit = EfficientSU2(4, entanglement="linear", reps=2).decompose()
circuit.assign_parameters([0.8] * len(circuit.parameters), inplace=True)
observable = SparsePauliOp(["ZZZZ"])
circuit.draw("mpl", scale=0.8)

partitioned_problem = partition_problem(
    circuit=circuit, partition_labels="AABB", observables=observable.paulis
)
subcircuits = partitioned_problem.subcircuits
bases = partitioned_problem.bases
subobservables = partitioned_problem.subobservables
subcircuits["A"].draw("mpl", scale=0.6)

subcircuits["B"].draw("mpl", scale=0.6)

subexperiments, coefficients = generate_cutting_experiments(
    circuits=subcircuits,
    observables=subobservables,
    num_samples=np.inf,
)
coefficients

print(f"Mapping probabilities for a CNOT decomposition: {bases[0].probabilities}")

from circuit_knitting.cutting.qpd import QPDBasis
from qiskit.circuit.library.standard_gates import CXGate

qpd_basis_cx = QPDBasis.from_instruction(CXGate())


def _min_nonzero(seq, /):
    """Return the minimum value in a sequence, ignoring values near zero."""
    return min(x for x in seq if not np.isclose(x, 0))


num_cx_cuts = 2

print(
    f"Number of samples needed to retrieve exact weights: {1 / _min_nonzero(qpd_basis_cx.probabilities)**num_cx_cuts}"
)

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