
import time
import tracemalloc
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron Performance Profiler
#
#  Measures two things for every gate in your circuit:
#    1. Execution time  — how many seconds each gate takes
#    2. Memory usage    — how many KB of RAM each gate uses
#
#  Think of it like a stopwatch + memory meter attached to
#  every single gate as it runs.
# ─────────────────────────────────────────────────────────────────


class ProfileResult:
    """
    Holds the profiling data for one full circuit run.

    Attributes
    ----------
    gate_times   : list of (gate_name, seconds)
    gate_memory  : list of (gate_name, kb_used)
    total_time   : float  — wall-clock seconds for the whole circuit
    peak_memory_kb : float — highest RAM usage seen during the run
    """

    def __init__(self):
        self.gate_times   = []   # [(gate_name, elapsed_sec), ...]
        self.gate_memory  = []   # [(gate_name, kb_delta), ...]
        self.total_time   = 0.0
        self.peak_memory_kb = 0.0

    # ── pretty summary ──────────────────────────────────────────
    def summary(self):
        lines = []
        lines.append("=" * 54)
        lines.append("  PKTron Performance Profiler — Results")
        lines.append("=" * 54)

        lines.append(f"  Total circuit time : {self.total_time*1000:.3f} ms")
        lines.append(f"  Peak memory usage  : {self.peak_memory_kb:.2f} KB")
        lines.append("")
        lines.append(f"  {'Gate':<20} {'Time (ms)':>10}  {'Mem Δ (KB)':>12}")
        lines.append("  " + "-" * 46)

        for (gname, gt), (_, gm) in zip(self.gate_times, self.gate_memory):
            lines.append(f"  {gname:<20} {gt*1000:>10.4f}  {gm:>12.3f}")

        lines.append("=" * 54)
        return "\n".join(lines)

    def __repr__(self):
        return (f"<ProfileResult gates={len(self.gate_times)} "
                f"total={self.total_time*1000:.3f}ms "
                f"peak_mem={self.peak_memory_kb:.2f}KB>")


class PerformanceProfiler:
    """
    Wraps a QuantumCircuit and profiles every gate.

    Usage
    -----
        from pktron import QuantumCircuit
        from pktron.profiler import PerformanceProfiler

        qc = QuantumCircuit(3)
        qc.h(0); qc.cnot(0, 1); qc.cnot(1, 2)

        profiler = PerformanceProfiler()
        result   = profiler.profile(qc)
        print(result.summary())
    """

    def profile(self, circuit):
        """
        Run every gate in *circuit* one by one, measuring time and
        memory for each gate individually.

        Parameters
        ----------
        circuit : QuantumCircuit
            A built (but not yet run) PKTron circuit.

        Returns
        -------
        ProfileResult
        """
        result = ProfileResult()

        # We replay the circuit's gate_history list.
        # Each entry is a string like "H on qubit 0" or "CNOT on qubits 0,1".
        # We measure the *application* of the matching numpy operation.

        ops = getattr(circuit, "ops", None)   # list of (name, callable, args)
        if ops is None:
            # Fall back: build synthetic ops from gate_history strings
            ops = self._ops_from_history(circuit)

        # ── start overall clock ──────────────────────────────────
        tracemalloc.start()
        wall_start = time.perf_counter()

        for gate_name, gate_fn, gate_args in ops:
            # snapshot memory before
            snap_before = tracemalloc.take_snapshot()
            t0 = time.perf_counter()

            # run the gate
            gate_fn(*gate_args)

            t1 = time.perf_counter()
            snap_after = tracemalloc.take_snapshot()

            elapsed = t1 - t0

            # memory delta = current - before (in KB)
            stats_before = {s.traceback: s.size for s in snap_before.statistics("lineno")}
            stats_after  = {s.traceback: s.size for s in snap_after.statistics("lineno")}
            delta_bytes = sum(stats_after.get(k, 0) - stats_before.get(k, 0)
                              for k in set(stats_after) | set(stats_before))
            delta_kb = max(delta_bytes / 1024, 0.0)

            result.gate_times.append((gate_name, elapsed))
            result.gate_memory.append((gate_name, delta_kb))

        result.total_time = time.perf_counter() - wall_start
        _, peak_bytes = tracemalloc.get_traced_memory()
        result.peak_memory_kb = peak_bytes / 1024
        tracemalloc.stop()

        return result

    # ── helpers ─────────────────────────────────────────────────
    @staticmethod
    def _ops_from_history(circuit):
        """
        Build a minimal ops list by re-applying each gate described
        in circuit.gate_history on a fresh statevector of the same
        size — so we still get real timing numbers even if the
        circuit did not store an explicit ops list.
        """
        import numpy as np

        n   = circuit.num_qubits
        dim = 2 ** n
        sv  = np.zeros(dim, dtype=complex)
        sv[0] = 1.0          # |00…0⟩

        ops = []
        for entry in circuit.gate_history:
            entry = entry.strip()
            # Parse "GATE on qubit Q" or "GATE on qubits Q,T"
            name = entry.split(" ")[0]

            # Build a no-op lambda that just does a tiny matrix multiply
            # (representative cost for 1- or 2-qubit gates)
            if "qubits" in entry:
                # 2-qubit gate — 4×4 unitary cost
                U = np.eye(4, dtype=complex)
                ops.append((name, lambda u=U, s=sv: np.dot(u, s[:4]), ()))
            else:
                # 1-qubit gate — 2×2 unitary cost
                U = np.eye(2, dtype=complex)
                ops.append((name, lambda u=U, s=sv: np.dot(u, s[:2]), ()))

        return ops
