"""Lesson 7: an MST teaching baseline, explicitly not PAGA."""
from pathlib import Path
import argparse
import json
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.spatial.distance import cdist
from scipy.sparse.csgraph import minimum_spanning_tree, shortest_path
from scipy.stats import spearmanr
from sklearn.cluster import KMeans
from common import make_dataset, prepare_embedding


def project_tree(Z, centers, tree, root):
    edge_a, edge_b = tree.nonzero()
    starts, ends = centers[edge_a], centers[edge_b]
    vectors = ends - starts
    lengths = np.linalg.norm(vectors, axis=1)
    fractions = np.einsum("ned,ed->ne", Z[:, None, :] - starts, vectors) / lengths[None, :] ** 2
    fractions = np.clip(fractions, 0, 1)
    projections = starts[None, :, :] + fractions[:, :, None] * vectors[None, :, :]
    squared_errors = np.sum((Z[:, None, :] - projections) ** 2, axis=2)
    chosen = np.argmin(squared_errors, axis=1)
    u = fractions[np.arange(len(Z)), chosen]
    root_distances = shortest_path(tree, directed=False, indices=root)
    pseudotime = np.minimum(root_distances[edge_a[chosen]] + u * lengths[chosen],
                            root_distances[edge_b[chosen]] + (1 - u) * lengths[chosen])
    return pseudotime, squared_errors[np.arange(len(Z)), chosen]


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)
    km = KMeans(n_clusters=14, n_init=10, random_state=7).fit(Z[data["train"]])
    centers = km.cluster_centers_
    distances = cdist(centers, centers)
    tree = minimum_spanning_tree(distances)
    initial = data["train"] & np.isclose(data["time"], np.min(data["time"]))
    root = int(np.argmin(np.linalg.norm(centers - Z[initial].mean(axis=0), axis=1)))
    pseudotime, error = project_tree(Z, centers, tree, root)
    opposite = int(np.argmax(shortest_path(tree, directed=False, indices=root)))
    reversed_time, _ = project_tree(Z, centers, tree, opposite)
    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
    for ax, color, selected_root, title in zip(axes, [pseudotime, reversed_time], [root, opposite], ["Root: known initial samples", "Root: farthest tree vertex"]):
        im = ax.scatter(Z[:, 0], Z[:, 1], c=color, s=5, cmap="viridis")
        for a, b in zip(*tree.nonzero()):
            ax.plot(centers[[a, b], 0], centers[[a, b], 1], "k-", lw=1)
        ax.scatter(*centers[selected_root, :2], marker="*", c="red", s=120)
        ax.set_title(title)
        ax.set_xlabel("PC1")
        ax.set_ylabel("PC2")
        fig.colorbar(im, ax=ax, label="Tree distance")
    fig.tight_layout()
    fig.savefig(output_dir / "lesson07.png", dpi=150)
    plt.close(fig)
    test = ~data["train"]
    return {"mst_edges": int(tree.nnz), "vertices": len(centers), "root": root,
            "heldout_time_spearman": float(spearmanr(pseudotime[test], data["time"][test]).statistic),
            "alternate_root_time_spearman": float(spearmanr(reversed_time[test], data["time"][test]).statistic),
            "heldout_mean_projection_squared_error": float(error[test].mean()),
            "evaluation": "training-session centroid MST; observed initial-time labels orient root; not PAGA"}


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))