"""Lecture 4: local geometry on shared data and a separate Swiss-roll diagnostic."""
import argparse
import json
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
from scipy import sparse
from scipy.sparse.csgraph import connected_components
from scipy.spatial.distance import pdist
from sklearn.datasets import make_swiss_roll
from sklearn.manifold import Isomap, trustworthiness
from sklearn.neighbors import kneighbors_graph
from sklearn.preprocessing import StandardScaler
from umap import UMAP
from common import make_dataset


def fit_manifolds(data):
    train = data['train']
    scaler = StandardScaler().fit(data['X'][train])
    X = scaler.transform(data['X'])
    # A larger neighborhood helps connect the two branches; report connectivity explicitly.
    graph = kneighbors_graph(X[train], 30, mode='distance', include_self=False)
    components = connected_components(graph.maximum(graph.T), directed=False)[0]
    if components != 1:
        raise ValueError('Training kNN graph disconnected; inspect k before interpreting Isomap.')
    isomap = Isomap(n_neighbors=30, n_components=2).fit(X[train])
    umap = UMAP(n_neighbors=30, n_components=2, min_dist=0.15,
                random_state=7, n_jobs=1).fit(X[train])
    embeddings = {'Isomap': isomap.transform(X[~train]), 'UMAP': umap.transform(X[~train])}
    metrics = {'training_graph_components': int(components)}
    for name, z in embeddings.items():
        metrics[name.lower() + '_test_trustworthiness'] = float(trustworthiness(X[~train], z, n_neighbors=10))
    return embeddings, metrics


def run(output_dir: Path) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    data = make_dataset()
    embeddings, metrics = fit_manifolds(data)
    # Separate educational diagnostic, not training data for the trajectory experiment.
    roll, roll_parameter = make_swiss_roll(n_samples=800, noise=0.03, random_state=7)
    clean_roll, _ = make_swiss_roll(n_samples=800, noise=0, random_state=7)
    # Known intrinsic coordinates are an oracle diagnostic, never Isomap inputs.
    arc_length = 0.5 * (roll_parameter * np.sqrt(1 + roll_parameter**2)
                        + np.arcsinh(roll_parameter))
    oracle_coordinates = np.column_stack([arc_length, clean_roll[:, 1]])
    oracle_distances = pdist(oracle_coordinates)
    roll_embeddings = {}
    for k in (6, 10):
        graph = kneighbors_graph(roll, k, mode='distance', include_self=False)
        components = int(connected_components(graph.maximum(graph.T), directed=False)[0])
        if components != 1:
            raise ValueError(f'Swiss-roll k={k} graph disconnected; no global unfolding claimed.')
        roll_z = Isomap(n_neighbors=k, n_components=2).fit_transform(roll)
        roll_embeddings[k] = roll_z
        embedded_distances = pdist(roll_z)
        scale = float(oracle_distances @ embedded_distances /
                      (embedded_distances @ embedded_distances))
        stress = np.sqrt(np.sum((scale * embedded_distances - oracle_distances)**2) /
                         np.sum(oracle_distances**2))
        metrics[f'swiss_roll_k{k}_graph_components'] = components
        edges = sparse.triu(graph.maximum(graph.T), k=1).tocoo()
        intrinsic_edge_lengths = np.linalg.norm(
            oracle_coordinates[edges.row] - oracle_coordinates[edges.col], axis=1)
        # Synthetic-only flag: an edge crossing rolled layers is much shorter in ambient space.
        metrics[f'swiss_roll_k{k}_oracle_shortcut_edges'] = int(
            np.sum(intrinsic_edge_lengths > 2 * edges.data))
        metrics[f'swiss_roll_k{k}_trustworthiness'] = float(trustworthiness(roll, roll_z, n_neighbors=10))
        metrics[f'swiss_roll_k{k}_oracle_distance_stress'] = float(stress)
    fig = plt.figure(figsize=(15, 8), constrained_layout=True)
    for pos, (name, z) in enumerate(embeddings.items(), start=1):
        ax = fig.add_subplot(2, 3, pos)
        ax.scatter(z[:, 0], z[:, 1], c=data['time'][~data['train']], s=9)
        ax.set(title=name + ': held-out sessions', xlabel='coordinate 1', ylabel='coordinate 2')
    ax = fig.add_subplot(2, 3, 3, projection='3d')
    ax.scatter(roll[:, 0], roll[:, 1], roll[:, 2], c=roll_parameter, s=6)
    ax.set_title('Swiss roll: 2D surface in 3D')
    for pos, k, label in ((4, 6, 'Unfolded sheet'), (5, 10, 'Shortcut failure')):
        ax = fig.add_subplot(2, 3, pos)
        roll_z = roll_embeddings[k]
        ax.scatter(roll_z[:, 0], roll_z[:, 1], c=roll_parameter, s=8)
        ax.set(title=f'{label}: Isomap k={k}', xlabel='coordinate 1', ylabel='coordinate 2')
        ax.set_aspect('equal', adjustable='datalim')
    ax = fig.add_subplot(2, 3, 6)
    ax.bar(['k=6', 'k=10'], [metrics[f'swiss_roll_k{k}_oracle_distance_stress'] for k in (6, 10)],
           color=['#16817a', '#bb5544'])
    ax.set(title='Oracle global-distance diagnostic', ylabel='Scale-adjusted relative stress (lower is better)')
    fig.savefig(output_dir / 'lesson04.png', dpi=160)
    plt.close(fig)
    return metrics


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