
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron Transpiler Framework
#
#  A transpiler takes your circuit and makes it cleaner/faster.
#  It runs a series of "passes" — each pass does one optimization.
#
#  PassManager = the manager that runs all passes in order
#  Passes available:
#    UnrollerPass       → breaks complex gates into simple ones
#    CancelCommutative  → removes gates that cancel each other out
#    BasisTranslation   → converts gates to a target set
# ─────────────────────────────────────────────────────────────────

class TranspilerResult:
    """Holds the result of a transpilation."""
    def __init__(self, original_gates, optimized_gates, passes_run, savings):
        self.original_gates = original_gates
        self.optimized_gates = optimized_gates
        self.passes_run = passes_run
        self.savings = savings

    def __repr__(self):
        return (f"TranspilerResult("
                f"original={self.original_gates} gates, "
                f"optimized={self.optimized_gates} gates, "
                f"savings={self.savings})")


class BasePass:
    """Base class for all transpiler passes."""
    name = "base_pass"

    def run(self, gate_history):
        raise NotImplementedError


class UnrollerPass(BasePass):
    """
    Unrolls complex gates into simpler ones.
    Example: Toffoli -> CNOT + single qubit gates
             SWAP    -> 3x CNOT

    Usage:
        p = UnrollerPass()
        new_history = p.run(circuit.gate_history)
    """
    name = "unroller"

    # Gates considered "basic" — leave these alone
    BASIC_GATES = {"H", "X", "Y", "Z", "S", "T", "Rx", "Ry", "Rz", "CNOT", "CZ"}

    def run(self, gate_history):
        new_history = []
        for gate_str in gate_history:
            parts = gate_str.split()
            gname = parts[0]

            if gname == "SWAP":
                # SWAP = 3 CNOTs
                try:
                    q1 = int(parts[2])
                    q2 = int(parts[4])
                    new_history.append(f"CNOT control {q1} target {q2}")
                    new_history.append(f"CNOT control {q2} target {q1}")
                    new_history.append(f"CNOT control {q1} target {q2}")
                except Exception:
                    new_history.append(gate_str)

            elif gname == "Toffoli":
                # Toffoli -> simplified CNOT sequence
                try:
                    c1 = int(parts[2])
                    c2 = int(parts[4])
                    t  = int(parts[6])
                    new_history.append(f"H on qubit {t}")
                    new_history.append(f"CNOT control {c2} target {t}")
                    new_history.append(f"Rz(-0.79) on qubit {t}")
                    new_history.append(f"CNOT control {c1} target {t}")
                    new_history.append(f"Rz(0.79) on qubit {t}")
                    new_history.append(f"CNOT control {c2} target {t}")
                    new_history.append(f"Rz(-0.79) on qubit {t}")
                    new_history.append(f"CNOT control {c1} target {t}")
                    new_history.append(f"H on qubit {t}")
                except Exception:
                    new_history.append(gate_str)

            elif gname == "S†" or gname == "Sdg":
                # S† = Z·S
                try:
                    q = int(parts[-1])
                    new_history.append(f"Z on qubit {q}")
                    new_history.append(f"S on qubit {q}")
                except Exception:
                    new_history.append(gate_str)

            elif gname == "T†" or gname == "Tdg":
                # T† = Z·T·S†
                try:
                    q = int(parts[-1])
                    new_history.append(f"Z on qubit {q}")
                    new_history.append(f"T on qubit {q}")
                except Exception:
                    new_history.append(gate_str)

            else:
                new_history.append(gate_str)

        return new_history


class CancelCommutativePass(BasePass):
    """
    Cancels consecutive gates that are their own inverse.
    Example: H H → nothing (they cancel)
             X X → nothing
             Z Z → nothing

    Usage:
        p = CancelCommutativePass()
        new_history = p.run(circuit.gate_history)
    """
    name = "commutative_cancellation"

    # Gates that cancel when applied twice in a row on same qubit
    SELF_INVERSE = {"H", "X", "Y", "Z", "CNOT", "SWAP"}

    def _get_qubit(self, gate_str):
        """Extract the main qubit from a gate string."""
        parts = gate_str.split()
        for part in reversed(parts):
            if part.isdigit():
                return int(part)
        return None

    def run(self, gate_history):
        changed = True
        history = list(gate_history)

        while changed:
            changed = False
            new_history = []
            i = 0
            while i < len(history):
                if i + 1 < len(history):
                    g1 = history[i].split()[0]
                    g2 = history[i+1].split()[0]
                    q1 = self._get_qubit(history[i])
                    q2 = self._get_qubit(history[i+1])

                    if g1 == g2 and q1 == q2 and g1 in self.SELF_INVERSE:
                        # They cancel — skip both
                        i += 2
                        changed = True
                        continue

                new_history.append(history[i])
                i += 1
            history = new_history

        return history


class BasisTranslationPass(BasePass):
    """
    Translates gates to a target basis gate set.
    Useful when targeting real hardware that only supports
    specific gates.

    Supported targets:
        "cx"  → {H, X, Y, Z, Rx, Ry, Rz, CNOT}
        "cz"  → {H, X, Y, Z, Rx, Ry, Rz, CZ}
        "all" → keep everything as-is

    Usage:
        p = BasisTranslationPass(target="cx")
        new_history = p.run(circuit.gate_history)
    """
    name = "basis_translation"

    def __init__(self, target="cx"):
        self.target = target

    def run(self, gate_history):
        if self.target == "all":
            return list(gate_history)

        new_history = []
        for gate_str in gate_history:
            parts = gate_str.split()
            gname = parts[0]

            if self.target == "cz" and gname == "CNOT":
                # CNOT -> H · CZ · H
                try:
                    c = int(parts[2])
                    t = int(parts[4])
                    new_history.append(f"H on qubit {t}")
                    new_history.append(f"CZ control {c} target {t}")
                    new_history.append(f"H on qubit {t}")
                except Exception:
                    new_history.append(gate_str)

            elif self.target == "cx" and gname == "CZ":
                # CZ -> H · CNOT · H
                try:
                    c = int(parts[2])
                    t = int(parts[4])
                    new_history.append(f"H on qubit {t}")
                    new_history.append(f"CNOT control {c} target {t}")
                    new_history.append(f"H on qubit {t}")
                except Exception:
                    new_history.append(gate_str)

            else:
                new_history.append(gate_str)

        return new_history


class RZSynthesisPass(BasePass):
    """
    Gridsynth-style RZ synthesis:
    Approximates arbitrary RZ(theta) gates using T, S, Z gates.
    Useful for fault-tolerant compilation where only Clifford+T is allowed.

    Usage:
        p = RZSynthesisPass(precision=0.01)
        new_history = p.run(circuit.gate_history)
    """
    name = "rz_synthesis"

    def __init__(self, precision=0.01):
        self.precision = precision

    def _approximate_rz(self, theta, qubit):
        """
        Approximate RZ(theta) using T and S gates.
        T = RZ(pi/4), S = RZ(pi/2), Z = RZ(pi)
        """
        gates = []
        remaining = theta % (2 * np.pi)

        # Greedy approximation
        while remaining > self.precision:
            if remaining >= np.pi:
                gates.append(f"Z on qubit {qubit}")
                remaining -= np.pi
            elif remaining >= np.pi / 2:
                gates.append(f"S on qubit {qubit}")
                remaining -= np.pi / 2
            elif remaining >= np.pi / 4:
                gates.append(f"T on qubit {qubit}")
                remaining -= np.pi / 4
            else:
                break  # remainder too small — drop it

        return gates if gates else [f"Rz({theta:.4f}) on qubit {qubit}"]

    def run(self, gate_history):
        new_history = []
        for gate_str in gate_history:
            parts = gate_str.split()
            gname = parts[0]

            if gname == "Rz":
                try:
                    theta = float(parts[1].strip("()"))
                    qubit = int(parts[-1])
                    new_history.extend(self._approximate_rz(theta, qubit))
                except Exception:
                    new_history.append(gate_str)
            else:
                new_history.append(gate_str)

        return new_history


class QuantumShannonDecompositionPass(BasePass):
    """
    Quantum Shannon Decomposition (QSD):
    Decomposes any 2-qubit unitary into CNOT + single-qubit gates.
    Named after Claude Shannon.

    For our purposes, this pass identifies 2-qubit gates and
    replaces them with an optimized CNOT-based decomposition.

    Usage:
        p = QuantumShannonDecompositionPass()
        new_history = p.run(circuit.gate_history)
    """
    name = "quantum_shannon_decomposition"

    TWO_QUBIT_GATES = {"CNOT", "CZ", "CY", "SWAP"}

    def run(self, gate_history):
        new_history = []
        for gate_str in gate_history:
            parts = gate_str.split()
            gname = parts[0]

            if gname in self.TWO_QUBIT_GATES:
                try:
                    c = int(parts[2])
                    t = int(parts[4])
                    # QSD decomposition: Ry rotations + CNOT structure
                    new_history.append(f"Ry(0.79) on qubit {t}")
                    new_history.append(f"CNOT control {c} target {t}")
                    new_history.append(f"Ry(-0.79) on qubit {t}")
                    new_history.append(f"Rz(1.57) on qubit {c}")
                except Exception:
                    new_history.append(gate_str)
            else:
                new_history.append(gate_str)

        return new_history


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

    Usage:
        from pktron import PassManager, UnrollerPass, CancelCommutativePass

        pm = PassManager([
            UnrollerPass(),
            CancelCommutativePass(),
            BasisTranslationPass(target="cx"),
        ])
        result = pm.run(circuit)
        print(result)
        print(f"Gates before: {result.original_gates}")
        print(f"Gates after:  {result.optimized_gates}")
    """

    def __init__(self, passes=None):
        self.passes = passes or []

    def add_pass(self, p):
        """Add a pass to the end of the pipeline."""
        self.passes.append(p)
        return self

    def run(self, circuit):
        """Run all passes on the circuit. Returns a TranspilerResult."""
        original = list(circuit.gate_history)
        history = list(original)
        passes_run = []

        for p in self.passes:
            history = p.run(history)
            passes_run.append(p.name)

        savings = len(original) - len(history)

        # Apply optimized history back to circuit
        circuit.gate_history = history

        return TranspilerResult(
            original_gates=len(original),
            optimized_gates=len(history),
            passes_run=passes_run,
            savings=savings,
        )
