"""Lesson 5: reversible, lazy diffusion on a symmetric kNN graph."""
from pathlib import Path
import argparse
import json
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy import sparse
from scipy.sparse.csgraph import connected_components
from scipy.sparse.linalg import eigsh
from sklearn.neighbors import NearestNeighbors
from common import make_dataset, prepare_embedding


def affinity_graph(Z, k=48):
    k = min(k, len(Z) - 1)
    distances, neighbors = NearestNeighbors(n_neighbors=k + 1).fit(Z).kneighbors(Z)
    distances, neighbors = distances[:, 1:], neighbors[:, 1:]
    sigma = max(float(np.median(distances[:, -1])), 1e-8)
    W = sparse.csr_matrix((np.exp(-(distances.ravel() / sigma) ** 2),
                           (np.repeat(np.arange(len(Z)), k), neighbors.ravel())),
                          shape=(len(Z), len(Z)))
    W = W.maximum(W.T)
    W.setdiag(0)
    W.eliminate_zeros()
    return W


def diffusion_coordinates(W, time=4, n_components=2):
    degree = np.asarray(W.sum(axis=1)).ravel()
    if np.any(degree <= 0):
        raise ValueError("Remove or analyze isolated vertices separately.")
    n_connected, _ = connected_components(W, directed=False)
    if n_connected != 1:
        raise ValueError(f"Graph has {n_connected} components; analyze components separately before a global diffusion map.")
    pi = degree / degree.sum()
    invsqrt = sparse.diags(1 / np.sqrt(degree))
    S = (sparse.eye(W.shape[0]) + invsqrt @ W @ invsqrt) / 2
    if not 1 <= n_components < W.shape[0]:
        raise ValueError("n_components must be between 1 and number of vertices minus 1.")
    n_eigs = n_components + 1
    if n_eigs >= W.shape[0]:
        values, vectors = np.linalg.eigh(S.toarray())
    else:
        values, vectors = eigsh(S, k=n_eigs, which="LA", v0=np.ones(W.shape[0]), tol=1e-9)
    order = np.argsort(values)[::-1]
    values, vectors = values[order], vectors[:, order]
    # Connectedness gives exactly one constant mode. Near-one modes can carry
    # the most important bottleneck geometry and must not be thresholded away.
    keep = np.arange(1, n_components + 1)
    values = values[keep]
    psi = vectors[:, keep] / np.sqrt(pi[:, None])
    coordinates = psi * values[None, :] ** time
    P = (sparse.eye(W.shape[0]) + sparse.diags(1 / degree) @ W) / 2
    return coordinates, values, P, pi, int(n_connected)


def run(output_dir: Path) -> dict:
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    data = make_dataset()
    _, Z, _, _ = prepare_embedding(data, n_components=6)
    W = affinity_graph(Z)
    D, values, P, pi, n_connected = diffusion_coordinates(W)
    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
    for ax, coordinates, title in zip(axes, [Z[:, :2], D], ["PCA, first two axes", "Diffusion coordinates, t=4"]):
        im = ax.scatter(*coordinates.T, c=data["time"], s=5, cmap="viridis")
        ax.set_title(title)
        ax.set_xlabel("coordinate 1")
        ax.set_ylabel("coordinate 2")
    fig.colorbar(im, ax=axes, label="Observed sampling time (evaluation only)")
    fig.savefig(output_dir / "lesson05.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    return {"connected_components": n_connected,
            "row_sum_max_error": float(np.max(np.abs(np.asarray(P.sum(axis=1)).ravel() - 1))),
            "stationarity_max_error": float(np.max(np.abs(np.asarray(pi @ P).ravel() - pi))),
            "nontrivial_eigenvalues": values.tolist(), "diffusion_time": 4,
            "evaluation": "transductive snapshot geometry; no future prediction"}


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=Path("outputs"))
    print(json.dumps(run(parser.parse_args().output), indent=2, allow_nan=False))