"""Datasets and plots for the live page on graphs as data (Homework 7, Problem 1).

The live page downloads this file into the in-browser Python session and runs
`from graph_tools import *`, so every exercise starts from the same helpers.
"""

import numpy as np
import matplotlib.pyplot as plt
import networkx as nx


# ---------------------------------------------------------------- datasets

def barbell_edges():
    """The barbell graph of the lecture, with vertices numbered 0, ..., 5.

    Lecture vertex k is Python vertex k - 1.
    """
    return [(0, 1), (0, 2), (1, 2), (3, 4), (3, 5), (4, 5), (2, 3)]


# Positions matching the barbell figure in the lecture.
BARBELL_POS = np.array([[0, 1.2], [0, 0], [1.0, 0.6], [2.4, 0.6], [3.4, 1.2], [3.4, 0]])


def two_rings(n=200, seed=0):
    """Two noisy concentric circles; the first n // 2 points are the inner ring."""
    rng = np.random.default_rng(seed)
    radius = np.where(np.arange(n) < n // 2, 1.0, 3.0)
    angle = rng.uniform(0, 2 * np.pi, n)
    ring = np.column_stack([radius * np.cos(angle), radius * np.sin(angle)])
    return ring + 0.15 * rng.standard_normal((n, 2))


def knn_graph(X, m):
    """Adjacency matrix of the m-nearest-neighbour graph of the rows of X.

    {i, j} is an edge when x_j is among the m nearest points to x_i, or vice versa.
    """
    sq = ((X[:, None, :] - X[None, :, :]) ** 2).sum(axis=-1)
    np.fill_diagonal(sq, np.inf)
    nearest = np.argsort(sq, axis=1)[:, :m]
    n = len(X)
    W = np.zeros((n, n))
    W[np.repeat(np.arange(n), m), nearest.ravel()] = 1
    return np.maximum(W, W.T)


def karate_club():
    """Zachary's karate club: adjacency matrix and the side each member joined.

    Returns W (34 x 34, unweighted) and a boolean array that is True for the
    members who joined Mr. Hi and False for those who stayed with the Officer.
    """
    G = nx.karate_club_graph()
    W = nx.to_numpy_array(G, nodelist=range(34), weight=None)
    mr_hi = np.array([G.nodes[i]["club"] == "Mr. Hi" for i in range(34)])
    return W, mr_hi


# ---------------------------------------------------------------- plots

def draw(W, pos=None, colors=None, title="", labels=True, ax=None, node_size=300, seed=0):
    """Draw the graph with adjacency matrix W.

    pos      n x 2 array of vertex positions; if None, a spring layout is used
    colors   length-n array used to color the vertices (numbers or booleans)
    labels   whether to write the vertex numbers on the vertices
    """
    G = nx.from_numpy_array(np.asarray(W))
    if pos is None:
        layout = nx.spring_layout(G, seed=seed)
    else:
        layout = {i: pos[i] for i in range(len(pos))}
    if ax is None:
        _, ax = plt.subplots(figsize=(6, 4.5))
    if colors is None:
        color_args = {"node_color": "#9ecae1"}
    else:
        color_args = {"node_color": np.asarray(colors, dtype=float), "cmap": "coolwarm"}
    nx.draw_networkx_edges(G, layout, ax=ax, edge_color="#999999", alpha=0.6)
    nx.draw_networkx_nodes(G, layout, ax=ax, node_size=node_size,
                           edgecolors="#444444", linewidths=0.5, **color_args)
    if labels:
        nx.draw_networkx_labels(G, layout, ax=ax, font_size=8)
    ax.set_title(title)
    ax.set_aspect("equal", adjustable="datalim")
    ax.axis("off")
    return ax


def draw_karate(W, fiedler, mr_hi, seed=0):
    """The karate club colored by the Fiedler vector; squares joined Mr. Hi."""
    G = nx.from_numpy_array(W)
    layout = nx.spring_layout(G, seed=seed)
    fig, ax = plt.subplots(figsize=(7, 5.5))
    nx.draw_networkx_edges(G, layout, ax=ax, edge_color="#999999", alpha=0.5)
    vmax = np.abs(fiedler).max()
    for side, shape in [(True, "s"), (False, "o")]:
        nodes = [i for i in range(len(W)) if mr_hi[i] == side]
        nx.draw_networkx_nodes(G, layout, nodelist=nodes, ax=ax, node_shape=shape,
                               node_color=fiedler[nodes], cmap="coolwarm", vmin=-vmax, vmax=vmax,
                               node_size=320, edgecolors="#444444", linewidths=0.5)
    nx.draw_networkx_labels(G, layout, ax=ax, font_size=8)
    ax.set_title("color: Fiedler vector    shape: square = joined Mr. Hi")
    ax.axis("off")
    plt.show()


def plot_spectrum(lam, count=10, title=""):
    """Plot the smallest `count` eigenvalues of a Laplacian."""
    fig, ax = plt.subplots(figsize=(5, 3))
    ax.plot(np.arange(1, count + 1), lam[:count], "o-")
    ax.axhline(0, color="#999999", lw=0.8)
    ax.set_xlabel("index")
    ax.set_ylabel("eigenvalue")
    ax.set_xticks(np.arange(1, count + 1))
    ax.set_title(title)
    ax.grid(alpha=0.3)
    fig.tight_layout()
    plt.show()