import gc
import re
import time
import numpy as np

try:
    import ipywidgets as widgets
    from IPython.display import display, clear_output
    _WIDGETS_OK = True
except ImportError:
    _WIDGETS_OK = False


# ---------------------------------------------------------
#  Statevector helpers
# ---------------------------------------------------------
def _apply1(sv, n, q, mat):
    dim = 2 ** n
    sv  = sv.reshape([2] * n)
    sv  = np.tensordot(mat, sv, axes=([1], [q]))
    sv  = np.moveaxis(sv, 0, q)
    return sv.reshape(dim)


def _apply_cx(sv, n, ctrl, tgt):
    dim = 2 ** n
    sv  = sv.reshape([2] * n)
    idx = [slice(None)] * n
    idx[ctrl] = 1
    tgt_local = tgt - (1 if tgt > ctrl else 0)
    sub = sv[tuple(idx)]
    X   = np.array([[0, 1], [1, 0]], dtype=complex)
    sub = np.tensordot(X, sub, axes=([1], [tgt_local]))
    sub = np.moveaxis(sub, 0, tgt_local)
    sv[tuple(idx)] = sub
    return sv.reshape(dim)


def _gate_matrix(gate_name, angle=0.0):
    g = gate_name.upper()
    a = float(angle) * np.pi
    gates = {
        "H" : np.array([[1, 1], [1, -1]]) / np.sqrt(2),
        "X" : np.array([[0, 1], [1,  0]], dtype=complex),
        "Y" : np.array([[0, -1j], [1j, 0]]),
        "Z" : np.array([[1, 0], [0, -1]], dtype=complex),
        "S" : np.array([[1, 0], [0, 1j]]),
        "T" : np.array([[1, 0], [0, np.exp(1j * np.pi / 4)]]),
        "RX": np.array([[np.cos(a/2), -1j*np.sin(a/2)],
                        [-1j*np.sin(a/2),  np.cos(a/2)]]),
        "RY": np.array([[np.cos(a/2), -np.sin(a/2)],
                        [ np.sin(a/2),  np.cos(a/2)]]),
        "RZ": np.array([[np.exp(-1j*a/2), 0],
                        [0, np.exp(1j*a/2)]]),
    }
    return gates.get(g, None)


def _parse_gate_history(gate_history):
    """
    Convert gate_history strings into structured step dicts.

    Handles all known pktron formats:
        "H on qubit 0"
        "CX on qubits 0,1"
        "CNOT control 0 target 1"
        "RX(0.5) on qubit 2"
        "SWAP on qubits 1,2"
    """
    steps = []
    for raw in gate_history:
        raw = raw.strip()
        matched = False

        # Format: "CNOT control 0 target 1"  (pktron native CNOT format)
        m = re.match(r"(\w+)\s+control\s+(\d+)\s+target\s+(\d+)", raw, re.IGNORECASE)
        if m:
            steps.append({
                "label" : raw,
                "gate"  : "CX",
                "qubit" : int(m.group(2)),
                "qubit2": int(m.group(3)),
                "angle" : 0.0,
            })
            continue

        # Format: "CX on qubits 0,1"  or  "SWAP on qubits 1, 2"
        m = re.match(r"(\w+)\s+on\s+qubits?\s+(\d+)[,\s]+(\d+)", raw, re.IGNORECASE)
        if m:
            steps.append({
                "label" : raw,
                "gate"  : m.group(1).upper(),
                "qubit" : int(m.group(2)),
                "qubit2": int(m.group(3)),
                "angle" : 0.0,
            })
            continue

        # Format: "RX(0.5) on qubit 2"
        m = re.match(r"(\w+)\(([0-9.]+)\)\s+on\s+qubit\s+(\d+)", raw, re.IGNORECASE)
        if m:
            steps.append({
                "label" : raw,
                "gate"  : m.group(1).upper(),
                "qubit" : int(m.group(3)),
                "qubit2": None,
                "angle" : float(m.group(2)),
            })
            continue

        # Format: "H on qubit 0"
        m = re.match(r"(\w+)\s+on\s+qubit\s+(\d+)", raw, re.IGNORECASE)
        if m:
            steps.append({
                "label" : raw,
                "gate"  : m.group(1).upper(),
                "qubit" : int(m.group(2)),
                "qubit2": None,
                "angle" : 0.0,
            })
            continue

        # Fallback — store as unknown, skip during simulation
        steps.append({
            "label" : raw,
            "gate"  : "??",
            "qubit" : 0,
            "qubit2": None,
            "angle" : 0.0,
        })

    return steps


# ---------------------------------------------------------
#  CircuitAnimator
# ---------------------------------------------------------
class CircuitAnimator:
    """
    Step-by-step circuit replay with Play / Pause / Rewind.

    Usage
    -----
        anim = CircuitAnimator(circuit)
        anim.launch()

        # headless:
        anim.step_forward()
        anim.step_backward()
        anim.rewind()
    """

    def __init__(self, circuit, auto_sim=True):
        if circuit.num_qubits > 12:
            raise ValueError(
                "Animation supports up to 12 qubits to stay within 12 GB RAM."
            )
        self.circuit    = circuit
        self.num_qubits = circuit.num_qubits
        self.steps      = _parse_gate_history(
            getattr(circuit, "gate_history", [])
        )
        self.total     = len(self.steps)
        self._step_idx = 0
        self._sv       = None
        self._playing  = False
        self._reset_sv()

    # -- statevector management ---------------------------
    def _reset_sv(self):
        dim        = 2 ** self.num_qubits
        self._sv   = np.zeros(dim, dtype=complex)
        self._sv[0] = 1.0

    def _apply_step(self, step):
        n  = self.num_qubits
        g  = step["gate"]
        q  = step["qubit"]
        q2 = step["qubit2"]
        a  = step["angle"]

        if g in ("CX", "CNOT"):
            self._sv = _apply_cx(self._sv, n, q, q2)
        elif g == "SWAP":
            self._sv = _apply_cx(self._sv, n, q, q2)
            self._sv = _apply_cx(self._sv, n, q2, q)
            self._sv = _apply_cx(self._sv, n, q, q2)
        else:
            mat = _gate_matrix(g, a)
            if mat is not None:
                self._sv = _apply1(self._sv, n, q, mat)

    def _rebuild_sv_to(self, target_idx):
        self._reset_sv()
        for i in range(target_idx):
            self._apply_step(self.steps[i])

    # -- navigation ---------------------------------------
    def step_forward(self):
        if self._step_idx >= self.total:
            return False
        self._apply_step(self.steps[self._step_idx])
        self._step_idx += 1
        if hasattr(self, "_ui_built"):
            self._refresh_display()
        return True

    def step_backward(self):
        if self._step_idx <= 0:
            return False
        self._step_idx -= 1
        self._rebuild_sv_to(self._step_idx)
        if hasattr(self, "_ui_built"):
            self._refresh_display()
        return True

    def rewind(self):
        self._step_idx = 0
        self._reset_sv()
        if hasattr(self, "_ui_built"):
            self._refresh_display()

    def jump_to(self, idx):
        idx = max(0, min(idx, self.total))
        self._step_idx = idx
        self._rebuild_sv_to(idx)
        if hasattr(self, "_ui_built"):
            self._refresh_display()

    # -- state info ---------------------------------------
    def current_probs(self):
        probs = np.abs(self._sv) ** 2
        n     = self.num_qubits
        fmt   = "{:0" + str(n) + "b}"
        return {fmt.format(i): round(float(p), 6)
                for i, p in enumerate(probs) if p > 1e-9}

    def state_vector_summary(self):
        n   = self.num_qubits
        fmt = "{:0" + str(n) + "b}"
        return [(fmt.format(i), amp)
                for i, amp in enumerate(self._sv) if abs(amp) > 1e-9]

    # =====================================================
    #  Interactive UI
    # =====================================================
    def launch(self):
        if not _WIDGETS_OK:
            raise ImportError("ipywidgets required: pip install ipywidgets")

        self._ui_built = True

        self._w_title = widgets.HTML(
            "<h3 style='color:#6C63FF;margin:0'>PKTron Circuit Animator</h3>"
            "<small style='color:#888'>"
            + str(self.num_qubits) + " qubits / "
            + str(self.total) + " gates</small>"
        )
        self._w_step     = widgets.HTML(self._step_html())
        self._w_progress = widgets.IntProgress(
            value=0, min=0, max=max(self.total, 1),
            bar_style="info",
            layout=widgets.Layout(width="100%")
        )
        self._w_circuit = widgets.Output()
        self._w_probs   = widgets.Output()

        def btn(desc, style, icon):
            return widgets.Button(
                description=desc, button_style=style,
                icon=icon, layout=widgets.Layout(width="110px")
            )

        self._b_rewind  = btn("Rewind",  "warning", "fast-backward")
        self._b_back    = btn("Back",    "",         "step-backward")
        self._b_play    = btn("Play",    "success",  "play")
        self._b_pause   = btn("Pause",   "danger",   "pause")
        self._b_forward = btn("Forward", "",         "step-forward")
        self._b_end     = btn("End",     "warning",  "fast-forward")

        self._w_speed = widgets.FloatSlider(
            value=0.8, min=0.1, max=3.0, step=0.1,
            description="Speed(s):",
            layout=widgets.Layout(width="280px")
        )
        self._w_jump = widgets.IntSlider(
            value=0, min=0, max=self.total,
            description="Jump to:",
            layout=widgets.Layout(width="280px")
        )

        self._b_rewind.on_click(lambda _: self.rewind())
        self._b_back.on_click(lambda _: self.step_backward())
        self._b_play.on_click(self._on_play)
        self._b_pause.on_click(self._on_pause)
        self._b_forward.on_click(lambda _: self.step_forward())
        self._b_end.on_click(lambda _: self.jump_to(self.total))
        self._w_jump.observe(self._on_jump_slider, names="value")

        controls = widgets.HBox([
            self._b_rewind, self._b_back,
            self._b_play,   self._b_pause,
            self._b_forward, self._b_end,
        ])
        sliders = widgets.HBox([self._w_speed, self._w_jump])

        self._w_ui = widgets.VBox([
            self._w_title,
            widgets.HTML("<hr style='margin:6px 0'>"),
            self._w_step,
            self._w_progress,
            controls, sliders,
            widgets.HTML("<b>Circuit steps:</b>"),
            self._w_circuit,
            widgets.HTML("<b>State probabilities:</b>"),
            self._w_probs,
        ])

        self._refresh_display()
        display(self._w_ui)

    def _on_play(self, _):
        self._playing = True
        while self._playing and self._step_idx < self.total:
            self.step_forward()
            time.sleep(self._w_speed.value)
        self._playing = False

    def _on_pause(self, _):
        self._playing = False

    def _on_jump_slider(self, change):
        self.jump_to(change["new"])

    def _refresh_display(self):
        self._w_step.value     = self._step_html()
        self._w_progress.value = self._step_idx
        if hasattr(self, "_w_jump"):
            self._w_jump.value = self._step_idx
        self._render_circuit()
        self._render_probs()

    def _step_html(self):
        if self._step_idx == 0:
            label = "Initial state |0...0>"
        elif self._step_idx <= self.total:
            label = self.steps[self._step_idx - 1]["label"]
        else:
            label = "Circuit complete"
        return (
            "<div style='padding:6px 10px;background:#1e1e2e;"
            "border-radius:6px;color:#cdd6f4;font-family:monospace'>"
            "Step " + str(self._step_idx) + "/" + str(self.total)
            + " -- " + label + "</div>"
        )

    def _render_circuit(self):
        with self._w_circuit:
            clear_output(wait=True)
            lines = []
            for i, s in enumerate(self.steps):
                arrow = ">" if i == self._step_idx - 1 else " "
                done  = "[x]" if i < self._step_idx else "[ ]"
                lines.append(
                    "  " + done + " " + arrow
                    + " " + str(i+1).rjust(2) + ". " + s["label"]
                )
            html = (
                "<pre style='background:#1e1e2e;color:#cdd6f4;"
                "padding:10px;border-radius:6px;font-size:12px;"
                "max-height:200px;overflow-y:auto'>"
                + ("\n".join(lines) if lines else "  (no gates)")
                + "</pre>"
            )
            display(widgets.HTML(html))

    def _render_probs(self):
        with self._w_probs:
            clear_output(wait=True)
            probs = self.current_probs()
            if not probs:
                display(widgets.HTML("<i>No amplitudes.</i>"))
                return
            max_p = max(probs.values()) or 1.0
            rows  = []
            for bs, p in sorted(probs.items()):
                bar_len = int(p / max_p * 28)
                bar     = "#" * bar_len + "." * (28 - bar_len)
                rows.append(
                    "  |" + bs + ">  " + bar + "  "
                    + str(round(p * 100, 2)).rjust(6) + "%"
                )
            html = (
                "<pre style='background:#1e1e2e;color:#a6e3a1;"
                "padding:10px;border-radius:6px;font-size:12px'>"
                + "\n".join(rows)
                + "</pre>"
            )
            display(widgets.HTML(html))

    def __repr__(self):
        return (
            "CircuitAnimator(qubits=" + str(self.num_qubits)
            + ", total_steps=" + str(self.total)
            + ", current_step=" + str(self._step_idx) + ")"
        )
