
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron Noise Modeling System
#
#  Think of noise as tiny errors that happen during real quantum
#  computing. We model them using "Kraus operators" — matrices that
#  describe how errors randomly affect qubits.
#
#  Supported noise types:
#    PauliError          → random X/Y/Z flip errors
#    DepolarizingError   → equal chance of any Pauli error
#    AmplitudeDamping    → qubit loses energy (|1> decays to |0>)
#    PhaseDamping        → qubit loses phase information
#    ThermalRelaxation   → combines T1 (energy) and T2 (phase) decay
# ─────────────────────────────────────────────────────────────────

class PauliError:
    """
    Random Pauli errors: X flip, Y flip, or Z flip.

    Example:
        # 1% X error, 1% Y error, 1% Z error
        err = PauliError([("X",0.01), ("Y",0.01), ("Z",0.01)])
    """
    def __init__(self, pauli_probs):
        self.pauli_probs = pauli_probs  # list of (gate_name, prob)
        self.kraus = self._build_kraus()

    def _build_kraus(self):
        gates = {
            "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),
        }
        total_error = sum(p for _, p in self.pauli_probs)
        kraus = []
        # Identity with probability of no error
        kraus.append(np.sqrt(1 - total_error) * np.eye(2, dtype=complex))
        for name, prob in self.pauli_probs:
            kraus.append(np.sqrt(prob) * gates[name])
        return kraus

    def get_kraus(self):
        return self.kraus


class DepolarizingError:
    """
    Depolarizing noise: equal probability of X, Y, or Z error.
    This is the most common noise model used in quantum computing.

    Example:
        err = DepolarizingError(prob=0.01)  # 1% depolarizing error
    """
    def __init__(self, prob):
        if prob < 0 or prob > 1:
            raise ValueError("Probability must be between 0 and 1")
        self.prob = prob
        self.kraus = self._build_kraus()

    def _build_kraus(self):
        p = self.prob
        return [
            np.sqrt(1 - p)     * np.eye(2, dtype=complex),
            np.sqrt(p / 3)     * np.array([[0,1],[1,0]], dtype=complex),
            np.sqrt(p / 3)     * np.array([[0,-1j],[1j,0]], dtype=complex),
            np.sqrt(p / 3)     * np.array([[1,0],[0,-1]], dtype=complex),
        ]

    def get_kraus(self):
        return self.kraus


class AmplitudeDampingError:
    """
    Amplitude damping: models energy loss (|1> slowly decays to |0>).
    Think of it like a qubit slowly "falling" to its ground state.

    Example:
        err = AmplitudeDampingError(gamma=0.05)  # 5% decay probability
    """
    def __init__(self, gamma):
        if gamma < 0 or gamma > 1:
            raise ValueError("Gamma must be between 0 and 1")
        self.gamma = gamma
        self.kraus = self._build_kraus()

    def _build_kraus(self):
        g = self.gamma
        K0 = np.array([[1, 0],[0, np.sqrt(1-g)]], dtype=complex)
        K1 = np.array([[0, np.sqrt(g)],[0, 0]], dtype=complex)
        return [K0, K1]

    def get_kraus(self):
        return self.kraus


class PhaseDampingError:
    """
    Phase damping: models loss of quantum phase (decoherence).
    The qubit stays in |0> or |1> but loses its superposition info.

    Example:
        err = PhaseDampingError(gamma=0.05)
    """
    def __init__(self, gamma):
        if gamma < 0 or gamma > 1:
            raise ValueError("Gamma must be between 0 and 1")
        self.gamma = gamma
        self.kraus = self._build_kraus()

    def _build_kraus(self):
        g = self.gamma
        K0 = np.array([[1,0],[0,np.sqrt(1-g)]], dtype=complex)
        K1 = np.array([[0,0],[0,np.sqrt(g)]], dtype=complex)
        return [K0, K1]

    def get_kraus(self):
        return self.kraus


class ThermalRelaxationError:
    """
    Thermal relaxation: combines T1 (energy decay) and T2 (phase decay).
    This is the most realistic noise model for real quantum hardware.

    T1 = energy relaxation time (how long qubit keeps its energy)
    T2 = dephasing time (how long qubit keeps its phase)
    t_gate = how long the gate takes (in same units as T1/T2)

    Example:
        err = ThermalRelaxationError(T1=100, T2=80, t_gate=1)
    """
    def __init__(self, T1, T2, t_gate):
        if T2 > 2 * T1:
            raise ValueError("T2 cannot exceed 2*T1")
        self.T1 = T1
        self.T2 = T2
        self.t_gate = t_gate
        self.kraus = self._build_kraus()

    def _build_kraus(self):
        T1, T2, t = self.T1, self.T2, self.t_gate
        p1 = np.exp(-t / T1)       # survival probability
        p2 = np.exp(-t / T2)       # coherence survival
        pz = (1 - p2) / 2

        K0 = np.array([[1,0],[0,np.sqrt(p1)*p2]], dtype=complex)
        K1 = np.array([[0,np.sqrt(1-p1)],[0,0]], dtype=complex)
        K2 = np.sqrt(pz) * np.array([[1,0],[0,-1]], dtype=complex)

        return [K0, K1, K2]

    def get_kraus(self):
        return self.kraus


# ─────────────────────────────────────────────────────────────────
#  NoiseModel — attach errors to specific gates or all qubits
# ─────────────────────────────────────────────────────────────────

class NoiseModel:
    """
    Attach noise to specific gates or all gates.

    Usage:
        noise = NoiseModel()

        # Add depolarizing noise to all H and CNOT gates
        noise.add_gate_error(DepolarizingError(0.01), ["H", "CNOT"])

        # Add amplitude damping to qubit 0 only
        noise.add_qubit_error(AmplitudeDampingError(0.05), [0])

        # Run with noise
        sim = AerSimulator(method="density_matrix")
        result = sim.run(circuit, noise_model=noise)
    """
    def __init__(self):
        # gate_name -> list of Kraus operator lists
        self._gate_errors = {}
        # qubit_index -> list of Kraus operator lists
        self._qubit_errors = {}

    def add_gate_error(self, error, gate_names):
        """Add an error channel that fires after specific gates."""
        if isinstance(gate_names, str):
            gate_names = [gate_names]
        for gate in gate_names:
            gate = gate.upper()
            if gate not in self._gate_errors:
                self._gate_errors[gate] = []
            self._gate_errors[gate].append(error.get_kraus())

    def add_qubit_error(self, error, qubits):
        """Add an error channel that fires on specific qubits."""
        if isinstance(qubits, int):
            qubits = [qubits]
        for q in qubits:
            if q not in self._qubit_errors:
                self._qubit_errors[q] = []
            self._qubit_errors[q].append(error.get_kraus())

    def get_kraus(self, gate_name, qubit):
        """
        Returns the Kraus operators to apply after a gate on a qubit.
        Used internally by AerSimulator density_matrix method.
        """
        kraus_list = []
        gate_upper = gate_name.upper()
        if gate_upper in self._gate_errors:
            for kset in self._gate_errors[gate_upper]:
                kraus_list.extend(kset)
        if qubit in self._qubit_errors:
            for kset in self._qubit_errors[qubit]:
                kraus_list.extend(kset)
        return kraus_list

    def apply(self, circuit):
        """
        Apply noise to a circuit by randomly inserting Pauli errors.
        Used by statevector method (approximate noise).
        Returns a new noisy circuit.
        """
        from .quantum import QuantumCircuit
        noisy = QuantumCircuit(circuit.num_qubits, "statevector")
        for gate_str in circuit.gate_history:
            parts = gate_str.split()
            gname = parts[0]
            if gname == "H":    noisy.h(int(parts[-1]))
            elif gname == "X":  noisy.x(int(parts[-1]))
            elif gname == "Y":  noisy.y(int(parts[-1]))
            elif gname == "Z":  noisy.z(int(parts[-1]))
            elif gname == "CNOT": noisy.cnot(int(parts[2]), int(parts[4]))
            elif gname == "CZ":   noisy.cz(int(parts[2]), int(parts[4]))
            elif gname == "SWAP": noisy.swap(int(parts[2]), int(parts[4]))
            # Apply random Pauli noise after gate
            g_up = gname.upper()
            if g_up in self._gate_errors:
                for kset in self._gate_errors[g_up]:
                    for qubit in range(circuit.num_qubits):
                        r = np.random.random()
                        cum = 0.0
                        for K in kset:
                            prob = float(np.real(np.trace(K.conj().T @ K))) / circuit.num_qubits
                            cum += prob
                            if r < cum:
                                # Apply the Pauli operator as a gate
                                if np.allclose(K/np.linalg.norm(K),
                                               np.array([[0,1],[1,0]])/np.sqrt(2)):
                                    noisy.x(qubit)
                                elif np.allclose(K/np.linalg.norm(K),
                                                 np.array([[1,0],[0,-1]])/np.sqrt(2)):
                                    noisy.z(qubit)
                                break
        return noisy

    def __repr__(self):
        return (f"NoiseModel("
                f"gate_errors={list(self._gate_errors.keys())}, "
                f"qubit_errors={list(self._qubit_errors.keys())})")
