"""Lesson 6: compare state partitions without mistaking them for truth."""
from pathlib import Path
import argparse
import json
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import networkx as nx
from sklearn.cluster import KMeans, AgglomerativeClustering, DBSCAN
from sklearn.metrics import adjusted_rand_score
from sklearn.neighbors import NearestNeighbors
from common import make_dataset, prepare_embedding
from lesson05 import affinity_graph


def community_labels(graph, resolution):
    groups = nx.community.louvain_communities(graph, weight="weight", resolution=resolution, seed=7)
    labels = np.empty(len(graph), dtype=int)
    for label, group in enumerate(groups):
        labels[list(group)] = label
    return labels, groups


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=6, n_init=10, random_state=7).fit(Z[data["train"]])
    kmeans_labels = km.predict(Z)
    ward_labels = AgglomerativeClustering(n_clusters=6, linkage="ward").fit_predict(Z)
    distances, _ = NearestNeighbors(n_neighbors=12).fit(Z).kneighbors(Z)
    eps = float(np.quantile(distances[:, -1], 0.75))
    dbscan_labels = DBSCAN(eps=eps, min_samples=12).fit_predict(Z)
    graph = nx.from_scipy_sparse_array(affinity_graph(Z), edge_attribute="weight")
    louvain_labels, groups = community_labels(graph, 1.0)
    fine_labels, fine_groups = community_labels(graph, 2.0)
    fig, axes = plt.subplots(2, 2, figsize=(10, 8))
    for ax, labels, title in zip(axes.ravel(), [kmeans_labels, ward_labels, dbscan_labels, louvain_labels],
                                 ["k-means (train fit, all assigned)", "Ward (transductive)", "DBSCAN (transductive)", "Louvain (transductive)"]):
        ax.scatter(Z[:, 0], Z[:, 1], c=labels, s=5, cmap="tab20")
        ax.set_title(title)
        ax.set_xlabel("PC1")
        ax.set_ylabel("PC2")
    fig.tight_layout()
    fig.savefig(output_dir / "lesson06.png", dpi=150)
    plt.close(fig)
    return {"louvain_communities": len(groups), "louvain_resolution_2_communities": len(fine_groups),
            "modularity_resolution_1": float(nx.community.modularity(graph, groups, weight="weight", resolution=1.0)),
            "resolution_partition_ARI": float(adjusted_rand_score(louvain_labels, fine_labels)),
            "dbscan_noise_fraction": float(np.mean(dbscan_labels == -1)),
            "dbscan_eps": eps,
            "heldout_kmeans_mean_squared_distance": float(np.mean(np.min(km.transform(Z[~data["train"]]) ** 2, axis=1))),
            "evaluation": "partition sensitivity; no ground-truth cell or narrative types asserted"}


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