
import numpy as np

# ── What is this? ─────────────────────────────────────────────────
# SamplerV2  → runs a circuit many times and counts results
# EstimatorV2 → measures the expectation value of an observable
# Both support "batched" runs (multiple circuits or parameter sets)
# ──────────────────────────────────────────────────────────────────

class SamplerResult:
    """Holds the results from a SamplerV2 run."""
    def __init__(self, counts, probabilities, shots):
        self.counts = counts            # e.g. {"00": 512, "11": 512}
        self.probabilities = probabilities  # e.g. {"00": 0.5, "11": 0.5}
        self.shots = shots

    def __repr__(self):
        return (f"SamplerResult(shots={self.shots}, "
                f"counts={self.counts})")


class EstimatorResult:
    """Holds the results from an EstimatorV2 run."""
    def __init__(self, expectation_value, variance, shots):
        self.expectation_value = expectation_value
        self.variance = variance
        self.shots = shots

    def __repr__(self):
        return (f"EstimatorResult("
                f"expval={self.expectation_value:.4f}, "
                f"variance={self.variance:.4f})")


class SamplerV2:
    """
    Runs a circuit (or batch of circuits) and returns measurement counts.

    Basic usage:
        sampler = SamplerV2(shots=1024)
        result = sampler.run(my_circuit)
        print(result.counts)

    Batched usage (multiple parameter sets):
        results = sampler.run_batch(my_circuit, [
            {theta: 0.5},
            {theta: 1.0},
            {theta: 1.5},
        ])
    """
    def __init__(self, shots=1024):
        self.shots = shots

    def run(self, circuit):
        """Run one circuit and return a SamplerResult."""
        counts = {}
        probs_sum = None
        n = circuit.num_qubits

        for _ in range(self.shots):
            outcome, probs = circuit.measure()
            key = format(outcome, f"0{n}b")
            counts[key] = counts.get(key, 0) + 1
            if probs_sum is None:
                probs_sum = probs.copy()

        probabilities = {format(i, f"0{n}b"): counts.get(format(i, f"0{n}b"), 0) / self.shots
                         for i in range(2**n)}
        return SamplerResult(counts, probabilities, self.shots)

    def run_batch(self, circuit, parameter_list):
        """
        Run the same circuit with different parameter values.
        parameter_list: list of dicts like [{theta: 0.5}, {theta: 1.0}]
        Returns a list of SamplerResult objects.
        """
        results = []
        for param_dict in parameter_list:
            bound = circuit.bind_parameters(param_dict)
            results.append(self.run(bound))
        return results


class EstimatorV2:
    """
    Measures the expectation value of a Pauli observable on a circuit.

    Basic usage:
        estimator = EstimatorV2(shots=1024)
        result = estimator.run(my_circuit, observable="Z", qubit=0)
        print(result.expectation_value)

    Batched usage:
        results = estimator.run_batch(my_circuit, "Z", 0, [
            {theta: 0.5},
            {theta: 1.0},
        ])
    """
    def __init__(self, shots=1024):
        self.shots = shots

    def run(self, circuit, observable="Z", qubit=0):
        """Run one circuit and return an EstimatorResult."""
        import numpy as np
        paulis = {
            "X": np.array([[0,1],[1,0]], dtype=complex),
            "Y": np.array([[0,-1j],[1j,0]], dtype=complex),
            "Z": np.array([[1,0],[0,-1]], dtype=complex),
        }
        if observable not in paulis:
            raise ValueError("Observable must be X, Y, or Z")

        state = circuit.get_state()
        pauli = paulis[observable]
        full_obs = 1.0
        for i in range(circuit.num_qubits):
            full_obs = np.kron(full_obs, pauli if i == qubit else np.eye(2))

        expval = float(np.real(state.conj() @ full_obs @ state))
        expval2 = float(np.real(state.conj() @ (full_obs @ full_obs) @ state))
        variance = expval2 - expval**2

        return EstimatorResult(expval, variance, self.shots)

    def run_batch(self, circuit, observable, qubit, parameter_list):
        """
        Run the same circuit with different parameter values.
        Returns a list of EstimatorResult objects.
        """
        results = []
        for param_dict in parameter_list:
            bound = circuit.bind_parameters(param_dict)
            results.append(self.run(bound, observable, qubit))
        return results
