# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/distance/cluster.ipynb.

# %% auto 0
__all__ = ['compute_thresholds_from_dist_limits', 'compute_clusters_indices_by_distances_thresholds',
           'compute_clusters_indices_by_quantiles', 'get_cluster_windows']

# %% ../../nbs/distance/cluster.ipynb 20
from typing import Iterable

import numpy as np

def compute_thresholds_from_dist_limits(
    dist_limits: Iterable[float],
):
    # ensure dist_limits is ascending
    dist_limits = sorted(dist_limits)
    assert dist_limits[0] >= 0
    
    n_clusters = len(dist_limits) + 1
    thresholds = []
    for cluster_idx in range(n_clusters):
        threshold_start = 0
        if cluster_idx != 0:
            threshold_start = dist_limits[cluster_idx - 1]

        threshold_end = np.inf
        if cluster_idx != n_clusters - 1:
            threshold_end = dist_limits[cluster_idx]

        thresholds.append((threshold_start, threshold_end))

    return thresholds

# %% ../../nbs/distance/cluster.ipynb 23
from typing import Iterable

import pandas as pd

def compute_clusters_indices_by_distances_thresholds(
    distances: pd.Series,
    dist_limits: Iterable[float],
):
    thresholds = compute_thresholds_from_dist_limits(dist_limits)
    cluster_indices = {}
    for cluster_idx, (threshold_start, threshold_end) in enumerate(thresholds):
        cluster_indices[cluster_idx] = distances[
            (threshold_start <= distances) & (distances <= threshold_end)
        ].index

    return cluster_indices

# %% ../../nbs/distance/cluster.ipynb 26
import pandas as pd
import numpy as np


def compute_clusters_indices_by_quantiles(
    distances: pd.Series,
    cluster_quantiles: list[int],
):
    # ensure that cluster_quantiles is ascending and between 0 and 1
    cluster_quantiles = sorted(cluster_quantiles)
    assert cluster_quantiles[0] >= 0 and cluster_quantiles[-1] <= 1
    cluster_quantiles

    # ensure that distances is a 1D array
    assert distances.ndim == 1

    cluster_thresholds = distances.quantile(cluster_quantiles).to_numpy()
    n_clusters = len(cluster_quantiles) + 1
    cluster_indices = {}
    for cluster_idx in range(n_clusters):
        threshold_start = 0
        if cluster_idx != 0:
            threshold_start = cluster_thresholds[cluster_idx - 1]

        threshold_end = np.inf
        if cluster_idx != n_clusters - 1:
            threshold_end = cluster_thresholds[cluster_idx]

        cluster_indices[cluster_idx] = distances[
            (threshold_start <= distances) & (distances <= threshold_end)
        ].index

    return cluster_indices

# %% ../../nbs/distance/cluster.ipynb 41
def get_cluster_windows(data: np.ndarray, cluster_indices: np.ndarray, seq_len: int):
    # ensure that data is a 1D array
    assert data.ndim == 1

    # ensure that cluster_indices is a 1D array of integers
    assert cluster_indices.ndim == 1
    assert cluster_indices.dtype == np.int64

    n_cluster_windows = len(cluster_indices) - seq_len + 1

    sample_idx_within_cluster = 0
    while sample_idx_within_cluster < n_cluster_windows:
        sample_idx_in_data = cluster_indices[sample_idx_within_cluster]
        yield data[sample_idx_in_data : sample_idx_in_data + seq_len]
        sample_idx_within_cluster += 1
