# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:sourmash-sbt2knn]
#     language: python
#     name: conda-env-sourmash-sbt2knn-py
# ---

# cd /home/olga/pureScratch/olgabot-maca/facs/

# ls -lha

# +
# # !ls -lha .sbt.tabula-muris-k21/ | head

# +
import glob
import itertools
from itertools import groupby
import math

import numpy as np
import scipy.sparse
from umap.umap_ import smooth_knn_dist, compute_membership_strengths

from umap.spectral import spectral_layout

# %load_ext autoreload
# %autoreload 2

import sourmash

from sourmash import signature as sig
from sourmash.sbt import Leaf
from sourmash.sbtmh import SigLeaf, create_sbt_index, load_sbt_index
from sourmash import sourmash_args
from sourmash.logging import notify
from sourmash import knn
from sourmash import umap
# -

# ### Settings for this session

# +

# --- Which signatures to use? --- #
ksize = 21
moltype = 'DNA'

# --- How to build sequence bloom tree? --- #
bf_size = 1e5
n_children = 2
scaled = False

# --- How to build K-nearest neighbor graph on similarities? --- #
n_neighbors = 5
ignore_abundance = True
downsample = False

# --- UMAP settings ---
n_epochs = 0
# -

# %time tree = load_sbt_index("tabula-muris-k21.sbt.json")
tree

len(tree.nodes)

set(type(node) for node in tree.nodes.values())

sum(1 for _ in knn.get_leaves(tree))

sum(1 for node in tree.nodes.values() if isinstance(node, (Leaf, SigLeaf)))

# +
# sum(1 for node in tree.nodes.values() if hasattr(node, 'data.minhash'))
# -

# %%time
adjacencies = tree.nearest_neighbor_adjacencies(n_neighbors=n_neighbors, 
                                                ignore_abundance=ignore_abundance, 
                                                downsample=downsample)
print(len(adjacencies))

len(adjacencies)

sum(1 for node in tree.nodes.values() if isinstance(node, (Leaf, SigLeaf)))

# +
# # %%time

# adjacencies = knn.nearest_neighbors(tree, n_neighbors=n_neighbors,
#                                              ignore_abundance=ignore_abundance,
#                                              downsample=downsample)
# print(len(adjacencies))
# -

len(adjacencies)

adjacencies[:10]

u_leaves = set(u for (u, v, similarity) in adjacencies)

'P9-MAA001884-3_38_F-1-1_S151' in u_leaves

leaves = list(knn.get_leaves(tree))
len(leaves)

sum(1 for _ in knn.get_leaves(tree))

len(tree.nodes)

tree.nodes[10]

set(type(node) for node in tree.nodes.values())

sum(1 for node in tree.nodes.values() if hasattr(node, 'minhash'))

leaf_to_position = dict(
        (node.data.name(), position) for position, node in tree.nodes.items()
        if isinstance(node, Leaf))
print(len(leaf_to_position))
leaf_to_index = dict((name, i) for i, name in enumerate(leaf_to_position.keys()))
len(leaf_to_index)

# +
knn_indices = []
knn_dists = []

for u, items in groupby(adjacencies, key=lambda x: x[0]):
    knn_indices_line = []
    knn_dists_line = []
    for u, v, similarity in items:
        index = leaf_to_index[v]
        if not isinstance(index, int):
            print(f"index {index} is not an integer")
        knn_indices_line.append(index)

        # Dissimilarity = 1-similarity
        dissimilarity = float(1 - similarity)
        if not isinstance(dissimilarity, float):
            print(f"dissimilarity {dissimilarity} is not a float")
        knn_dists_line.append(dissimilarity)
#     knn_indices_line = np.array(knn_indices_line)
#     knn_dists_line = np.array(knn_dists_line)
    knn_indices.append(knn_indices_line)
    knn_dists.append(knn_dists_line)
knn_indices = np.array(knn_indices, dtype='int')
knn_dists = np.array(knn_dists)
# -

[x for x in knn_indices if len(x) != n_neighbors]

knn_indices.dtype

knn_indices.astype(int)

knn_dists.dtype

knn_dists[0].dtype

knn_indices

knn_dists.astype(float)

knn_dists[:10]

leaf_to_index['P9-MAA001884-3_38_F-1-1_S151']

leaf_to_position['P9-MAA001884-3_38_F-1-1_S151']

knn_indices.dtypes

graph = umap.fuzzy_simplicial_set(knn_indices, knn_dists, n_neighbors)

# +
# knn_indices, knn_dists, leaf_to_index = \
#     sourmash_knn.similarity_adjacency_to_knn(adjacencies, tree)
# print(f"knn_indices.shape: {knn_indices.shape}")
# print(f"knn_dists.shape: {knn_dists.shape}")
# print(f"len(leaf_to_index): {len(leaf_to_index)}")
# -

# %debug

index_to_leaf = dict(zip(leaf_to_index.values(), leaf_to_index.keys()))

graph = sourmash_knn.make_search_graph(knn_indices, knn_dists)

# +
graph, n_epochs = sourmash_knn._clean_graph(graph, n_epochs)


# -

# ## Get connected components and similarity_adjacency_to_knnem



# +
def convert_name_to_tree_position(tree, name):
    for position, node in tree.nodes.items():
        if node.data.name() == name:
            return position

label_to_tree_position = dict((name, convert_name_to_tree_position(name)) for name in leaf_to_index.keys())
len(label_to_tree_position)
label_to_tree_position['P9-MAA001884-3_38_F-1-1_S151']

# +
n_connected_components, labels = scipy.sparse.csgraph.connected_components(graph)

for label in range(n_connected_components):
    graph_index_of_label = np.where(component_labels == label)
    tree_position
    
    first_sig = tree.get(graph)
    for i in graph_index_of_label[1:]:
        
# -

np.where(component_labels == label)



StarcoderdataJupyterScriptsDedupFiltered