
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron SparseObservable + Advanced Estimator
#
#  SparseObservable lets you define multi-Pauli observables
#  efficiently. Instead of storing a huge matrix, we store
#  a list of (coefficient, Pauli_string) pairs.
#
#  Example: H = 0.5*ZZ + 0.3*XI - 0.1*IZ
#    obs = SparseObservable([("ZZ", 0.5), ("XI", 0.3), ("IZ", -0.1)])
# ─────────────────────────────────────────────────────────────────

_PAULI_MATS = {
    "I": np.eye(2, dtype=complex),
    "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),
}

def _pauli_string_matrix(label):
    """Build full matrix for a Pauli string like 'XZI'."""
    mat = None
    for ch in label.upper():
        g = _PAULI_MATS.get(ch)
        if g is None:
            raise ValueError(f"Unknown Pauli: {ch}")
        mat = g if mat is None else np.kron(mat, g)
    return mat


class SparseObservable:
    """
    An observable defined as a weighted sum of Pauli strings.
    Much more memory-efficient than storing the full matrix.

    Usage:
        # Simple Z on qubit 0
        obs = SparseObservable([("Z", 1.0)])

        # Heisenberg interaction: ZZ + XX + YY
        obs = SparseObservable([("ZZ", 1.0), ("XX", 1.0), ("YY", 1.0)])

        # With coefficients
        obs = SparseObservable([("ZI", 0.5), ("IZ", 0.5), ("XX", 0.25)])

        # Get the full matrix
        print(obs.to_matrix())

        # Expectation value
        print(obs.expectation_value(statevector))
    """

    def __init__(self, terms):
        """
        terms: list of (pauli_string, coefficient) tuples
               e.g. [("ZZ", 0.5), ("XI", -0.3)]
        """
        self.terms = []
        for item in terms:
            if isinstance(item, (list, tuple)) and len(item) == 2:
                label, coeff = item
            else:
                raise ValueError("Each term must be (pauli_string, coefficient)")
            self.terms.append((str(label).upper(), complex(coeff)))

        # Infer num_qubits from first term
        if self.terms:
            self.num_qubits = len(self.terms[0][0])
            for label, _ in self.terms:
                if len(label) != self.num_qubits:
                    raise ValueError(
                        f"All Pauli strings must have same length. "
                        f"Expected {self.num_qubits}, got {len(label)}"
                    )
        else:
            self.num_qubits = 0

    def to_matrix(self):
        """Build the full matrix representation."""
        dim = 2 ** self.num_qubits
        mat = np.zeros((dim, dim), dtype=complex)
        for label, coeff in self.terms:
            mat += coeff * _pauli_string_matrix(label)
        return mat

    def expectation_value(self, state):
        """
        Compute <state|Observable|state> efficiently term by term.
        state: numpy array, Statevector, or DensityMatrix
        """
        from .quantum_info import Statevector, DensityMatrix

        if isinstance(state, Statevector):
            vec = state.data
        elif isinstance(state, DensityMatrix):
            # Tr(rho * O)
            total = 0.0 + 0j
            for label, coeff in self.terms:
                P = _pauli_string_matrix(label)
                total += coeff * np.trace(state.data @ P)
            return float(np.real(total))
        elif isinstance(state, np.ndarray):
            vec = state
        else:
            raise ValueError("state must be Statevector, DensityMatrix, or numpy array")

        total = 0.0 + 0j
        for label, coeff in self.terms:
            P = _pauli_string_matrix(label)
            total += coeff * (vec.conj() @ P @ vec)
        return float(np.real(total))

    def add_term(self, pauli_string, coefficient):
        """Add a new Pauli term to the observable."""
        self.terms.append((pauli_string.upper(), complex(coefficient)))
        return self

    def simplify(self):
        """Merge duplicate Pauli strings."""
        merged = {}
        for label, coeff in self.terms:
            merged[label] = merged.get(label, 0) + coeff
        self.terms = [(k, v) for k, v in merged.items() if abs(v) > 1e-12]
        return self

    def num_terms(self):
        return len(self.terms)

    def __repr__(self):
        terms_str = " + ".join(
            f"({c:.3f})*{l}" for l, c in self.terms
        )
        return f"SparseObservable({terms_str})"


class AdvancedEstimatorResult:
    """Holds results from AdvancedEstimatorV2."""
    def __init__(self, expectation_values, variances, observable, shots):
        self.expectation_values = expectation_values
        self.variances = variances
        self.observable = observable
        self.shots = shots
        # Metadata
        self.metadata = {
            "shots": shots,
            "num_terms": observable.num_terms() if hasattr(observable, "num_terms") else 1,
        }

    def __repr__(self):
        return (f"AdvancedEstimatorResult("
                f"expval={self.expectation_values:.6f}, "
                f"variance={self.variances:.6f})")


class AdvancedEstimatorV2:
    """
    Improved EstimatorV2 that supports SparseObservable.
    Computes expectation values term by term — efficient for
    chemistry and optimization problems with many Pauli terms.

    Usage:
        obs = SparseObservable([("ZZ", 0.5), ("XX", 0.3)])
        estimator = AdvancedEstimatorV2(shots=1024)
        result = estimator.run(circuit, obs)
        print(result.expectation_values)

        # Batched over parameter sets
        results = estimator.run_batch(circuit, obs, [
            {theta: 0.5},
            {theta: 1.0},
        ])
    """

    def __init__(self, shots=1024):
        self.shots = shots

    def run(self, circuit, observable):
        """
        Run the circuit and compute expectation value of observable.
        observable: SparseObservable, or a single Pauli string like "ZZ"
        """
        if isinstance(observable, str):
            observable = SparseObservable([(observable, 1.0)])

        state = circuit.get_state()

        expval = observable.expectation_value(state)

        # Variance: <O^2> - <O>^2
        mat = observable.to_matrix()
        expval2 = float(np.real(state.conj() @ mat @ mat @ state))
        variance = expval2 - expval ** 2

        return AdvancedEstimatorResult(expval, variance, observable, self.shots)

    def run_batch(self, circuit, observable, parameter_list):
        """
        Run with multiple parameter sets.
        parameter_list: list of dicts like [{theta: 0.5}, {theta: 1.0}]
        Returns list of AdvancedEstimatorResult.
        """
        results = []
        for param_dict in parameter_list:
            bound = circuit.bind_parameters(param_dict)
            results.append(self.run(bound, observable))
        return results


class PauliProductMeasurement:
    """
    Instruction to measure a product of Pauli operators.
    Used in advanced synthesis for fault-tolerant circuits.

    Usage:
        ppm = PauliProductMeasurement("ZZII")
        result = ppm.measure(circuit)
        print(result)  # +1 or -1
    """

    def __init__(self, pauli_string):
        self.pauli_string = pauli_string.upper()
        self.num_qubits = len(pauli_string)
        self._matrix = _pauli_string_matrix(pauli_string)

    def measure(self, circuit):
        """
        Measure the Pauli product on the circuit state.
        Returns +1 or -1 (eigenvalue of the Pauli product).
        """
        state = circuit.get_state()
        expval = float(np.real(state.conj() @ self._matrix @ state))
        # Collapse to nearest eigenvalue
        return +1 if expval >= 0 else -1

    def expectation_value(self, circuit):
        """Returns the continuous expectation value (between -1 and +1)."""
        state = circuit.get_state()
        return float(np.real(state.conj() @ self._matrix @ state))

    def __repr__(self):
        return f"PauliProductMeasurement({self.pauli_string})"
