
import numpy as np

class BitFlipError:
    """Random X (bit-flip) error with probability p."""
    def __init__(self, p):
        self.p = p
    def kraus(self):
        return [np.sqrt(1-self.p)*np.eye(2, dtype=complex),
                np.sqrt(self.p)*np.array([[0,1],[1,0]], dtype=complex)]

class PhaseFlipError:
    """Random Z (phase-flip) error with probability p."""
    def __init__(self, p):
        self.p = p
    def kraus(self):
        return [np.sqrt(1-self.p)*np.eye(2, dtype=complex),
                np.sqrt(self.p)*np.array([[1,0],[0,-1]], dtype=complex)]

class DepolarizingError:
    """Equal chance of X, Y, or Z error — most common noise model."""
    def __init__(self, p):
        self.p = p
    def kraus(self):
        return [
            np.sqrt(1-self.p)   * np.eye(2, dtype=complex),
            np.sqrt(self.p/3)   * np.array([[0,1],[1,0]], dtype=complex),
            np.sqrt(self.p/3)   * np.array([[0,-1j],[1j,0]], dtype=complex),
            np.sqrt(self.p/3)   * np.array([[1,0],[0,-1]], dtype=complex),
        ]

class AmplitudeDampingError:
    """Qubit loses energy — |1> slowly decays to |0>."""
    def __init__(self, gamma):
        self.gamma = gamma
    def kraus(self):
        g = self.gamma
        return [np.array([[1,0],[0,np.sqrt(1-g)]], dtype=complex),
                np.array([[0,np.sqrt(g)],[0,0]], dtype=complex)]

class PhaseDampingError:
    """Qubit loses phase — superposition decays without energy loss."""
    def __init__(self, gamma):
        self.gamma = gamma
    def kraus(self):
        g = self.gamma
        return [np.array([[1,0],[0,np.sqrt(1-g)]], dtype=complex),
                np.array([[0,0],[0,np.sqrt(g)]], dtype=complex)]

class NoiseModel:
    """
    Attach noise channels to specific gates.

    Usage:
        nm = NoiseModel()
        nm.add_error(DepolarizingError(0.01), ["H", "CNOT"])
        result = sim.run(circuit, shots=1024, noise_model=nm)
    """
    def __init__(self):
        self._errors = {}   # gate_name -> list of Kraus lists

    def add_error(self, error, gates):
        if isinstance(gates, str):
            gates = [gates]
        for g in gates:
            self._errors.setdefault(g.upper(), []).append(error.kraus())

    def get_kraus(self, gate_name):
        ops = []
        for kset in self._errors.get(gate_name.upper(), []):
            ops.extend(kset)
        return ops

    def apply_noise_to_state(self, state, gate_name, qubit, n):
        """Apply Kraus noise to a statevector after a gate."""
        from .fast_statevector import apply_1q
        kraus_ops = self.get_kraus(gate_name)
        if not kraus_ops:
            return state
        # Convert statevector to density matrix, apply Kraus, sample back
        rho = np.outer(state, state.conj())
        new_rho = np.zeros_like(rho)
        for K in kraus_ops:
            full = 1.0
            for i in range(n):
                full = np.kron(full, K if i == qubit else np.eye(2))
            new_rho += full @ rho @ full.conj().T
        # Sample a new pure state from the mixed state (approximate)
        probs = np.real(np.diag(new_rho))
        probs = np.abs(probs) / np.sum(np.abs(probs))
        idx = np.random.choice(len(probs), p=probs)
        new_state = new_rho[:, idx] / np.linalg.norm(new_rho[:, idx])
        return new_state

    def __repr__(self):
        return f"NoiseModel(gates={list(self._errors.keys())})"
