import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
from matplotlib.cm import ScalarMappable
import os

def run_chawla_ranking(file_path=None, steps=100, coin_type='fourier'):
    """
    Discrete-time quantum walk algorithm for ranking nodes on a network
    Chawla, P ; Mangal, R  ; Chandrashekar, CM
    QUANTUM INFORMATION PROCESSING
    19(5): 158, 2020.
    """
    if file_path and os.path.exists(file_path):
        G = nx.read_edgelist(file_path, create_using=nx.DiGraph())
    else:
        G = nx.gnp_random_graph(15, 0.2, directed=True, seed=42)

    nodes = list(G.nodes())
    N = len(nodes)
    mapping = {node: i for i, node in enumerate(nodes)}
    edges = list(G.edges())
    M = len(edges)

    if M == 0: raise ValueError("Graph must have edges.")

    idx_map = {edge: i for i, edge in enumerate(edges)}
    U = np.zeros((M, M), dtype=np.complex128)

    for v in nodes:
        in_edges = [e for e in edges if e[1] == v]
        out_edges = [e for e in edges if e[0] == v]
        d_in = len(in_edges); d_out = len(out_edges)

        if d_in > 0 and d_out > 0:
            if coin_type == 'fourier':
                for k in range(d_out):
                    for l in range(d_in):
                        phase = np.exp(2j * np.pi * k * l / max(d_in, d_out))
                        U[idx_map[out_edges[k]], idx_map[in_edges[l]]] = phase / np.sqrt(d_in)
            else:
                for k in range(d_out):
                    for l in range(d_in):
                        U[idx_map[out_edges[k]], idx_map[in_edges[l]]] = np.sqrt(1.0 / d_in)

    state = np.ones(M, dtype=np.complex128) / np.sqrt(M)
    prob_accumulator = np.zeros(M)
    curr_state = state

    for t in range(steps):
        curr_state = np.dot(U, curr_state)
        norm = np.linalg.norm(curr_state)
        if norm > 0: curr_state /= norm
        prob_accumulator += np.abs(curr_state) ** 2

    node_ranks = np.zeros(N)
    for i, (u, v) in enumerate(edges):
        node_ranks[mapping[v]] += prob_accumulator[i] / steps

    pos = nx.kamada_kawai_layout(G)
    fig, ax = plt.subplots(figsize=(6, 4))
    cmap = plt.cm.YlOrRd
    norm = Normalize(vmin=node_ranks.min(), vmax=node_ranks.max())

    nx.draw_networkx_edges(G, pos, ax=ax, edge_color='#dddddd', alpha=0.7, arrowsize=12, connectionstyle='arc3,rad=0.1')
    nx.draw_networkx_nodes(G, pos, ax=ax, node_size=220, node_color=node_ranks, cmap=cmap, edgecolors='#444444', linewidths=0.8)

    for i, node in enumerate(nodes):
        score = node_ranks[i]
        rgb = cmap(norm(score))[:3]
        brightness = 0.299 * rgb[0] + 0.587 * rgb[1] + 0.114 * rgb[2]
        t_color = 'white' if brightness < 0.45 else 'black'
        ax.text(pos[node][0], pos[node][1], s=str(node), color=t_color, fontsize=8, ha='center', va='center', fontweight='bold')

    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    cb = plt.colorbar(sm, ax=ax, fraction=0.046, pad=0.04)
    cb.outline.set_visible(False)
    cb.ax.tick_params(labelsize=8)
    ax.axis('off')
    plt.tight_layout()
    plt.show()
    return node_ranks

def run_szegedy_service(G=None, steps=10, initial_type='uniform', initial_node=None):
    """
    Quantum speed-up of Markov chain based algorithms
    Szegedy M
    Proceedings of the 45th Annual IEEE Symposium on Foundations of Computer Science
    2004: 32–41.
    """
    if G is None: G = nx.karate_club_graph()
    nodes = list(G.nodes())
    N = len(nodes)
    mapping = {node: i for i, node in enumerate(nodes)}
    A = nx.to_numpy_array(G, nodelist=nodes)
    degrees = np.sum(A, axis=0)
    degrees[degrees == 0] = 1
    P = A / degrees

    psi = np.zeros((N * N, N), dtype=np.complex128)
    for i in range(N):
        vec_psi_i = np.zeros((N, N), dtype=np.complex128)
        vec_psi_i[:, i] = np.sqrt(P[:, i])
        psi[:, i] = vec_psi_i.flatten()

    S = 2 * np.dot(psi, psi.conj().T) - np.eye(N * N)
    SWAP = np.zeros((N * N, N * N))
    for i in range(N):
        for j in range(N): SWAP[i * N + j, j * N + i] = 1
    U = np.dot(SWAP, S)

    if initial_type == 'node' and initial_node in mapping:
        idx = mapping[initial_node]
        init_matrix = np.zeros((N, N), dtype=np.complex128)
        init_matrix[:, idx] = np.sqrt(P[:, idx])
        state_vec = init_matrix.flatten()
    else:
        state_vec = np.sum(psi, axis=1) / np.sqrt(N)

    for _ in range(steps):
        state_vec = np.dot(U, state_vec)

    node_probs = np.sum(np.abs(state_vec.reshape((N, N))) ** 2, axis=1)
    node_probs /= np.sum(node_probs)

    pos = nx.spring_layout(G, seed=42)
    fig, ax = plt.subplots(figsize=(6, 4))
    cmap = plt.cm.YlGnBu
    norm = Normalize(vmin=node_probs.min(), vmax=node_probs.max())
    nx.draw_networkx_edges(G, pos, ax=ax, edge_color='#bbbbbb', width=1.0)
    nx.draw_networkx_nodes(G, pos, ax=ax, node_size=180, node_color=node_probs, cmap=cmap, edgecolors='#444444', linewidths=1.0)

    for node in nodes:
        prob = node_probs[mapping[node]]
        rgb = cmap(norm(prob))[:3]
        brightness = 0.299 * rgb[0] + 0.587 * rgb[1] + 0.114 * rgb[2]
        t_color = 'white' if brightness < 0.38 else 'black'
        ax.text(pos[node][0], pos[node][1], s=str(node), color=t_color, fontsize=9, ha='center', va='center')

    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    plt.colorbar(sm, ax=ax, fraction=0.046, pad=0.04).outline.set_visible(False)
    ax.axis('off')
    plt.tight_layout()
    plt.show()
    return node_probs

def run_community_detection(G=None, initial_state='uniform', coin_type='fourier'):
    """
    Discrete-time quantum walk on complex networks for community detection
    Kanae Mukai and Naomichi Hatano
    Phys. Rev. Research 2, 023378 (2020)
    """
    if G is None: G = nx.karate_club_graph()
    nodes = list(G.nodes())
    N = len(nodes)
    mapping = {node: i for i, node in enumerate(nodes)}
    edges = list(G.edges())

    dir_edges = []
    for u, v in edges:
        dir_edges.append((u, v)); dir_edges.append((v, u))
    idx_map = {edge: i for i, edge in enumerate(dir_edges)}
    dim = len(dir_edges)

    C = np.zeros((dim, dim), dtype=np.complex128)
    for j in nodes:
        neighbors = list(G.neighbors(j))
        d_j = len(neighbors)
        if d_j == 0: continue
        
        if coin_type.lower() == 'fourier':
            coin_matrix = np.exp(2j * np.pi * np.outer(np.arange(d_j), np.arange(d_j)) / d_j) / np.sqrt(d_j)
        elif coin_type.lower() == 'grover':
            coin_matrix = np.full((d_j, d_j), 2.0 / d_j) - np.eye(d_j)
        else: raise ValueError("Unsupported coin_type")

        for k in range(d_j):
            for l in range(d_j):
                row = idx_map[(j, neighbors[k])]
                col = idx_map[(neighbors[l], j)]
                C[row, col] = coin_matrix[k, l]

    S = np.zeros((dim, dim))
    for i, j in dir_edges:
        S[idx_map[(j, i)], idx_map[(i, j)]] = 1
    U = S @ C

    if initial_state == 'uniform':
        phi_0 = np.ones(dim, dtype=np.complex128) / np.sqrt(dim)
    else:
        phi_0 = np.zeros(dim, dtype=np.complex128)
        # Note: 'initial_state' when not uniform implies an index or node logic, keeping simple here
        # Assuming initial_state passed is a valid node index if it's an integer
        if isinstance(initial_state, int) and 0 <= initial_state < N:
             seed_neighbors = list(G.neighbors(nodes[initial_state]))
             for nb in seed_neighbors:
                 phi_0[idx_map[(nodes[initial_state], nb)]] = 1.0 / np.sqrt(len(seed_neighbors))

    vals, vecs = np.linalg.eig(U)
    unique_vals, inverse_indices = np.unique(np.round(vals, 10), return_inverse=True)
    inf_probs = np.zeros(dim)

    for i in range(len(unique_vals)):
        indices = np.where(inverse_indices == i)[0]
        subspace_vecs = vecs[:, indices]
        proj = subspace_vecs @ (np.conj(subspace_vecs.T) @ phi_0)
        inf_probs += np.abs(proj) ** 2

    node_scores = np.zeros(N)
    for i, (u, v) in enumerate(dir_edges):
        node_scores[mapping[v]] += np.real(inf_probs[i])

    return dict(zip(nodes, node_scores))

def run_quantum_pagerank(file_path=None, steps=20, alpha=0.85, initial_type='uniform'):
    """
    Quantum Google in a Complex Network (Paparo et al.)
    """
    if file_path and os.path.exists(file_path):
        G = nx.read_edgelist(file_path, create_using=nx.Graph())
    else:
        G = nx.karate_club_graph()

    nodes = list(G.nodes())
    N = len(nodes)
    mapping = {node: i for i, node in enumerate(nodes)}
    A = nx.to_numpy_array(G, nodelist=nodes)
    degrees = np.sum(A, axis=1)

    P_classic = np.zeros((N, N))
    for i in range(N):
        if degrees[i] > 0: P_classic[:, i] = A[i, :] / degrees[i]
        else: P_classic[:, i] = 1.0 / N
    P = alpha * P_classic + (1 - alpha) / N * np.ones((N, N))

    psi = np.zeros((N * N, N), dtype=np.complex128)
    for j in range(N):
        column_j = np.zeros((N, N), dtype=np.complex128)
        column_j[:, j] = np.sqrt(P[:, j])
        psi[:, j] = column_j.flatten()

    S = 2 * np.dot(psi, psi.conj().T) - np.eye(N * N)
    SWAP = np.zeros((N * N, N * N))
    for i in range(N):
        for j in range(N): SWAP[i * N + j, j * N + i] = 1
    U = np.dot(SWAP, S)

    state = np.sum(psi, axis=1) / np.sqrt(N) if initial_type == 'uniform' else psi[:, 0]
    prob_accumulator = np.zeros(N)
    curr_state = state

    for t in range(1, steps + 1):
        curr_state = np.dot(U, curr_state)
        probs = np.sum(np.abs(curr_state.reshape((N, N))) ** 2, axis=1)
        prob_accumulator += probs
    q_pagerank_scores = prob_accumulator / steps

    pos = nx.spring_layout(G, seed=42)
    fig, ax = plt.subplots(figsize=(6, 4))
    cmap = plt.cm.YlGnBu
    norm = Normalize(vmin=q_pagerank_scores.min(), vmax=q_pagerank_scores.max())
    nx.draw_networkx_edges(G, pos, ax=ax, edge_color='#bbbbbb', alpha=0.4, arrows=False)
    nx.draw_networkx_nodes(G, pos, ax=ax, node_size=180, node_color=q_pagerank_scores, cmap=cmap, edgecolors='#444444', linewidths=0.8)

    for node in nodes:
        score = q_pagerank_scores[mapping[node]]
        rgb = cmap(norm(score))[:3]
        brightness = 0.299 * rgb[0] + 0.587 * rgb[1] + 0.114 * rgb[2]
        t_color = 'white' if brightness < 0.38 else 'black'
        ax.text(pos[node][0], pos[node][1], s=str(node), color=t_color, fontsize=8, ha='center', va='center')

    plt.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=norm), ax=ax, fraction=0.046, pad=0.04)
    ax.axis('off')
    plt.tight_layout()
    plt.show()
    return q_pagerank_scores

def run_hadamard(steps=100, top_p=0.05):
    """
    A Hadamard walk model and its application in identification of important edges in complex networks
    Liang et al.
    """
    G = nx.karate_club_graph()
    nodes, edges = list(G.nodes()), list(G.edges())
    N, M = len(nodes), len(edges)

    directed_edges = []
    for u, v in edges: directed_edges.extend([(u, v), (v, u)])
    dim = len(directed_edges)
    idx_map = {edge: i for i, edge in enumerate(directed_edges)}

    H = np.array([[1, 1], [1, -1]]) / np.sqrt(2)
    C = np.zeros((dim, dim))
    for i in range(0, dim, 2): C[i:i + 2, i:i + 2] = H
    S = np.zeros((dim, dim))
    for u, v in directed_edges: S[idx_map[(v, u)], idx_map[(u, v)]] = 1
    U = S @ C

    psi = np.ones(dim, dtype=np.complex128) / np.sqrt(dim)
    node_probs_accum = np.zeros(N)
    for _ in range(steps):
        psi = U @ psi
        curr_p = np.abs(psi) ** 2
        for i, (u, v) in enumerate(directed_edges):
            node_probs_accum[nodes.index(v)] += curr_p[i]
    node_avg_p = node_probs_accum / steps

    edge_scores = np.array([node_avg_p[nodes.index(u)] + node_avg_p[nodes.index(v)] for u, v in edges])
    num_top = max(1, int(M * top_p))
    top_indices = np.argsort(edge_scores)[-num_top:]
    top_edges = [edges[i] for i in top_indices]
    top_scores = edge_scores[top_indices]

    pos = nx.spring_layout(G, seed=42)
    fig, ax = plt.subplots(figsize=(10, 7))
    target_cmap = plt.cm.plasma
    norm = Normalize(vmin=top_scores.min(), vmax=top_scores.max())

    nx.draw_networkx_nodes(G, pos, ax=ax, node_size=250, node_color='green', edgecolors='white', linewidths=0.6)
    nx.draw_networkx_edges(G, pos, ax=ax, edgelist=edges, edge_color='#444444', width=0.8, alpha=0.3, arrows=False)
    nx.draw_networkx_edges(G, pos, ax=ax, edgelist=top_edges, edge_color=top_scores, edge_cmap=target_cmap,
                           edge_vmin=top_scores.min(), edge_vmax=top_scores.max(), width=2.5, alpha=1.0, arrows=False)
    nx.draw_networkx_labels(G, pos, font_size=8, font_color='white', font_weight='bold')
    ax.set_title(f"Hadamard Walk: Top {num_top}", fontsize=12)
    ax.axis('off')
    plt.tight_layout()
    plt.show()
    return edge_scores
