"""Datasets, embeddings, and plots for the live page on spectral embeddings (Homework 8, Problem 1).

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

import numpy as np
import matplotlib.pyplot as plt


# ---------------------------------------------------------------- graphs

def pairwise_sq_dists(X):
    """The n x n matrix of squared distances between the rows of X."""
    sq = (X**2).sum(axis=1)
    return np.maximum(sq[:, None] + sq[None, :] - 2 * X @ X.T, 0)


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 = pairwise_sq_dists(X)
    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 gaussian_graph(X, eps):
    """Gaussian similarity graph: w_ij = exp(-|x_i - x_j|^2 / (2 eps)), no self-loops."""
    W = np.exp(-pairwise_sq_dists(X) / (2 * eps))
    np.fill_diagonal(W, 0)
    return W


def laplacian(W):
    """The Laplacian L = D - W."""
    return np.diag(W.sum(axis=1)) - W


def normalized_embedding(W, k):
    """The k smallest eigenvalues of L_sym and the matrix Phi = [phi_1, ..., phi_k].

    phi_l = D^{-1/2} v_l, where v_l are orthonormal eigenvectors of
    L_sym = I - D^{-1/2} W D^{-1/2}.  Row i of Phi is the spectral embedding of vertex i.
    """
    d = W.sum(axis=1)
    L_sym = np.eye(len(d)) - W / np.sqrt(np.outer(d, d))
    lam, V = np.linalg.eigh(L_sym)    # increasing order
    return lam[:k], V[:, :k] / np.sqrt(d)[:, None]


def eigenmap(W, m):
    """The Laplacian eigenmap in m dimensions: Phi without its first column."""
    lam, Phi = normalized_embedding(W, m + 1)
    return lam, Phi[:, 1:]


def energy(Y, W):
    """E(Y) = (1/2) sum_{i,j} w_ij |y_i - y_j|^2, computed straight from the definition."""
    return 0.5 * np.sum(W * pairwise_sq_dists(Y))


def random_orthogonal(m, seed=0):
    """A random m x m orthogonal matrix (the Q factor of a random Gaussian matrix)."""
    Q, _ = np.linalg.qr(np.random.default_rng(seed).standard_normal((m, m)))
    return Q


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

def blobs_with_outlier(outlier=(2.0, 4.0), n=100, seed=0):
    """Two Gaussian blobs of n points each, centred at (0, 0) and (4, 0), plus one outlier.

    Returns X and labels: 0 and 1 for the blobs, 2 for the outlier (the last row).
    """
    rng = np.random.default_rng(seed)
    X = np.vstack([rng.normal([0, 0], 0.6, (n, 2)),
                   rng.normal([4, 0], 0.6, (n, 2)),
                   [outlier]])
    return X, np.r_[np.zeros(n), np.ones(n), 2].astype(int)


def rectangle(a, n=500, seed=0):
    """n points uniform in the rectangle [0, a] x [0, 1]."""
    return np.random.default_rng(seed).uniform(size=(n, 2)) * [a, 1]


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

def scatter(Y, colors=None, title="", ax=None, cmap="viridis", size=10, equal=True):
    """Scatter plot of the first two columns of Y."""
    if ax is None:
        _, ax = plt.subplots(figsize=(4.5, 4))
    if colors is not None:
        colors = np.asarray(colors, dtype=float)    # booleans and integer labels work too
    ax.scatter(Y[:, 0], Y[:, 1], c=colors, cmap=cmap, s=size)
    ax.set_title(title, fontsize=10)
    if equal:
        ax.set_aspect("equal", adjustable="datalim")
    ax.grid(alpha=0.3)
    return ax


def panels(count, width=4.0, height=3.6):
    """A row of `count` empty axes."""
    fig, axes = plt.subplots(1, count, figsize=(width * count, height))
    return fig, np.atleast_1d(axes)