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

# %% auto 0
__all__ = ['get_gpus', 'compute_normalized_euclidean_distance', 'dist_to_df']

# %% ../../nbs/distance/dist.ipynb 1
import stumpy
from numba import cuda
import numpy as np
import pandas as pd

# %% ../../nbs/distance/dist.ipynb 2
def get_gpus(debug=False):
    gpus = []
    try:
        gpus = [device.id for device in cuda.list_devices()]
        debug and print("GPU device found!")
        debug and print(gpus)
    except cuda.CudaSupportError:
        debug and print("GPU device not found!")
    return gpus

# %% ../../nbs/distance/dist.ipynb 4
def compute_normalized_euclidean_distance(data, window_size):
    gpus = get_gpus()
    if len(gpus) > 0:
        return stumpy.gpu_stump(
            data.astype(np.float64), 
            m=window_size, 
            device_id=gpus,
        )
    return stumpy.stump(
        data.astype(np.float64), 
        m=window_size,
    )

# %% ../../nbs/distance/dist.ipynb 5
def dist_to_df(dist: np.ndarray):
    return pd.DataFrame(dist, columns=["distance", "index", "left_index", "right_index"])
