
import numpy as np
from .quantum import QuantumCircuit

# ─────────────────────────────────────────────────────────────────
#  PKTron Mid-Circuit Measurement & Dynamic Circuits
#
#  Classical registers store measurement results mid-circuit.
#  Conditional gates fire only if a classical bit equals a value.
#  If/While/Switch control flow lets circuits adapt as they run.
# ─────────────────────────────────────────────────────────────────

class ClassicalRegister:
    """
    Stores classical bits (0 or 1) from mid-circuit measurements.

    Example:
        cr = ClassicalRegister(2, name="c")
        print(cr[0])  # 0
    """
    def __init__(self, size, name="cr"):
        self.size = size
        self.name = name
        self.bits = [0] * size

    def __getitem__(self, index):
        return self.bits[index]

    def __setitem__(self, index, value):
        self.bits[index] = int(value)

    def __repr__(self):
        return f"ClassicalRegister({self.name}, bits={self.bits})"


class DynamicCircuit(QuantumCircuit):
    """
    A QuantumCircuit that supports:
      - Mid-circuit measurement (measure one qubit, store in classical bit)
      - Conditional gates (only run if classical bit == value)
      - If / While / Switch control flow

    Example:
        qc = DynamicCircuit(2, 2)
        qc.h(0)
        qc.measure_mid(0, 0)          # measure qubit 0 into classical bit 0
        qc.if_then(0, 1, "x", 1)     # if classical bit 0 == 1, apply X to qubit 1
        outcome, probs = qc.measure()
    """

    def __init__(self, num_qubits, num_classical=0, backend_type="statevector"):
        super().__init__(num_qubits, backend_type)
        self.classical = ClassicalRegister(max(num_classical, 1))
        self.control_log = []   # log of all control flow decisions

    # ── Mid-Circuit Measurement ───────────────────────────────────
    def measure_mid(self, qubit, classical_bit):
        """
        Measure one qubit right now and store result in a classical bit.
        The qubit collapses to |0> or |1> based on probabilities.
        """
        self._allocate_state()
        n = self.num_qubits
        probs = np.abs(self._state) ** 2

        # Probability that qubit is |1>
        prob_one = sum(
            probs[i] for i in range(2**n) if (i >> qubit) & 1
        )
        result = int(np.random.random() < prob_one)
        self.classical[classical_bit] = result

        # Collapse the state
        new_state = self._state.copy()
        for i in range(2**n):
            bit_val = (i >> qubit) & 1
            if bit_val != result:
                new_state[i] = 0.0
        norm = np.sqrt(np.sum(np.abs(new_state)**2))
        if norm > 0:
            new_state /= norm
        self._state = new_state

        self.gate_history.append(
            f"MeasureMid qubit {qubit} -> classical[{classical_bit}]={result}"
        )
        self.control_log.append(
            f"MeasureMid: qubit {qubit} collapsed to |{result}>"
        )
        return result

    # ── Conditional Gate ─────────────────────────────────────────
    def if_then(self, classical_bit, value, gate_name, qubit):
        """
        Apply a gate ONLY IF classical_bit equals value.

        Example:
            qc.if_then(0, 1, "x", 1)
            # if classical bit 0 is 1, apply X gate to qubit 1
        """
        actual = self.classical[classical_bit]
        fired = actual == value
        if fired:
            gate_fn = getattr(self, gate_name.lower(), None)
            if gate_fn:
                gate_fn(qubit)
        self.gate_history.append(
            f"IfThen classical[{classical_bit}]=={value} -> "
            f"{'FIRED' if fired else 'SKIPPED'} {gate_name.upper()} on qubit {qubit}"
        )
        self.control_log.append(
            f"IfThen: bit[{classical_bit}]={actual}, condition={value}, fired={fired}"
        )
        return fired

    # ── If/Else ───────────────────────────────────────────────────
    def if_else(self, classical_bit, value, gate_if, qubit_if, gate_else, qubit_else):
        """
        If classical_bit == value: apply gate_if to qubit_if
        Else: apply gate_else to qubit_else

        Example:
            qc.if_else(0, 1, "x", 1, "z", 0)
        """
        actual = self.classical[classical_bit]
        if actual == value:
            getattr(self, gate_if.lower())(qubit_if)
            self.control_log.append(f"IfElse: took IF branch -> {gate_if} on qubit {qubit_if}")
        else:
            getattr(self, gate_else.lower())(qubit_else)
            self.control_log.append(f"IfElse: took ELSE branch -> {gate_else} on qubit {qubit_else}")
        self.gate_history.append(
            f"IfElse classical[{classical_bit}]=={value} actual={actual}"
        )

    # ── While Loop ────────────────────────────────────────────────
    def while_loop(self, classical_bit, value, gate_name, qubit, max_iter=10):
        """
        Keep applying a gate WHILE classical_bit == value.
        Remeasures the qubit each iteration.
        max_iter prevents infinite loops.

        Example:
            qc.while_loop(0, 1, "h", 0)
        """
        iterations = 0
        while self.classical[classical_bit] == value and iterations < max_iter:
            getattr(self, gate_name.lower())(qubit)
            self.measure_mid(qubit, classical_bit)
            iterations += 1
        self.gate_history.append(
            f"WhileLoop classical[{classical_bit}]=={value} "
            f"-> {gate_name.upper()} on qubit {qubit} x{iterations}"
        )
        self.control_log.append(
            f"WhileLoop: ran {iterations} iterations"
        )
        return iterations

    # ── Switch/Case ───────────────────────────────────────────────
    def switch(self, classical_bit, cases):
        """
        Switch on the value of a classical bit and apply matching gate.
        cases: dict like {0: ("x", 1), 1: ("h", 0)}
        meaning: if bit==0 apply X to qubit 1, if bit==1 apply H to qubit 0

        Example:
            qc.switch(0, {0: ("x", 1), 1: ("h", 0)})
        """
        actual = self.classical[classical_bit]
        if actual in cases:
            gate_name, qubit = cases[actual]
            getattr(self, gate_name.lower())(qubit)
            self.control_log.append(
                f"Switch: bit={actual} matched -> {gate_name} on qubit {qubit}"
            )
        else:
            self.control_log.append(f"Switch: bit={actual} no match")
        self.gate_history.append(
            f"Switch classical[{classical_bit}]={actual}"
        )

    # ── Feed-Forward ──────────────────────────────────────────────
    def feed_forward(self, qubit, classical_bit, correct_gate="x"):
        """
        Measure a qubit mid-circuit, then automatically apply a
        correction gate if the result was |1>. Common in teleportation.

        Example:
            qc.feed_forward(0, 0)
            # measures qubit 0, if result=1 applies X correction
        """
        result = self.measure_mid(qubit, classical_bit)
        if result == 1:
            getattr(self, correct_gate.lower())(qubit)
            self.control_log.append(
                f"FeedForward: measured |1>, applied {correct_gate} correction"
            )
        else:
            self.control_log.append(
                f"FeedForward: measured |0>, no correction needed"
            )
        self.gate_history.append(
            f"FeedForward qubit {qubit} classical[{classical_bit}]={result}"
        )
        return result

    def get_classical_state(self):
        """Returns the current classical register bits."""
        return self.classical.bits.copy()

    def get_control_log(self):
        """Returns a full log of all control flow decisions."""
        return self.control_log.copy()
