
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 = {}
        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()

    # ── Standard gates (ALL return self for chaining) ──
    def h(self, qubit):             h(self, qubit);             return self
    def x(self, qubit):             x(self, qubit);             return self
    def y(self, qubit):             y(self, qubit);             return self
    def z(self, qubit):             z(self, qubit);             return self
    def s(self, qubit):             s(self, qubit);             return self
    def sdg(self, qubit):           sdg(self, qubit);           return self
    def t(self, qubit):             t(self, qubit);             return self
    def tdg(self, qubit):           tdg(self, qubit);           return self
    def cnot(self, control, target): cnot(self, control, target); return self
    def cz(self, control, target):   cz(self, control, target);   return self
    def cy(self, control, target):   cy(self, control, target);   return self
    def swap(self, qubit1, qubit2):  swap(self, qubit1, qubit2);  return self
    def toffoli(self, c1, c2, tgt):  toffoli(self, c1, c2, tgt); return self

    # ── Rotation gates (theta FIRST, qubit SECOND) ──
    def rx(self, theta, qubit):
        if isinstance(theta, Parameter):
            theta = self._parameters.get(theta.name, theta)
        rx(self, qubit, float(theta))
        return self

    def ry(self, theta, qubit):
        if isinstance(theta, Parameter):
            theta = self._parameters.get(theta.name, theta)
        ry(self, qubit, float(theta))
        return self

    def rz(self, theta, qubit):
        if isinstance(theta, Parameter):
            theta = self._parameters.get(theta.name, theta)
        rz(self, qubit, float(theta))
        return self

    # ── Parameterized gates (return self) ──
    def rx_param(self, qubit, param):
        self._param_gates.append(("rx", qubit, param))
        self._parameters[param.name] = param
        self.gate_history.append(f"Rx({param.name}) on qubit {qubit} [unbound]")
        return self

    def ry_param(self, qubit, param):
        self._param_gates.append(("ry", qubit, param))
        self._parameters[param.name] = param
        self.gate_history.append(f"Ry({param.name}) on qubit {qubit} [unbound]")
        return self

    def rz_param(self, qubit, param):
        self._param_gates.append(("rz", qubit, param))
        self._parameters[param.name] = param
        self.gate_history.append(f"Rz({param.name}) on qubit {qubit} [unbound]")
        return self

    def bind_parameters(self, value_dict):
        new_circuit = QuantumCircuit(self.num_qubits, self.backend_type)
        new_circuit._state = self._state.copy() if self._state is not None else None
        new_circuit.gate_history = self.gate_history.copy()
        new_circuit._param_gates = self._param_gates.copy()
        new_circuit._parameters = self._parameters.copy()
        for param, value in value_dict.items():
            param_name = param.name if isinstance(param, Parameter) else str(param)
            new_circuit._parameters[param_name] = float(value)
        return new_circuit

    @property
    def parameters(self):
        return [k for k, v in self._parameters.items() if isinstance(v, Parameter)]

    def measure(self):
        self._allocate_state()
        sv = self._state
        probs = np.abs(sv)**2
        probs /= probs.sum()
        outcome_idx = np.random.choice(2**self.num_qubits, p=probs)
        return outcome_idx, probs

    def measure_all(self):
        outcome_idx, probs = self.measure()
        result = format(outcome_idx, f"0{self.num_qubits}b")
        self.gate_history.append(f"Measure all -> {result}")
        return result

    def save_statevector(self):
        self._allocate_state()
        self._saved_statevector = self._state.copy()
        return self._saved_statevector

    def copy(self):
        new_qc = QuantumCircuit(self.num_qubits, self.backend_type)
        new_qc._state = self._state.copy() if self._state is not None else None
        new_qc.gate_history = self.gate_history.copy()
        new_qc._param_gates = self._param_gates.copy()
        new_qc._parameters = self._parameters.copy()
        return new_qc

    def __repr__(self):
        return f"QuantumCircuit({self.num_qubits} qubits, {len(self.gate_history)} gates)"
