
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron Quantum Information Module
#
#  Professional quantum info tools:
#    Statevector  → wraps a state vector with useful methods
#    DensityMatrix → wraps a density matrix with useful methods
#    Operator     → represents a quantum operator/unitary
#    Pauli        → single or multi-qubit Pauli operators
#    Clifford     → Clifford group elements
# ─────────────────────────────────────────────────────────────────

class Statevector:
    """
    Wraps a quantum state vector with useful methods.

    Usage:
        sv = Statevector([1, 0])          # |0>
        sv = Statevector([1, 1]) / sqrt(2) # |+>
        sv = Statevector.from_circuit(qc)
    """

    def __init__(self, data):
        arr = np.array(data, dtype=complex)
        norm = np.linalg.norm(arr)
        if norm > 0:
            arr = arr / norm
        self.data = arr
        self.num_qubits = int(np.log2(len(arr)))

    @classmethod
    def from_circuit(cls, circuit):
        """Build a Statevector from a QuantumCircuit."""
        circuit._allocate_state()
        return cls(circuit._state.copy())

    def __truediv__(self, scalar):
        return Statevector(self.data / scalar)

    def __mul__(self, scalar):
        return Statevector(self.data * scalar)

    def probabilities(self):
        """Returns measurement probabilities for each basis state."""
        return np.abs(self.data) ** 2

    def measure(self):
        """Sample one outcome from this state."""
        probs = self.probabilities()
        return np.random.choice(len(probs), p=probs)

    def fidelity(self, other):
        """
        Fidelity between this state and another.
        fidelity = |<psi|phi>|^2
        1.0 = identical, 0.0 = orthogonal
        """
        if isinstance(other, Statevector):
            return float(np.abs(np.dot(self.data.conj(), other.data)) ** 2)
        raise ValueError("Can only compute fidelity with another Statevector")

    def to_density_matrix(self):
        """Convert |psi> to density matrix rho = |psi><psi|"""
        rho = np.outer(self.data, self.data.conj())
        return DensityMatrix(rho)

    def inner_product(self, other):
        """<self|other> inner product."""
        return complex(np.dot(self.data.conj(), other.data))

    def is_valid(self):
        """Check if state is normalized."""
        return bool(np.isclose(np.sum(np.abs(self.data)**2), 1.0))

    def dims(self):
        return (2,) * self.num_qubits

    def __repr__(self):
        return f"Statevector({np.round(self.data, 4)}, num_qubits={self.num_qubits})"


class DensityMatrix:
    """
    Wraps a density matrix with useful methods.

    Usage:
        dm = DensityMatrix([[0.5, 0.5], [0.5, 0.5]])
        dm = DensityMatrix.from_circuit(qc)
        dm = DensityMatrix.from_statevector(sv)
    """

    def __init__(self, data):
        self.data = np.array(data, dtype=complex)
        self.num_qubits = int(np.log2(self.data.shape[0]))

    @classmethod
    def from_statevector(cls, sv):
        """Build DensityMatrix from a Statevector."""
        rho = np.outer(sv.data, sv.data.conj())
        return cls(rho)

    @classmethod
    def from_circuit(cls, circuit):
        """Build DensityMatrix from a QuantumCircuit."""
        sv = Statevector.from_circuit(circuit)
        return cls.from_statevector(sv)

    def purity(self):
        """Tr(rho^2). 1.0 = pure state, <1 = mixed state."""
        return float(np.real(np.trace(self.data @ self.data)))

    def trace(self):
        """Should always be 1.0 for a valid density matrix."""
        return float(np.real(np.trace(self.data)))

    def fidelity(self, other):
        """
        Fidelity between two density matrices.
        Uses sqrt(rho) method for mixed states.
        """
        if isinstance(other, Statevector):
            other = DensityMatrix.from_statevector(other)
        rho = self.data
        sigma = other.data
        sqrt_rho = _matrix_sqrt(rho)
        M = sqrt_rho @ sigma @ sqrt_rho
        sqrt_M = _matrix_sqrt(M)
        return float(np.real(np.trace(sqrt_M)) ** 2)

    def is_valid(self):
        """Check if density matrix is valid (trace=1, PSD, Hermitian)."""
        trace_ok = np.isclose(self.trace(), 1.0)
        herm_ok = np.allclose(self.data, self.data.conj().T)
        eigs = np.linalg.eigvalsh(self.data)
        psd_ok = bool(np.all(eigs >= -1e-10))
        return bool(trace_ok and herm_ok and psd_ok)

    def partial_trace(self, keep_qubit):
        """Partial trace — keep only one qubit's reduced state."""
        n = self.num_qubits
        rho = self.data.reshape([2] * (2 * n))
        keep = set([keep_qubit])
        traced = []
        for i in range(n):
            if i not in keep:
                ri = i - len([j for j in traced if j < i])
                ci = (n + i) - len([j for j in traced if j < n + i])
                rho = np.trace(rho, axis1=ri, axis2=ci)
                traced += [i, n + i]
        return DensityMatrix(rho)

    def von_neumann_entropy(self):
        """
        Von Neumann entropy S = -Tr(rho log rho).
        0.0 = pure state, log(d) = maximally mixed.
        """
        eigs = np.linalg.eigvalsh(self.data)
        eigs = eigs[eigs > 1e-12]
        return float(-np.sum(eigs * np.log2(eigs)))

    def expectation_value(self, operator):
        """Tr(rho * operator)"""
        if isinstance(operator, Operator):
            return complex(np.trace(self.data @ operator.data))
        return complex(np.trace(self.data @ operator))

    def __repr__(self):
        return (f"DensityMatrix(shape={self.data.shape}, "
                f"purity={self.purity():.4f}, "
                f"num_qubits={self.num_qubits})")


class Operator:
    """
    Represents a quantum operator (matrix).
    Can be unitary (gates) or Hermitian (observables).

    Usage:
        op = Operator([[0, 1], [1, 0]])  # X gate
        op = Operator.from_label("X")
    """

    _LABELS = {
        "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),
        "H": np.array([[1,1],[1,-1]], dtype=complex) / np.sqrt(2),
        "S": np.array([[1,0],[0,1j]], dtype=complex),
        "T": np.array([[1,0],[0,np.exp(1j*np.pi/4)]], dtype=complex),
    }

    def __init__(self, data):
        self.data = np.array(data, dtype=complex)
        self.num_qubits = int(np.log2(self.data.shape[0]))

    @classmethod
    def from_label(cls, label):
        """
        Build operator from a string label.
        Single qubit: "X", "Y", "Z", "H", "S", "T", "I"
        Multi-qubit:  "XX", "XZ", "IZ" etc. (tensor product)
        """
        mat = None
        for ch in label:
            g = cls._LABELS.get(ch.upper())
            if g is None:
                raise ValueError(f"Unknown label: {ch}")
            mat = g if mat is None else np.kron(mat, g)
        return cls(mat)

    def is_unitary(self):
        """Check if operator is unitary: U†U = I"""
        prod = self.data.conj().T @ self.data
        return bool(np.allclose(prod, np.eye(len(self.data))))

    def is_hermitian(self):
        """Check if operator is Hermitian: O = O†"""
        return bool(np.allclose(self.data, self.data.conj().T))

    def conjugate(self):
        """Return the conjugate transpose (dagger) of this operator."""
        return Operator(self.data.conj().T)

    def compose(self, other):
        """Compose two operators: self @ other"""
        return Operator(self.data @ other.data)

    def tensor(self, other):
        """Tensor product: self ⊗ other"""
        return Operator(np.kron(self.data, other.data))

    def eigenvalues(self):
        """Return eigenvalues."""
        return np.linalg.eigvals(self.data)

    def trace(self):
        return complex(np.trace(self.data))

    def __repr__(self):
        return (f"Operator(shape={self.data.shape}, "
                f"unitary={self.is_unitary()}, "
                f"hermitian={self.is_hermitian()})")


class Pauli:
    """
    Single or multi-qubit Pauli operator.
    Supports X, Y, Z, I and tensor products.

    Usage:
        p = Pauli("X")
        p = Pauli("XZ")   # X on qubit 0, Z on qubit 1
        p = Pauli("IZX")  # 3-qubit Pauli
        print(p.matrix)
        print(p.num_qubits)
    """

    _SINGLE = {
        "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 __init__(self, label):
        self.label = label.upper()
        self.num_qubits = len(self.label)
        self.matrix = self._build()

    def _build(self):
        mat = None
        for ch in self.label:
            g = self._SINGLE.get(ch)
            if g is None:
                raise ValueError(f"Invalid Pauli: {ch}. Use I, X, Y, Z")
            mat = g if mat is None else np.kron(mat, g)
        return mat

    def compose(self, other):
        """Multiply two Pauli operators."""
        return Operator(self.matrix @ other.matrix)

    def tensor(self, other):
        """Tensor product of two Paulis."""
        return Pauli.__new__(Pauli)

    def expectation_value(self, state):
        """<state|Pauli|state>"""
        if isinstance(state, Statevector):
            return float(np.real(state.data.conj() @ self.matrix @ state.data))
        elif isinstance(state, DensityMatrix):
            return float(np.real(np.trace(state.data @ self.matrix)))
        raise ValueError("state must be Statevector or DensityMatrix")

    def anticommutes_with(self, other):
        """Check if two Paulis anticommute: {A,B} = 0"""
        AB = self.matrix @ other.matrix
        BA = other.matrix @ self.matrix
        return bool(np.allclose(AB, -BA))

    def __repr__(self):
        return f"Pauli({self.label}, num_qubits={self.num_qubits})"


class Clifford:
    """
    Clifford group element — gates that map Paulis to Paulis.
    The Clifford group is generated by H, S, and CNOT.

    Usage:
        c = Clifford.from_label("H")
        c = Clifford.from_label("S")
        c = Clifford.from_circuit(qc)
    """

    CLIFFORD_GATES = {
        "H": np.array([[1,1],[1,-1]], dtype=complex) / np.sqrt(2),
        "S": np.array([[1,0],[0,1j]], 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),
    }

    NON_CLIFFORD = {"T", "Rx", "Ry", "Rz"}

    def __init__(self, unitary, label="custom"):
        self.unitary = np.array(unitary, dtype=complex)
        self.label = label

    @classmethod
    def from_label(cls, label):
        """Build a Clifford from a gate label."""
        gate = cls.CLIFFORD_GATES.get(label.upper())
        if gate is None:
            raise ValueError(f"{label} is not a Clifford gate")
        return cls(gate, label)

    @classmethod
    def from_circuit(cls, circuit):
        """
        Build a Clifford from a circuit — fails if non-Clifford gates present.
        """
        n = circuit.num_qubits
        U = np.eye(2**n, dtype=complex)
        for gate_str in circuit.gate_history:
            parts = gate_str.split()
            gname = parts[0]
            if gname in cls.NON_CLIFFORD:
                raise ValueError(
                    f"Gate {gname} is not a Clifford gate. "
                    f"Clifford circuits may only use H, S, X, Y, Z, CNOT."
                )
            gate = cls.CLIFFORD_GATES.get(gname)
            if gate is not None:
                full = 1.0
                q = int(parts[-1])
                for i in range(n):
                    full = np.kron(full, gate if i == q else np.eye(2))
                U = full @ U
        return cls(U, "circuit")

    def conjugate_pauli(self, pauli):
        """
        Conjugate a Pauli by this Clifford: C P C†
        Returns the resulting operator.
        """
        if isinstance(pauli, Pauli):
            result = self.unitary @ pauli.matrix @ self.unitary.conj().T
            return Operator(result)
        raise ValueError("Must pass a Pauli object")

    def is_clifford(self):
        """Always True for this class."""
        return True

    def __repr__(self):
        return f"Clifford(label={self.label}, shape={self.unitary.shape})"


# ─── Helper ───────────────────────────────────────────────────────

def _matrix_sqrt(M):
    """Compute matrix square root via eigendecomposition."""
    eigs, vecs = np.linalg.eigh(M)
    eigs = np.maximum(eigs, 0)
    return vecs @ np.diag(np.sqrt(eigs)) @ vecs.conj().T


def state_fidelity(state1, state2):
    """
    Convenience function: fidelity between any two states.
    Accepts Statevector or DensityMatrix.
    """
    if isinstance(state1, Statevector) and isinstance(state2, Statevector):
        return state1.fidelity(state2)
    if isinstance(state1, DensityMatrix):
        return state1.fidelity(state2)
    if isinstance(state1, Statevector):
        return DensityMatrix.from_statevector(state1).fidelity(state2)
    raise ValueError("Inputs must be Statevector or DensityMatrix")


def partial_trace(state, keep_qubits):
    """
    Partial trace over all qubits except keep_qubits.
    Returns a DensityMatrix.
    """
    if isinstance(state, Statevector):
        state = DensityMatrix.from_statevector(state)
    n = state.num_qubits
    rho = state.data.reshape([2] * (2 * n))
    keep = set(keep_qubits)
    traced = []
    for i in range(n):
        if i not in keep:
            ri = i - len([j for j in traced if j < i])
            ci = (n + i) - len([j for j in traced if j < n + i])
            rho = np.trace(rho, axis1=ri, axis2=ci)
            traced += [i, n + i]
    size = 2 ** len(keep_qubits)
    return DensityMatrix(rho.reshape(size, size))
