# %%
# A basic funtion to convert MNI coords to harvard oxford labels using spherical masking
# input: 
#   - list of coords of size N*3
# output: 
#   - list of labels of size N


import numpy as np
from nilearn import datasets
from nilearn.input_data import NiftiSpheresMasker



def coords2labels(coords: np.array, sphere_radius=2):
  """
  
  Args:
  -----
    coords: 2d numpy array of size (N,3).

  Returns:
  -----
    a list strings of size N.
  """

  atlas = datasets.fetch_atlas_harvard_oxford('cort-maxprob-thr25-2mm')

  masker = NiftiSpheresMasker(seeds=coords,radius=sphere_radius,allow_overlap=True)

  masker.fit()
  label_indices = masker.transform_single_imgs(atlas.maps).astype(int)

  labels = np.array(atlas.labels)[label_indices.flatten()].tolist()
  return labels


# example use case:
sample_coords = [[0,0,35],
                 [4,5,6],
                 [-12, 20, -20]
                ]

coords2labels(sample_coords)