
import numpy as np
import matplotlib.pyplot as plt
from .gates import h, x, y, z, s, sdg, t, tdg, rx, ry, rz, cnot, cz, cy, swap, toffoli
from .parameters import Parameter, ParameterVector

class QuantumCircuit:
    def __init__(self, num_qubits, backend_type="statevector"):
        self.num_qubits = num_qubits
        self.backend_type = backend_type.lower()
        self._state = None
        self.gate_history = []
        self._param_gates = []
        self._parameters = {}
        # NEW v0.1.7: saved results storage
        self._saved_statevector = None
        self._saved_probabilities = None
        self._saved_expectation_values = {}

    def _allocate_state(self):
        if self._state is None:
            self._state = np.zeros(2**self.num_qubits, dtype=complex)
            self._state[0] = 1.0

    def get_state(self):
        self._allocate_state()
        return self._state

    def normalize(self):
        if self._state is None:
            return
        norm = np.sqrt(np.sum(np.abs(self._state)**2))
        if norm > 0:
            self._state /= norm

    def _apply_single_qubit_gate(self, gate, qubit):
        if qubit >= self.num_qubits:
            raise ValueError(f"Qubit {qubit} out of range")
        self._allocate_state()
        full_gate = 1.0
        for i in range(self.num_qubits):
            full_gate = np.kron(full_gate, gate if i == qubit else np.eye(2))
        self._state = np.dot(full_gate, self._state)
        self.normalize()

    def _apply_two_qubit_gate(self, gate, qubit1, qubit2):
        if qubit1 == qubit2 or max(qubit1, qubit2) >= self.num_qubits:
            raise ValueError(f"Invalid qubits {qubit1}, {qubit2}")
        self._allocate_state()
        qubits = sorted([qubit1, qubit2])
        tensor = np.eye(2**self.num_qubits, dtype=complex).reshape([2]*(2*self.num_qubits))
        gate_indices = [slice(None)]*(2*self.num_qubits)
        for q in qubits:
            gate_indices[q] = slice(0,2)
            gate_indices[self.num_qubits+q] = slice(0,2)
        tensor[tuple(gate_indices)] = gate.reshape(2,2,2,2)
        self._state = np.dot(tensor.reshape(2**self.num_qubits, 2**self.num_qubits), self._state)
        self.normalize()

    def _apply_three_qubit_gate(self, gate, qubit1, qubit2, qubit3):
        if len(set([qubit1,qubit2,qubit3])) != 3 or max(qubit1,qubit2,qubit3) >= self.num_qubits:
            raise ValueError(f"Invalid qubits {qubit1},{qubit2},{qubit3}")
        self._allocate_state()
        qubits = sorted([qubit1,qubit2,qubit3])
        tensor = np.eye(2**self.num_qubits, dtype=complex).reshape([2]*(2*self.num_qubits))
        gate_indices = [slice(None)]*(2*self.num_qubits)
        for q in qubits:
            gate_indices[q] = slice(0,2)
            gate_indices[self.num_qubits+q] = slice(0,2)
        tensor[tuple(gate_indices)] = gate.reshape(2,2,2,2,2,2)
        self._state = np.dot(tensor.reshape(2**self.num_qubits, 2**self.num_qubits), self._state)
        self.normalize()

    def _apply_gate(self, gate, qubits):
        if len(set(qubits)) != len(qubits) or max(qubits, default=-1) >= self.num_qubits:
            raise ValueError(f"Invalid qubits {qubits}")
        self._allocate_state()
        n = len(qubits)
        tensor = np.eye(2**self.num_qubits, dtype=complex).reshape([2]*(2*self.num_qubits))
        gate_indices = [slice(None)]*(2*self.num_qubits)
        for q in sorted(qubits):
            gate_indices[q] = slice(0,2)
            gate_indices[self.num_qubits+q] = slice(0,2)
        tensor[tuple(gate_indices)] = gate.reshape([2]*(2*n))
        self._state = np.dot(tensor.reshape(2**self.num_qubits, 2**self.num_qubits), self._state)
        self.normalize()

    # Standard gates
    def h(self, qubit): h(self, qubit)
    def x(self, qubit): x(self, qubit)
    def y(self, qubit): y(self, qubit)
    def z(self, qubit): z(self, qubit)
    def s(self, qubit): s(self, qubit)
    def sdg(self, qubit): sdg(self, qubit)
    def t(self, qubit): t(self, qubit)
    def tdg(self, qubit): tdg(self, qubit)
    def rx(self, qubit, theta): rx(self, qubit, theta)
    def ry(self, qubit, theta): ry(self, qubit, theta)
    def rz(self, qubit, theta): rz(self, qubit, theta)
    def cnot(self, control, target): cnot(self, control, target)
    def cz(self, control, target): cz(self, control, target)
    def cy(self, control, target): cy(self, control, target)
    def swap(self, qubit1, qubit2): swap(self, qubit1, qubit2)
    def toffoli(self, c1, c2, target): toffoli(self, c1, c2, target)

    # ── NEW v0.1.7: Parameterized gates ──────────────────────────
    def rx_param(self, qubit, param):
        """Rx gate using a Parameter object instead of a number."""
        self._param_gates.append(("rx", qubit, param))
        self._parameters[param.name] = param
        self.gate_history.append(f"Rx({param.name}) on qubit {qubit} [unbound]")

    def ry_param(self, qubit, param):
        """Ry gate using a Parameter object instead of a number."""
        self._param_gates.append(("ry", qubit, param))
        self._parameters[param.name] = param
        self.gate_history.append(f"Ry({param.name}) on qubit {qubit} [unbound]")

    def rz_param(self, qubit, param):
        """Rz gate using a Parameter object instead of a number."""
        self._param_gates.append(("rz", qubit, param))
        self._parameters[param.name] = param
        self.gate_history.append(f"Rz({param.name}) on qubit {qubit} [unbound]")

    def bind_parameters(self, value_dict):
        """
        Give real number values to the parameters and return a
        new circuit ready to run.

        Example:
            theta = Parameter("theta")
            qc.rx_param(0, theta)
            bound = qc.bind_parameters({theta: 1.57})
            outcome, probs = bound.measure()
        """
        name_to_value = {}
        for key, val in value_dict.items():
            name_to_value[key.name if isinstance(key, Parameter) else str(key)] = float(val)

        bound = QuantumCircuit(self.num_qubits, self.backend_type)
        for (gate_name, qubit, param) in self._param_gates:
            if param.name not in name_to_value:
                raise ValueError(f"No value given for parameter: {param.name}")
            val = name_to_value[param.name]
            if gate_name == "rx": bound.rx(qubit, val)
            elif gate_name == "ry": bound.ry(qubit, val)
            elif gate_name == "rz": bound.rz(qubit, val)
        return bound

    @property
    def parameters(self):
        """Shows which parameters still need values."""
        return set(self._parameters.values())

    # ── NEW v0.1.7: Save Methods ──────────────────────────────────
    def save_statevector(self):
        """Save a snapshot of the current statevector."""
        self._allocate_state()
        self._saved_statevector = self._state.copy()
        self.gate_history.append("SaveStatevector")
        return self._saved_statevector

    def save_probabilities(self):
        """Save a snapshot of the current measurement probabilities."""
        self._allocate_state()
        self._saved_probabilities = np.abs(self._state)**2
        self.gate_history.append("SaveProbabilities")
        return self._saved_probabilities

    def save_expectation_value(self, observable_name="Z", qubit=0):
        """
        Save the expectation value of a Pauli operator on one qubit.
        observable_name: "X", "Y", or "Z"
        """
        self._allocate_state()
        paulis = {
            "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),
        }
        if observable_name not in paulis:
            raise ValueError("Observable must be X, Y, or Z")
        pauli = paulis[observable_name]
        full_obs = 1.0
        for i in range(self.num_qubits):
            full_obs = np.kron(full_obs, pauli if i == qubit else np.eye(2))
        expval = float(np.real(self._state.conj() @ full_obs @ self._state))
        key = f"{observable_name}{qubit}"
        self._saved_expectation_values[key] = expval
        self.gate_history.append(f"SaveExpVal({observable_name} qubit {qubit})={expval:.4f}")
        return expval

    def get_saved_statevector(self):
        if self._saved_statevector is None:
            raise ValueError("No statevector saved. Call save_statevector() first.")
        return self._saved_statevector

    def get_saved_probabilities(self):
        if self._saved_probabilities is None:
            raise ValueError("No probabilities saved. Call save_probabilities() first.")
        return self._saved_probabilities

    def get_saved_expectation_values(self):
        if not self._saved_expectation_values:
            raise ValueError("No expectation values saved. Call save_expectation_value() first.")
        return self._saved_expectation_values

    def measure(self):
        if self._param_gates and self._state is None:
            raise ValueError("Circuit has unbound parameters. Call bind_parameters() first.")
        self._allocate_state()
        probs = np.abs(self._state)**2
        probs = probs / probs.sum()
        outcome = np.random.choice(range(2**self.num_qubits), p=probs)
        return outcome, probs

    def draw(self):
        fig, ax = plt.subplots(figsize=(max(10, len(self.gate_history)*2), self.num_qubits*0.5+2))
        ax.set_xlim(-0.5, len(self.gate_history)+1.5)
        ax.set_ylim(-0.5, self.num_qubits-0.5)
        ax.set_axis_off()
        for q in range(self.num_qubits):
            ax.plot([0, len(self.gate_history)+1], [q,q], "k-", linewidth=1)
            ax.text(-0.2, q, f"q{q}", fontsize=12, ha="right", va="center")
        for layer, gate in enumerate(self.gate_history):
            parts = gate.split()
            gname = parts[0]
            if gname in ["H","X","Y","Z","S","T"]:
                qubit = int(parts[-1])
                ax.add_patch(plt.Rectangle((layer+0.2, qubit-0.2), 0.6, 0.4, facecolor="lightblue", edgecolor="black"))
                ax.text(layer+0.5, qubit, gname, fontsize=8, ha="center", va="center")
            elif gname in ["Rx","Ry","Rz"]:
                qubit = int(parts[-1])
                ax.add_patch(plt.Rectangle((layer+0.2, qubit-0.2), 0.6, 0.4, facecolor="lightgreen", edgecolor="black"))
                ax.text(layer+0.5, qubit, gname[:2], fontsize=8, ha="center", va="center")
            elif gname == "CNOT":
                control, target = int(parts[2]), int(parts[4])
                ax.plot([layer+0.5, layer+0.5], [control, target], "k-", linewidth=1)
                ax.add_patch(plt.Circle((layer+0.5, control), 0.1, facecolor="black"))
                ax.add_patch(plt.Circle((layer+0.5, target), 0.2, facecolor="white", edgecolor="black"))
                ax.text(layer+0.5, target, "X", fontsize=10, ha="center", va="center")
        ax.set_title("PKTron v0.1.7 Circuit Diagram", fontsize=14, color="darkblue", pad=20)
        plt.tight_layout()
        plt.show()
