
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron General Quantum Walks
#
#  What is a Quantum Walk?
#    Like a classical random walk (a drunkard stumbling left/right)
#    but quantum — the walker is in SUPERPOSITION of all positions
#    at once. This gives quadratic speedup over classical walks.
#
#  Two types supported:
#    1. DISCRETE  — step by step, uses a coin + shift operator
#    2. CONTINUOUS — evolves smoothly with e^(-iHt)
#
#  Works on ANY graph:
#    - Line (1D)
#    - Cycle (ring)
#    - Grid (2D)
#    - Hypercube
#    - Complete
#    - Custom (you provide edges)
#
#  Key quantities:
#    - Probability distribution over nodes at each step
#    - Mixing time  (steps to reach near-uniform distribution)
#    - Hitting time (steps to first reach a target node)
#    - Quantum vs classical comparison
#
#  Usage:
#      qw = QuantumWalk("cycle", n_nodes=8)
#      result = qw.run(steps=20)
#      result.summary()
#      plot_walk(result)
# ─────────────────────────────────────────────────────────────────


# ── Graph builders ────────────────────────────────────────────────

def build_graph(graph_type: str, n_nodes: int,
                edges: list = None) -> np.ndarray:
    """
    Build an adjacency matrix for a named graph type.

    Args:
        graph_type : "line","cycle","grid","hypercube",
                     "complete","custom"
        n_nodes    : number of nodes
        edges      : list of (i,j) tuples for custom graphs

    Returns:
        adjacency matrix, shape (n_nodes, n_nodes)

    Example:
        A = build_graph("cycle", 8)
        A = build_graph("custom", 5,
                edges=[(0,1),(1,2),(2,3),(3,4),(4,0),(0,2)])
    """
    if graph_type == "line":
        A = np.zeros((n_nodes, n_nodes))
        for i in range(n_nodes - 1):
            A[i, i+1] = A[i+1, i] = 1
        return A

    elif graph_type == "cycle":
        A = np.zeros((n_nodes, n_nodes))
        for i in range(n_nodes):
            A[i, (i+1) % n_nodes] = 1
            A[(i+1) % n_nodes, i] = 1
        return A

    elif graph_type == "complete":
        return np.ones((n_nodes, n_nodes)) - np.eye(n_nodes)

    elif graph_type == "grid":
        side = int(np.round(np.sqrt(n_nodes)))
        if side * side != n_nodes:
            raise ValueError(
                f"grid needs n_nodes = perfect square. "
                f"Got {n_nodes}. Try 4,9,16,25 ...")
        A = np.zeros((n_nodes, n_nodes))
        for r in range(side):
            for c in range(side):
                node = r * side + c
                if c + 1 < side:
                    A[node, node+1] = A[node+1, node] = 1
                if r + 1 < side:
                    A[node, node+side] = A[node+side, node] = 1
        return A

    elif graph_type == "hypercube":
        n_bits = int(np.round(np.log2(n_nodes)))
        if 2**n_bits != n_nodes:
            raise ValueError(
                f"hypercube needs n_nodes = 2^k. "
                f"Got {n_nodes}. Try 2,4,8,16 ...")
        A = np.zeros((n_nodes, n_nodes))
        for i in range(n_nodes):
            for b in range(n_bits):
                j = i ^ (1 << b)
                A[i, j] = 1
        return A

    elif graph_type == "custom":
        if edges is None:
            raise ValueError("custom graph needs edges list")
        A = np.zeros((n_nodes, n_nodes))
        for (i, j) in edges:
            A[i, j] = A[j, i] = 1
        return A

    else:
        raise ValueError(
            f"Unknown graph: {graph_type!r}. "
            f"Choose: line,cycle,grid,hypercube,complete,custom")


# ── Coin operators ────────────────────────────────────────────────

def hadamard_coin(d: int) -> np.ndarray:
    """Generalised Hadamard (DFT) coin for degree-d nodes."""
    H = np.zeros((d, d), dtype=complex)
    for j in range(d):
        for k in range(d):
            H[j, k] = np.exp(2j * np.pi * j * k / d) / np.sqrt(d)
    return H

def grover_coin(d: int) -> np.ndarray:
    """Grover diffusion coin — faster mixing than Hadamard."""
    return (2.0 / d) * np.ones((d, d)) - np.eye(d)

def dft_coin(d: int) -> np.ndarray:
    """DFT coin — alias for hadamard_coin."""
    return hadamard_coin(d)


# ── Result class ──────────────────────────────────────────────────

class QuantumWalkResult:
    """
    Stores all results from a quantum walk simulation.

    Attributes:
        probs_history  : (steps+1, n_nodes) array
        final_probs    : probability at final step
        classical_probs: classical random walk comparison
        mixing_time    : estimated steps to near-uniform dist
        hitting_times  : dict {target: hitting_time}
        graph_type     : name of graph
        n_nodes        : number of nodes
        steps          : total steps simulated
        walk_type      : "discrete" or "continuous"
        start_node     : starting node
        adjacency      : adjacency matrix
    """
    def __init__(self, probs_history, classical_probs,
                 graph_type, n_nodes, steps, walk_type,
                 start_node, adjacency):
        self.probs_history   = np.array(probs_history)
        self.final_probs     = self.probs_history[-1]
        self.classical_probs = classical_probs
        self.graph_type      = graph_type
        self.n_nodes         = n_nodes
        self.steps           = steps
        self.walk_type       = walk_type
        self.start_node      = start_node
        self.adjacency       = adjacency
        self.hitting_times   = {}
        self._compute_mixing_time()

    def _compute_mixing_time(self, epsilon: float = 0.05):
        uniform          = np.ones(self.n_nodes) / self.n_nodes
        self.mixing_time = self.steps
        for t, probs in enumerate(self.probs_history):
            if 0.5 * np.sum(np.abs(probs - uniform)) < epsilon:
                self.mixing_time = t
                break

    def compute_hitting_time(self, target_node: int) -> float:
        threshold = 1.0 / self.n_nodes
        for t, probs in enumerate(self.probs_history):
            if probs[target_node] > threshold:
                self.hitting_times[target_node] = float(t)
                return float(t)
        self.hitting_times[target_node] = float(self.steps)
        return float(self.steps)

    def summary(self):
        print("=" * 60)
        print("  PKTron Quantum Walk — Results")
        print("=" * 60)
        print(f"  Graph type   : {self.graph_type}")
        print(f"  Nodes        : {self.n_nodes}")
        print(f"  Walk type    : {self.walk_type}")
        print(f"  Start node   : {self.start_node}")
        print(f"  Steps        : {self.steps}")
        print(f"  Mixing time  : {self.mixing_time} steps")
        print()
        print(f"  {'Node':>5}  {'Quantum':>9}  {'Classical':>10}"
              f"  {'Diff':>8}")
        print(f"  {'-'*40}")
        for i, (qp, cp) in enumerate(
                zip(self.final_probs, self.classical_probs)):
            bar    = '▓' * int(qp * 20)
            marker = ' ◀ start' if i == self.start_node else ''
            print(f"  {i:>5}  {qp:>9.5f}  {cp:>10.5f}"
                  f"  {qp-cp:>+8.5f}  {bar}{marker}")
        print()
        uniform = 1.0 / self.n_nodes
        tv_q = 0.5 * np.sum(np.abs(self.final_probs - uniform))
        tv_c = 0.5 * np.sum(np.abs(self.classical_probs - uniform))
        print(f"  TV from uniform — Quantum: {tv_q:.5f} | "
              f"Classical: {tv_c:.5f}")
        print("=" * 60)


# ── Discrete Quantum Walk ─────────────────────────────────────────

class DiscreteQuantumWalk:
    """
    Discrete-time quantum walk on a graph.

    Step = Coin · Shift
    State space = |node⟩ ⊗ |direction⟩

    Args:
        adjacency : adjacency matrix (n_nodes × n_nodes)
        coin      : "hadamard", "grover", or custom numpy matrix
        start_node: starting node

    Example:
        A   = build_graph("cycle", 8)
        dqw = DiscreteQuantumWalk(A, coin="hadamard")
        result = dqw.run(steps=20)
    """

    def __init__(self, adjacency: np.ndarray,
                 coin: str = "hadamard",
                 start_node: int = 0):
        self.A          = np.array(adjacency)
        self.n_nodes    = len(adjacency)
        self.coin_type  = coin
        self.start_node = start_node

        # Degree of each node
        self.degrees = self.A.sum(axis=1).astype(int)
        max_deg      = int(self.degrees.max())

        # IMPORTANT: set dim and max_deg BEFORE calling
        # _build_coin/_build_shift which need them
        self.max_deg = max_deg
        self.dim     = self.n_nodes * self.max_deg

        # Build operators
        self._build_coin()
        self._build_shift()

    def _build_coin(self):
        """Build block-diagonal coin operator."""
        C = np.zeros((self.dim, self.dim), dtype=complex)
        for node in range(self.n_nodes):
            d = int(self.degrees[node])
            if d == 0:
                continue

            if self.coin_type == "hadamard":
                coin_mat = hadamard_coin(d)
            elif self.coin_type == "grover":
                coin_mat = grover_coin(d)
            else:
                coin_mat = np.array(self.coin_type, dtype=complex)

            # Pad to max_deg if needed
            if d < self.max_deg:
                padded = np.eye(self.max_deg, dtype=complex)
                padded[:d, :d] = coin_mat
                coin_mat = padded

            s = node * self.max_deg
            C[s:s+self.max_deg, s:s+self.max_deg] = coin_mat

        self.C = C

    def _build_shift(self):
        """Build conditional shift operator."""
        S = np.zeros((self.dim, self.dim), dtype=complex)
        for node in range(self.n_nodes):
            neighbours = np.where(self.A[node] > 0)[0]
            for dir_idx, neighbour in enumerate(neighbours):
                row_from    = node * self.max_deg + dir_idx
                nb_neighs   = np.where(self.A[neighbour] > 0)[0]
                rev_dirs    = np.where(nb_neighs == node)[0]
                rev_dir     = (rev_dirs[0] if len(rev_dirs) > 0
                               else dir_idx % max(1, int(
                                   self.degrees[neighbour])))
                row_to      = neighbour * self.max_deg + rev_dir
                S[row_to, row_from] = 1.0
        self.S = S

    def run(self, steps: int = 20) -> QuantumWalkResult:
        """Run the discrete walk for given steps."""
        # Initial state: uniform superposition over directions
        psi = np.zeros(self.dim, dtype=complex)
        d   = int(self.degrees[self.start_node])
        for k in range(d):
            psi[self.start_node * self.max_deg + k] = (
                1.0 / np.sqrt(max(d, 1)))

        U             = self.S @ self.C
        probs_history = [self._node_probs(psi)]

        for _ in range(steps):
            psi = U @ psi
            probs_history.append(self._node_probs(psi))

        return QuantumWalkResult(
            probs_history  = probs_history,
            classical_probs = self._classical_walk(steps),
            graph_type     = "custom",
            n_nodes        = self.n_nodes,
            steps          = steps,
            walk_type      = "discrete",
            start_node     = self.start_node,
            adjacency      = self.A,
        )

    def _node_probs(self, state: np.ndarray) -> np.ndarray:
        p = np.zeros(self.n_nodes)
        for node in range(self.n_nodes):
            sl = state[node*self.max_deg:(node+1)*self.max_deg]
            p[node] = float(np.sum(np.abs(sl)**2))
        return p

    def _classical_walk(self, steps: int) -> np.ndarray:
        T = np.zeros_like(self.A, dtype=float)
        for i in range(self.n_nodes):
            if self.degrees[i] > 0:
                T[:, i] = self.A[:, i] / self.degrees[i]
        p = np.zeros(self.n_nodes)
        p[self.start_node] = 1.0
        for _ in range(steps):
            p = T @ p
        return p


# ── Continuous Quantum Walk ───────────────────────────────────────

class ContinuousQuantumWalk:
    """
    Continuous-time quantum walk.

    Evolution: |ψ(t)⟩ = e^(-iHt)|ψ(0)⟩
    H = adjacency matrix or Laplacian

    Args:
        adjacency  : adjacency matrix
        hamiltonian: "adjacency" or "laplacian"
        start_node : starting node

    Example:
        A   = build_graph("hypercube", 8)
        cqw = ContinuousQuantumWalk(A)
        result = cqw.run(t_max=5.0, n_steps=50)
    """

    def __init__(self, adjacency: np.ndarray,
                 hamiltonian: str = "adjacency",
                 start_node: int = 0):
        self.A          = np.array(adjacency, dtype=float)
        self.n_nodes    = len(adjacency)
        self.start_node = start_node

        if hamiltonian == "adjacency":
            self.H = self.A.astype(complex)
        elif hamiltonian == "laplacian":
            D      = np.diag(self.A.sum(axis=1))
            self.H = (D - self.A).astype(complex)
        else:
            raise ValueError(
                "hamiltonian must be 'adjacency' or 'laplacian'")

        # Pre-diagonalise for fast exponentiation
        self._evals, self._evecs = np.linalg.eigh(self.H)

    def _evolve(self, t: float) -> np.ndarray:
        psi0   = np.zeros(self.n_nodes, dtype=complex)
        psi0[self.start_node] = 1.0
        phases = np.exp(-1j * self._evals * t)
        return self._evecs @ (phases * (
            self._evecs.conj().T @ psi0))

    def run(self, t_max: float = 5.0,
            n_steps: int = 50) -> QuantumWalkResult:
        """Evolve from t=0 to t=t_max."""
        times         = np.linspace(0, t_max, n_steps + 1)
        probs_history = [np.abs(self._evolve(t))**2
                         for t in times]
        result = QuantumWalkResult(
            probs_history  = probs_history,
            classical_probs = self._classical_walk(t_max, n_steps),
            graph_type     = "custom",
            n_nodes        = self.n_nodes,
            steps          = n_steps,
            walk_type      = f"continuous (t_max={t_max})",
            start_node     = self.start_node,
            adjacency      = self.A,
        )
        result.times = times
        return result

    def _classical_walk(self, t_max: float,
                        n_steps: int) -> np.ndarray:
        degrees = self.A.sum(axis=1)
        T       = np.zeros_like(self.A, dtype=float)
        for i in range(self.n_nodes):
            if degrees[i] > 0:
                T[:, i] = self.A[:, i] / degrees[i]
        p  = np.zeros(self.n_nodes)
        p[self.start_node] = 1.0
        dt = t_max / max(n_steps, 1)
        for _ in range(n_steps):
            p = p + dt * (T @ p - p)
            p = np.maximum(p, 0)
            s = p.sum()
            if s > 1e-12:
                p /= s
        return p


# ── Unified interface ─────────────────────────────────────────────

class QuantumWalk:
    """
    Unified interface for discrete and continuous quantum walks
    on any graph.

    Args:
        graph_type : "line","cycle","grid","hypercube",
                     "complete","custom"
        n_nodes    : number of nodes
        walk_type  : "discrete" or "continuous"
        coin       : "hadamard" or "grover" (discrete only)
        hamiltonian: "adjacency" or "laplacian" (continuous only)
        start_node : starting node
        edges      : list of (i,j) tuples for custom graphs

    Example:
        qw = QuantumWalk("cycle", n_nodes=8)
        result = qw.run(steps=30)
        result.summary()
        plot_walk(result)
    """

    def __init__(self, graph_type: str = "cycle",
                 n_nodes: int = 8,
                 walk_type: str = "discrete",
                 coin: str = "hadamard",
                 hamiltonian: str = "adjacency",
                 start_node: int = 0,
                 edges: list = None):
        self.graph_type  = graph_type
        self.n_nodes     = n_nodes
        self.walk_type   = walk_type
        self.start_node  = start_node

        self.A = build_graph(graph_type, n_nodes, edges)

        if walk_type == "discrete":
            self._walker = DiscreteQuantumWalk(
                self.A, coin=coin, start_node=start_node)
        elif walk_type == "continuous":
            self._walker = ContinuousQuantumWalk(
                self.A, hamiltonian=hamiltonian,
                start_node=start_node)
        else:
            raise ValueError(
                "walk_type must be 'discrete' or 'continuous'")

        print(f"[QuantumWalk] {walk_type.capitalize()} walk on "
              f"{graph_type} graph ({n_nodes} nodes) | "
              f"start={start_node}")

    def run(self, steps: int = 20, t_max: float = 5.0,
            n_steps: int = 50) -> QuantumWalkResult:
        """
        Run the walk.
        Discrete  → uses `steps`
        Continuous → uses `t_max` and `n_steps`
        """
        if self.walk_type == "discrete":
            result = self._walker.run(steps=steps)
        else:
            result = self._walker.run(t_max=t_max,
                                      n_steps=n_steps)
        result.graph_type = self.graph_type
        return result


# ── Hitting time analysis ─────────────────────────────────────────

def hitting_time_analysis(result: QuantumWalkResult,
                          targets: list = None) -> dict:
    """
    Compute hitting times to target nodes.

    Hitting time = first step where P(target) > 1/n_nodes

    Args:
        result : QuantumWalkResult
        targets: list of node indices (default: all nodes)

    Returns:
        dict {node: hitting_time}

    Example:
        ht = hitting_time_analysis(result, targets=[3,7])
    """
    if targets is None:
        targets = list(range(result.n_nodes))

    print("=" * 52)
    print("  Hitting Time Analysis")
    print("=" * 52)
    print(f"  {'Target':>6}  {'Hitting Time':>13}  Status")
    print(f"  {'-'*42}")

    hitting = {}
    for t in targets:
        ht     = result.compute_hitting_time(t)
        status = "✅ Hit" if ht < result.steps else "❌ Not reached"
        marker = " ◀ start" if t == result.start_node else ""
        hitting[t] = ht
        print(f"  {t:>6}  {ht:>12.0f}  {status}{marker}")

    print("=" * 52)
    print(f"  Average hitting time: "
          f"{np.mean(list(hitting.values())):.2f} steps")
    return hitting


# ── Plotting ──────────────────────────────────────────────────────

def plot_walk(result: QuantumWalkResult):
    """
    4-panel visualization:
      1. Probability evolution heatmap
      2. Final quantum vs classical distribution
      3. P(start node) over time with mixing time
      4. Graph adjacency matrix

    Args:
        result: QuantumWalkResult
    """
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(2, 2, figsize=(14, 10))

    # Panel 1: Heatmap
    ax  = axes[0, 0]
    im  = ax.imshow(result.probs_history.T, aspect='auto',
                    cmap='YlOrRd', interpolation='nearest')
    ax.set_xlabel("Step")
    ax.set_ylabel("Node")
    ax.set_title("Probability Evolution\n(bright = high prob)")
    plt.colorbar(im, ax=ax, label="Probability")
    ax.axhline(result.start_node, color='cyan',
               linestyle='--', linewidth=1.2, alpha=0.8,
               label=f'Start={result.start_node}')
    ax.legend(fontsize=8)

    # Panel 2: Final distribution
    ax2   = axes[0, 1]
    nodes = np.arange(result.n_nodes)
    w     = 0.35
    ax2.bar(nodes - w/2, result.final_probs, w,
            label='Quantum', color='steelblue',
            edgecolor='navy', linewidth=0.5)
    ax2.bar(nodes + w/2, result.classical_probs, w,
            label='Classical', color='tomato',
            edgecolor='darkred', linewidth=0.5)
    ax2.axhline(1/result.n_nodes, color='green',
                linestyle='--', alpha=0.7, label='Uniform')
    ax2.set_xlabel("Node")
    ax2.set_ylabel("Probability")
    ax2.set_title("Final: Quantum vs Classical")
    ax2.legend(fontsize=8)
    ax2.set_xticks(nodes)
    ax2.grid(axis='y', alpha=0.3)

    # Panel 3: P(start) over time
    ax3     = axes[1, 0]
    p_start = result.probs_history[:, result.start_node]
    ax3.plot(p_start, color='steelblue', linewidth=1.8,
             label=f'P(node {result.start_node})')
    ax3.axhline(1/result.n_nodes, color='green',
                linestyle='--', alpha=0.7, label='Uniform')
    ax3.axvline(result.mixing_time, color='orange',
                linestyle=':', alpha=0.9,
                label=f'Mix time={result.mixing_time}')
    ax3.set_xlabel("Step")
    ax3.set_ylabel("Probability")
    ax3.set_title(f"P(start node {result.start_node}) over Time")
    ax3.legend(fontsize=8)
    ax3.grid(alpha=0.3)

    # Panel 4: Adjacency matrix
    ax4 = axes[1, 1]
    ax4.imshow(result.adjacency, cmap='Blues',
               interpolation='nearest')
    ax4.set_title(f"Graph: {result.graph_type} "
                  f"({result.n_nodes} nodes)")
    ax4.set_xlabel("Node")
    ax4.set_ylabel("Node")
    s = result.start_node
    for spine_ax in [ax4]:
        spine_ax.axhline(s - 0.5, color='red',
                         linewidth=1.2, alpha=0.6)
        spine_ax.axhline(s + 0.5, color='red',
                         linewidth=1.2, alpha=0.6)
        spine_ax.axvline(s - 0.5, color='red',
                         linewidth=1.2, alpha=0.6)
        spine_ax.axvline(s + 0.5, color='red',
                         linewidth=1.2, alpha=0.6)

    plt.suptitle(
        f"PKTron Quantum Walk — {result.graph_type} | "
        f"{result.walk_type} | {result.n_nodes} nodes",
        fontsize=12, fontweight='bold')
    plt.tight_layout()
    plt.savefig('pktron_quantum_walk.png', dpi=120,
                bbox_inches='tight')
    plt.show()
    print("[QuantumWalk] Plot saved as pktron_quantum_walk.png")


def plot_walk_animation(result: QuantumWalkResult,
                        max_frames: int = 20):
    """
    Interactive Plotly animation of probability spreading.
    Play/pause and step slider included.
    """
    try:
        import plotly.graph_objects as go

        n      = result.n_nodes
        nodes  = list(range(n))
        step   = max(1, len(result.probs_history) // max_frames)
        frames_data = result.probs_history[::step]

        frames = [
            go.Frame(
                data=[go.Bar(
                    x=nodes,
                    y=probs.tolist(),
                    marker_color='steelblue',
                    text=[f"{p:.3f}" for p in probs],
                    textposition='outside',
                )],
                name=str(i),
                layout=go.Layout(
                    title_text=(f"Quantum Walk — "
                                f"Step {i*step}"))
            )
            for i, probs in enumerate(frames_data)
        ]

        fig = go.Figure(
            data=[go.Bar(x=nodes,
                         y=frames_data[0].tolist(),
                         marker_color='steelblue')],
            frames=frames,
            layout=go.Layout(
                title="PKTron Quantum Walk Animation",
                xaxis_title="Node",
                yaxis_title="Probability",
                yaxis=dict(range=[0, 1.05]),
                updatemenus=[dict(
                    type="buttons",
                    showactive=False,
                    buttons=[
                        dict(label="▶ Play",
                             method="animate",
                             args=[None,
                                   {"frame": {"duration": 300},
                                    "fromcurrent": True}]),
                        dict(label="⏸ Pause",
                             method="animate",
                             args=[[None],
                                   {"mode": "immediate"}]),
                    ]
                )],
                sliders=[dict(
                    steps=[
                        dict(method="animate",
                             args=[[str(i)],
                                   {"mode": "immediate"}],
                             label=f"t={i*step}")
                        for i in range(len(frames))
                    ],
                    currentvalue={"prefix": "Step: "},
                )]
            )
        )
        fig.show()

    except ImportError:
        print("[QuantumWalk] plotly not available. "
              "Use plot_walk() instead.")
