
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  Transpiler: a pipeline of "passes" that clean up a circuit.
#
#  Pass 1 — UnrollToffoli:    Toffoli → H + CNOT + T gates
#  Pass 2 — MergeRotations:   Rx(a) followed by Rx(b) → Rx(a+b)
#  Pass 3 — CancelInverse:    H followed by H → remove both
#  Pass 4 — BasisTranslation: rewrite any gate as U3 + CNOT
# ─────────────────────────────────────────────────────────────────

class TranspilerPass:
    """Base class — every pass must implement run()."""
    name = "BasePass"
    def run(self, ops):
        return ops   # ops = list of (gate_name, qubits, params)


class UnrollToffoliPass(TranspilerPass):
    """Decompose Toffoli (CCX) into basic H + T + CNOT gates."""
    name = "UnrollToffoli"

    def run(self, ops):
        out = []
        for gate, qubits, params in ops:
            if gate == "Toffoli" and len(qubits) == 3:
                c1, c2, t = qubits
                # Standard Toffoli decomposition
                out += [
                    ("H",    [t],      []),
                    ("CNOT", [c2, t],  []),
                    ("Tdg",  [t],      []),
                    ("CNOT", [c1, t],  []),
                    ("T",    [t],      []),
                    ("CNOT", [c2, t],  []),
                    ("Tdg",  [t],      []),
                    ("CNOT", [c1, t],  []),
                    ("T",    [c2],     []),
                    ("T",    [t],      []),
                    ("H",    [t],      []),
                    ("CNOT", [c1, c2], []),
                    ("T",    [c1],     []),
                    ("Tdg",  [c2],     []),
                    ("CNOT", [c1, c2], []),
                ]
            else:
                out.append((gate, qubits, params))
        return out


class MergeRotationsPass(TranspilerPass):
    """Merge consecutive same-axis rotations on same qubit. Rx(a)+Rx(b)→Rx(a+b)."""
    name = "MergeRotations"
    MERGE_GATES = {"Rx", "Ry", "Rz"}

    def run(self, ops):
        if not ops:
            return ops
        out = [ops[0]]
        for gate, qubits, params in ops[1:]:
            prev_gate, prev_qubits, prev_params = out[-1]
            # Same rotation axis, same qubit → merge
            if (gate in self.MERGE_GATES
                    and gate == prev_gate
                    and qubits == prev_qubits
                    and params and prev_params):
                merged_angle = float(prev_params[0]) + float(params[0])
                out[-1] = (gate, qubits, [merged_angle])
            else:
                out.append((gate, qubits, params))
        return out


class CancelInversePass(TranspilerPass):
    """Remove back-to-back self-inverse gates: H+H, X+X, CNOT+CNOT → nothing."""
    name = "CancelInverse"
    SELF_INVERSE = {"H", "X", "Y", "Z", "CNOT", "CZ", "SWAP"}

    def run(self, ops):
        out = []
        for gate, qubits, params in ops:
            if (out
                    and gate in self.SELF_INVERSE
                    and out[-1][0] == gate
                    and out[-1][1] == qubits):
                out.pop()   # cancel the pair
            else:
                out.append((gate, qubits, params))
        return out


class BasisTranslationPass(TranspilerPass):
    """Translate all gates to U3 + CNOT basis (universal gate set)."""
    name = "BasisTranslation"

    _TO_U3 = {
        "H":   (np.pi/2, 0,       np.pi),
        "X":   (np.pi,   0,       np.pi),
        "Y":   (np.pi,   np.pi/2, np.pi/2),
        "Z":   (0,       0,       np.pi),
        "S":   (0,       0,       np.pi/2),
        "Sdg": (0,       0,      -np.pi/2),
        "T":   (0,       0,       np.pi/4),
        "Tdg": (0,       0,      -np.pi/4),
    }

    def run(self, ops):
        out = []
        for gate, qubits, params in ops:
            if gate in self._TO_U3 and len(qubits) == 1:
                out.append(("U3", qubits, list(self._TO_U3[gate])))
            elif gate == "Rx" and params:
                out.append(("U3", qubits, [float(params[0]), -np.pi/2, np.pi/2]))
            elif gate == "Ry" and params:
                out.append(("U3", qubits, [float(params[0]), 0, 0]))
            elif gate == "Rz" and params:
                out.append(("U3", qubits, [0, 0, float(params[0])]))
            else:
                out.append((gate, qubits, params))
        return out


class PassManager:
    """
    Runs a list of transpiler passes in order.

    Usage:
        pm = PassManager([
            UnrollToffoliPass(),
            MergeRotationsPass(),
            CancelInversePass(),
        ])
        optimized_ops = pm.run(circuit)
        print(f"Gates before: {len(circuit.gate_history)}")
        print(f"Gates after:  {len(optimized_ops)}")
    """
    def __init__(self, passes=None):
        self.passes = passes or []

    def add_pass(self, p):
        self.passes.append(p)

    def run(self, circuit):
        """Parse gate_history and run all passes. Returns cleaned ops list."""
        ops = self._parse(circuit.gate_history)
        for p in self.passes:
            ops = p.run(ops)
        return ops

    def _parse(self, gate_history):
        """Convert gate_history strings to (name, qubits, params) tuples."""
        ops = []
        for entry in gate_history:
            if isinstance(entry, tuple):
                ops.append(entry)
                continue
            parts = entry.split()
            name = parts[0]
            nums = []
            for p in parts[1:]:
                try:
                    nums.append(float(p.strip("(),")))
                except ValueError:
                    pass
            if name in ("CNOT","CZ","CY","SWAP") and len(nums) >= 2:
                ops.append((name, [int(nums[0]), int(nums[1])], []))
            elif name == "Toffoli" and len(nums) >= 3:
                ops.append((name, [int(nums[0]), int(nums[1]), int(nums[2])], []))
            elif name in ("Rx","Ry","Rz") and len(nums) >= 2:
                ops.append((name, [int(nums[-1])], [nums[0]]))
            elif nums:
                ops.append((name, [int(nums[-1])], []))
            else:
                ops.append((name, [], []))
        return ops

    def summary(self, circuit):
        """Print before/after gate counts."""
        before = len(circuit.gate_history)
        after  = len(self.run(circuit))
        print(f"Transpiler summary: {before} gates → {after} gates")
        for p in self.passes:
            print(f"  Pass: {p.name}")
