
import numpy as np
from .quantum import QuantumCircuit

# ─────────────────────────────────────────────────────────────────
#  PKTron AerSimulator-like class
#  Supports three simulation methods:
#    "statevector"    → pure state vector (default, already in v0.1.6)
#    "density_matrix" → full density matrix (supports mixed/noisy states)
#    "unitary"        → builds the full unitary matrix of the circuit
# ─────────────────────────────────────────────────────────────────

class AerSimulator:
    """
    A configurable simulator backend — similar to Qiskit Aer.

    Usage:
        sim = AerSimulator(method="statevector")
        sim = AerSimulator(method="density_matrix")
        sim = AerSimulator(method="unitary")

        result = sim.run(circuit, shots=1024)
        print(result.get_counts())
        print(result.get_statevector())
        print(result.get_density_matrix())
        print(result.get_unitary())
    """

    SUPPORTED_METHODS = ["statevector", "density_matrix", "unitary"]

    def __init__(self, method="statevector"):
        method = method.lower()
        if method not in self.SUPPORTED_METHODS:
            raise ValueError(
                f"Unknown method: {method}. "
                f"Choose from: {self.SUPPORTED_METHODS}"
            )
        self.method = method
        self.name = f"aer_simulator_{method}"

    def properties(self):
        return {
            "name": self.name,
            "method": self.method,
            "max_qubits": 100,
            "simulator": True,
            "basis_gates": ["h","x","y","z","rx","ry","rz",
                            "cnot","cz","swap","toffoli"],
        }

    def run(self, circuit, shots=1024, noise_model=None):
        """
        Run the circuit using the chosen simulation method.
        Returns an AerResult object.
        """
        if self.method == "statevector":
            return self._run_statevector(circuit, shots, noise_model)
        elif self.method == "density_matrix":
            return self._run_density_matrix(circuit, shots, noise_model)
        elif self.method == "unitary":
            return self._run_unitary(circuit)

    # ── Method 1: Statevector ─────────────────────────────────────
    def _run_statevector(self, circuit, shots, noise_model):
        counts = {}
        for _ in range(shots):
            # Apply noise if provided
            if noise_model is not None:
                tmp = noise_model.apply(circuit)
                outcome, _ = tmp.measure()
            else:
                outcome, _ = circuit.measure()
            key = format(outcome, f"0{circuit.num_qubits}b")
            counts[key] = counts.get(key, 0) + 1

        state = circuit.get_state().copy()
        return AerResult(
            method="statevector",
            counts=counts,
            statevector=state,
            density_matrix=None,
            unitary=None,
            num_qubits=circuit.num_qubits,
            shots=shots,
        )

    # ── Method 2: Density Matrix ──────────────────────────────────
    def _run_density_matrix(self, circuit, shots, noise_model):
        """
        Simulates the circuit as a density matrix (rho = |psi><psi|
        for pure states, mixed for noisy states).
        """
        dm_sim = DensityMatrixSimulator(circuit.num_qubits)

        # Replay gate history onto density matrix simulator
        for gate_str in circuit.gate_history:
            parts = gate_str.split()
            gname = parts[0]

            if gname == "H":
                dm_sim.apply_single_gate(_H(), int(parts[-1]))
            elif gname == "X":
                dm_sim.apply_single_gate(_X(), int(parts[-1]))
            elif gname == "Y":
                dm_sim.apply_single_gate(_Y(), int(parts[-1]))
            elif gname == "Z":
                dm_sim.apply_single_gate(_Z(), int(parts[-1]))
            elif gname == "S":
                dm_sim.apply_single_gate(_S(), int(parts[-1]))
            elif gname == "T":
                dm_sim.apply_single_gate(_T(), int(parts[-1]))
            elif gname == "Rx":
                theta = float(parts[1].strip("()"))
                dm_sim.apply_single_gate(_Rx(theta), int(parts[-1]))
            elif gname == "Ry":
                theta = float(parts[1].strip("()"))
                dm_sim.apply_single_gate(_Ry(theta), int(parts[-1]))
            elif gname == "Rz":
                theta = float(parts[1].strip("()"))
                dm_sim.apply_single_gate(_Rz(theta), int(parts[-1]))
            elif gname == "CNOT":
                dm_sim.apply_two_gate(_CNOT(), int(parts[2]), int(parts[4]))
            elif gname == "CZ":
                dm_sim.apply_two_gate(_CZ(), int(parts[2]), int(parts[4]))
            elif gname == "SWAP":
                dm_sim.apply_two_gate(_SWAP(), int(parts[2]), int(parts[4]))

            # Apply noise after each gate if noise model provided
            if noise_model is not None:
                for qubit in range(circuit.num_qubits):
                    kraus_list = noise_model.get_kraus(gname, qubit)
                    if kraus_list:
                        dm_sim.apply_kraus(kraus_list, qubit)

        rho = dm_sim.density_matrix

        # Sample counts from density matrix probabilities
        probs = np.real(np.diag(rho))
        probs = np.abs(probs) / np.sum(np.abs(probs))
        counts = {}
        for _ in range(shots):
            outcome = np.random.choice(range(2**circuit.num_qubits), p=probs)
            key = format(outcome, f"0{circuit.num_qubits}b")
            counts[key] = counts.get(key, 0) + 1

        return AerResult(
            method="density_matrix",
            counts=counts,
            statevector=None,
            density_matrix=rho,
            unitary=None,
            num_qubits=circuit.num_qubits,
            shots=shots,
        )

    # ── Method 3: Unitary ─────────────────────────────────────────
    def _run_unitary(self, circuit):
        """
        Builds the full unitary matrix U of the circuit.
        U|0> gives the output state.
        """
        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]
            gate_u = None

            if gname == "H":
                gate_u = _full_gate(_H(), int(parts[-1]), n)
            elif gname == "X":
                gate_u = _full_gate(_X(), int(parts[-1]), n)
            elif gname == "Y":
                gate_u = _full_gate(_Y(), int(parts[-1]), n)
            elif gname == "Z":
                gate_u = _full_gate(_Z(), int(parts[-1]), n)
            elif gname == "Rx":
                theta = float(parts[1].strip("()"))
                gate_u = _full_gate(_Rx(theta), int(parts[-1]), n)
            elif gname == "Ry":
                theta = float(parts[1].strip("()"))
                gate_u = _full_gate(_Ry(theta), int(parts[-1]), n)
            elif gname == "Rz":
                theta = float(parts[1].strip("()"))
                gate_u = _full_gate(_Rz(theta), int(parts[-1]), n)

            if gate_u is not None:
                U = gate_u @ U

        return AerResult(
            method="unitary",
            counts={},
            statevector=None,
            density_matrix=None,
            unitary=U,
            num_qubits=n,
            shots=0,
        )


# ─────────────────────────────────────────────────────────────────
#  Density Matrix Simulator (used internally by AerSimulator)
# ─────────────────────────────────────────────────────────────────

class DensityMatrixSimulator:
    """
    Simulates a quantum circuit using density matrices.
    rho starts as |0><0| and evolves as rho -> U rho U†
    For noise: applies Kraus operators sum_k K_k rho K_k†
    """
    def __init__(self, num_qubits):
        self.num_qubits = num_qubits
        dim = 2**num_qubits
        self.density_matrix = np.zeros((dim, dim), dtype=complex)
        self.density_matrix[0, 0] = 1.0  # Start in |0><0|

    def apply_single_gate(self, gate, qubit):
        """Apply a single-qubit gate to the density matrix."""
        n = self.num_qubits
        full = 1.0
        for i in range(n):
            full = np.kron(full, gate if i == qubit else np.eye(2))
        rho = self.density_matrix
        self.density_matrix = full @ rho @ full.conj().T

    def apply_two_gate(self, gate, qubit1, qubit2):
        """Apply a two-qubit gate to the density matrix."""
        n = self.num_qubits
        dim = 2**n
        tensor = np.eye(dim, dtype=complex).reshape([2]*(2*n))
        gate_indices = [slice(None)]*(2*n)
        for q in sorted([qubit1, qubit2]):
            gate_indices[q] = slice(0,2)
            gate_indices[n+q] = slice(0,2)
        tensor[tuple(gate_indices)] = gate.reshape(2,2,2,2)
        full = tensor.reshape(dim, dim)
        rho = self.density_matrix
        self.density_matrix = full @ rho @ full.conj().T

    def apply_kraus(self, kraus_ops, qubit):
        """
        Apply a Kraus channel to one qubit.
        kraus_ops: list of 2x2 matrices [K0, K1, ...]
        rho -> sum_k (K_k ⊗ I) rho (K_k ⊗ I)†
        """
        n = self.num_qubits
        new_rho = np.zeros_like(self.density_matrix)
        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 @ self.density_matrix @ full.conj().T
        self.density_matrix = new_rho

    def get_probabilities(self):
        probs = np.real(np.diag(self.density_matrix))
        return np.abs(probs) / np.sum(np.abs(probs))

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

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


# ─────────────────────────────────────────────────────────────────
#  AerResult — holds results from any simulation method
# ─────────────────────────────────────────────────────────────────

class AerResult:
    def __init__(self, method, counts, statevector,
                 density_matrix, unitary, num_qubits, shots):
        self.method = method
        self._counts = counts
        self._statevector = statevector
        self._density_matrix = density_matrix
        self._unitary = unitary
        self.num_qubits = num_qubits
        self.shots = shots

    def get_counts(self):
        return self._counts

    def get_statevector(self):
        if self._statevector is None:
            raise ValueError("No statevector. Use method='statevector'")
        return self._statevector

    def get_density_matrix(self):
        if self._density_matrix is None:
            raise ValueError("No density matrix. Use method='density_matrix'")
        return self._density_matrix

    def get_unitary(self):
        if self._unitary is None:
            raise ValueError("No unitary. Use method='unitary'")
        return self._unitary

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


# ─────────────────────────────────────────────────────────────────
#  Helper gate matrices (used internally)
# ─────────────────────────────────────────────────────────────────

def _H():
    return np.array([[1,1],[1,-1]], dtype=complex) / np.sqrt(2)

def _X():
    return np.array([[0,1],[1,0]], dtype=complex)

def _Y():
    return np.array([[0,-1j],[1j,0]], dtype=complex)

def _Z():
    return np.array([[1,0],[0,-1]], dtype=complex)

def _S():
    return np.array([[1,0],[0,1j]], dtype=complex)

def _T():
    return np.array([[1,0],[0,np.exp(1j*np.pi/4)]], dtype=complex)

def _Rx(theta):
    c, s = np.cos(theta/2), np.sin(theta/2)
    return np.array([[c, -1j*s],[-1j*s, c]], dtype=complex)

def _Ry(theta):
    c, s = np.cos(theta/2), np.sin(theta/2)
    return np.array([[c,-s],[s,c]], dtype=complex)

def _Rz(theta):
    return np.array([[np.exp(-1j*theta/2),0],
                     [0,np.exp(1j*theta/2)]], dtype=complex)

def _CNOT():
    return np.array([[1,0,0,0],[0,1,0,0],
                     [0,0,0,1],[0,0,1,0]], dtype=complex)

def _CZ():
    return np.array([[1,0,0,0],[0,1,0,0],
                     [0,0,1,0],[0,0,0,-1]], dtype=complex)

def _SWAP():
    return np.array([[1,0,0,0],[0,0,1,0],
                     [0,1,0,0],[0,0,0,1]], dtype=complex)

def _full_gate(gate, qubit, num_qubits):
    """Expand a single-qubit gate to act on the full n-qubit space."""
    full = 1.0
    for i in range(num_qubits):
        full = np.kron(full, gate if i == qubit else np.eye(2))
    return full
