Compute modules

class BaseCompute():
    '''Base class for analysis compute jobs.

    Subclasses implement specific analysis pipelines (spike sorting, cell
    segmentation, pose estimation, etc.) by overriding ``_compute()``.  The
    base class handles job claiming, file retrieval, status tracking, and
    safe-exit cleanup.

    Jobs are dispatched via ``place_tasks_in_queue()`` and consumed by worker
    nodes running inside Apptainer containers.  The worker calls
    ``handle_compute()`` which instantiates the appropriate subclass and calls
    ``compute()``.

    Class Attributes
    ----------------
    name : str or None
        Short identifier used to select this class from the command line.
    container : str
        Apptainer container image name.
    cuda : bool
        Whether the container requires a GPU.
    ec2 : dict
        EC2 instance type presets (``'small'``, ``'large'``) for cloud runs.

    Parameters
    ----------
    job_id : int
        Primary key of the ``ComputeTask`` row to process.
    project : str or None, optional
        Project name used to load the correct schema module.
    allow_s3 : bool or None, optional
        Whether to download missing files from S3.  Defaults to
        ``prefs['allow_s3_download']``.
    keep_intermediate : bool, optional
        Keep intermediate files after the job finishes.  Default is False.

    See Also
    --------
    labdata.compute.ephys.SpksCompute : Spike sorting with Kilosort.
    labdata.compute.suite2p.Suite2pCompute : Cell segmentation with Suite2p.
    labdata.compute.caiman.CaimanCompute : Cell segmentation with CaImAn.
    labdata.compute.pose.DeeplabcutCompute : Pose estimation with DeepLabCut.
    '''
    name = None
    container = 'labdata-base'
    cuda = False
    ec2 = dict(small=dict(instance_type='g4dn.2xlarge'),  # 8 cpus, 32 GB mem, 200 GB nvme, 1 gpu
               large=dict(instance_type='g6.4xlarge',     # 16 cpus, 64 GB mem, 600 GB nvme, 1 gpu  
                          availability_zone='us-west-2b'))

    def __init__(self, job_id, project=None, allow_s3=None, keep_intermediate=False):
        '''
        Executes a computation on a dataset, that can be remote or local.
        Uses a singularity/apptainer image if possible.
        '''
        self.file_filters = ['.'] # selects all files...
        self.parameters = dict()
        self.schema = load_project_schema(project)
        self.keep_intermediate = keep_intermediate
        self.job_id = job_id
        if not self.job_id is None:
            self._check_if_taken()

        self.paths = None
        self.local_path = Path(prefs['local_paths'][0])
        self.scratch_path = Path(prefs['scratch_path'])
        self.assigned_files = None
        self.dataset_key = None
        self.is_container = False
        if allow_s3 is None:
            allow_s3 = prefs['allow_s3_download']
        self.allow_s3 = allow_s3
        if 'LABDATA_CONTAINER' in os.environ.keys():
            # then it is running inside a container
            self.is_container = True
        #self.is_ec2 = False # then files should be taken from s3

    def _init_job(self): # to run in the init function
        if not self.job_id is None:
            with self.schema.dj.conn().transaction:
                self.jobquery = (self.schema.ComputeTask() & dict(job_id = self.job_id))
                job_status = self.jobquery.fetch(as_dict = True)
                if len(job_status):
                    if not job_status[0]['task_waiting']:
                        print(f'Checking job_status - task was not waiting: {job_status}', flush = True)
                        if 'SLURM_RESTART_COUNT' in os.environ.keys():
                            # then the job is running on slurm.. its a putative restart, try to run it..
                            self.set_job_status(job_status = 'WORKING',
                                                job_waiting = 0)
                        else:
                            print(f"Compute task [{self.job_id}] is already taken.")
                            print(job_status, flush = True)
                            return # exit.
                    else:
                        self.set_job_status(job_status = 'WORKING',
                                            job_waiting = 0)

                        def cleanup_function(job_id=self.job_id):
                            '''Reset the job to WAITING if the process exits abnormally.'''
                            print('Running the cleanup function.', flush=True)
                            status = (self.schema.ComputeTask() & dict(job_id = job_id)).fetch(as_dict = True)[0]
                            if status['task_status'] in ['WORKING']:
                                self.schema.ComputeTask.update1(dict(job_id = job_id,
                                                                     task_status = 'CANCELLED',
                                                                     task_waiting = 1,
                                                                     task_endtime = datetime.now()))

                        self.cleanup_function = cleanup_function       
                        self.register_safe_exit()
                        par = json.loads(job_status[0]['task_parameters'])
                        for k in par.keys():
                            self.parameters[k] = par[k]
                        self.assigned_files = pd.DataFrame((self.schema.ComputeTask.AssignedFiles() & dict(job_id = self.job_id)).fetch())
                        self.dataset_key = dict(subject_name = job_status[0]['subject_name'],
                                                session_name = job_status[0]['session_name'],
                                                dataset_name = job_status[0]['dataset_name'])
                        if '--multisession' in job_status[0]['task_cmd']:
                            self.dataset_key = (self.schema.Dataset() &
                                                (self.schema.Dataset.DataFiles & self.assigned_files)).proj().fetch(as_dict = True)
                        # delete the job if has --delete-on-complete
                        job_status = self.jobquery.fetch(as_dict = True)
                        if '--keep-files' in job_status[0]['task_cmd']:
                            self.keep_intermediate = True
                else:
                    # that should just be a problem to fix
                    raise ValueError(f'job_id {self.job_id} does not exist.')

    def register_safe_exit(self):
        '''Register ``cleanup_function`` with ``safe_exit`` so it runs on abnormal termination.'''
        import safe_exit
        safe_exit.register(self.cleanup_function)

    def unregister_safe_exit(self):
        '''Unregister the cleanup function when the job finishes normally.'''
        import safe_exit
        safe_exit.unregister(self.cleanup_function)

    def get_files(self, dset, allowed_extensions=[]):
        '''
        Gets the paths and downloads from S3 if needed.
        '''
        if type(dset) is list:
            # then it is a list of dicts, convert to DataFrame
            dset = pd.DataFrame(dset)
        files = dset.file_path.values
        print('---')
        print(files)
        print('---')
        storage = dset.storage.values
        localpath = Path(prefs['local_paths'][0])
        self.files_existed = True
        localfiles = [find_local_filepath(f,
                                          allowed_extensions = allowed_extensions) for f in files]
        localfiles = np.unique(list(filter(lambda x: not x is None,localfiles)))
        if not len(localfiles) >= len(self.dataset_key): # tries to have at least one file per dataset (check this assumption, it is here because of the "allowed extensions")
            # then you can try downloading the files
            if self.allow_s3: # get the files from s3
                #TODO: then it should download using "File"
                from ..s3 import copy_from_s3
                for s in np.unique(storage):
                    # so it can work with multiple storages
                    srcfiles = [f for f in files[storage == s]]
                    dstfiles = [localpath/f for f in srcfiles]
                print(f'Downloading {len(srcfiles)} files from S3 [{s}].')
                copy_from_s3(srcfiles,dstfiles,storage_name = s)
                localfiles = np.unique([find_local_filepath(
                    f,
                    allowed_extensions = allowed_extensions) for f in files])
                if len(localfiles):
                    self.files_existed = False # delete the files in the end if they were not local.
            else:
                print(files, localpath)
                raise(ValueError('Files not found locally, set allow_s3 in the preferences to download.'))
        return localfiles

    def place_tasks_in_queue(self, datasets, task_cmd=None, force_submit=False, multisession=False):
        '''Submit compute tasks for each dataset to the ``ComputeTask`` queue.

        Delegates to ``_place_tasks_in_queue()``.  Override in subclasses to
        customise submission logic (e.g. ``SpksCompute`` handles multi-probe).

        Parameters
        ----------
        datasets : list of dict or None
            Dataset keys to process.  ``None`` submits a single job with no
            dataset association.
        task_cmd : str or None, optional
            Full command string stored in ``ComputeTask.task_cmd``.
        force_submit : bool, optional
            Re-submit even if an identical job already exists.  Default False.
        multisession : bool, optional
            Combine all datasets into a single multi-session job.  Default False.

        Returns
        -------
        list of int
            ``ComputeTask.job_id`` values for the submitted jobs.
        '''
        return self._place_tasks_in_queue(datasets, task_cmd=task_cmd,
                                          force_submit=force_submit,
                                          multisession=multisession)

    def _place_tasks_in_queue(self,datasets, task_cmd = None, force_submit = False, multisession = False, parameters = None):
        ''' This will put the tasks in the queue for each dataset.
        If the task and parameters are the same it will return the job_id instead.

        '''
        if parameters is None:
            parameters = self.parameters # so we can pass multiple parameters (e.g. in the multiprobe case)
        job_ids = []
        if datasets is None:
            datasets = [None]
        if multisession:
            print('Combining data from multiple sessions/datasets.')
            datasets = [datasets]
        for dataset in datasets:
            if not dataset is None: # then there are no associated files.
                files = pd.DataFrame((self.schema.Dataset.DataFiles() & dataset).fetch())
                idx = []
                for f in self.file_filters:
                    idx += list(filter(lambda x: not x is None,[i if f in s else None for i,s in enumerate(
                        files.file_path.values)]))
                if len(idx) == 0:
                    raise ValueError(f'Could not find valid Dataset.DataFiles for {dataset}')
                files = files.iloc[idx]
                if type(dataset) is dict:
                    key = dict(dataset,task_name = self.name)
                else:
                    key = dict(dataset[0],task_name = self.name)
                exists = self.schema.ComputeTask() & key
                if len(exists):
                    d = pd.DataFrame(exists.fetch())
                    # if any(d.task_status.values=='WORKING'):
                    #     print('A task was running for this dataset, stop or delete it first.')
                    #     print(key)
                    #     continue

                    idx = np.where(np.array(d.task_parameters.values) == json.dumps(parameters))[0]
                    if len(idx):
                        job_id = d.iloc[idx].job_id.values[0]
                        print(f'There is a task to analyse dataset {key} with the same parameters. [{job_id}]')
                        if force_submit:
                            print('Deleting the previous job because force_submit is set.')
                            with self.schema.dj.conn().transaction:
                                self.schema.dj.config['safemode'] = False
                                # delete part table because the reference priviledge is not sufficient
                                (self.schema.ComputeTask.AssignedFiles & f'job_id = {job_id}').delete(force = True) 
                                (self.schema.ComputeTask & f'job_id = {job_id}').delete() # deleting a previous job because of force_submit
                                self.schema.dj.config['safemode'] = True
                        else:
                            continue
            else:
                key = dict(task_name = self.name)
                files = None

            with self.schema.dj.conn().transaction:
                job_id = self.schema.ComputeTask().fetch('job_id')
                if len(job_id):
                    job_id = np.max(job_id) + 1 
                else:
                    job_id = 1
                if not task_cmd is None:
                    if len(task_cmd) >1999:
                        task_cmd = task_cmd[:1999]
                self.schema.ComputeTask().insert1(dict(key,
                                                       job_id = job_id,
                                                       task_waiting = 1,
                                                       task_status = "WAITING",
                                                       task_target = None,
                                                       task_host = None,
                                                       task_cmd = task_cmd,
                                                       task_parameters = json.dumps(parameters),
                                                       task_log = None))
                if not files is None:
                    self.schema.ComputeTask.AssignedFiles().insert([dict(job_id = job_id,
                                                                         storage = f.storage,
                                                                         file_path = f.file_path)
                                                                for i,f in files.iterrows()])
                job_ids.append(job_id)
        return job_ids

    def find_datasets(self,subject_name = None, session_name = None, dataset_name = None):
        '''
        Find datasets to analyze, this function will search in the proper tables if datasets are available.
        Has to be implemented per Compute class since it varies.
        '''
        raise NotImplementedError('The find_datasets method has to be implemented.')

    def secondary_parse(self, secondary_arguments, parameter_number=None):
        '''Parse secondary command-line arguments for this compute class.

        Delegates to ``_secondary_parse()`` which subclasses override to add
        algorithm-specific CLI options (e.g. motion correction flags, parameter
        set selection, batch size).

        Parameters
        ----------
        secondary_arguments : list of str
            Arguments that appear after ``--`` on the command line.
        parameter_number : int or None, optional
            Force selection of a specific parameter set number.
        '''
        self._secondary_parse(secondary_arguments, parameter_number)

    def _secondary_parse(self,secondary_arguments):
        return

    def _check_if_taken(self):
        if not self.job_id is None:
            self.jobquery = (self.schema.ComputeTask() & dict(job_id = self.job_id))
            job_status = self.jobquery.fetch(as_dict = True)
            if len(job_status):
                if job_status[0]['task_waiting']:
                    return
                else:
                    print(job_status, flush = True)
                    raise ValueError(f'job_id {self.job_id} is already taken.')
                    return # exit.
            else:
                raise ValueError(f'job_id {self.job_id} does not exist.')
            # get the paths?
            #self.src_paths = pd.DataFrame((ComputeTask.AssignedFiles() &
            #                               dict(job_id = self.job_id)).fetch())
            #if not len(self.src_paths):
            #    self.set_job_status(job_status = 'FAILED',
            #                        job_log = f'Could not find files for {self.job_id} in ComputeTask.AssignedFiles.')
            #    raise ValueError(f'Could not find files for {self.job_id} in ComputeTask.AssignedFiles.')
        else:
            raise ValueError(f'Compute: job_id not specified.')

    def compute(self):
        '''This calls the compute function. 
If "use_s3" is true it will download the files from s3 when needed.'''
        try:
            if not self.job_id is None:
                dd = dict(job_id = self.job_id,
                          task_starttime = datetime.now())
                self.schema.ComputeTask().update1(dd)
            self._compute() # can use the src_paths
        except Exception as err:
            # log the error
            print(f'There was an error processing job {self.job_id}.')
            err =  str(traceback.format_exc()) + "ERROR" +str(err)
            print(err)

            if len(err) > 1999: # then get only the last part of the error.
                err = err[-1900:]
            if type(err) is str:
                # avoid encoding errors
                err = err.encode('utf-8')
            self.set_job_status(job_status = 'FAILED',
                                job_log = f'{err}')
            return

        # get the job from the DB if the status is not failed, mark completed (remember to clean the log)
        self.jobquery = (self.schema.ComputeTask() & dict(job_id = self.job_id))
        job_status = self.jobquery.fetch(as_dict = True)

        if not job_status[0]['task_status'] in ['FAILED']:
            self._post_compute() # so the rules can insert tables and all.
            # delete the job if has --delete-on-complete
            if '--delete-on-complete' in job_status[0]['task_cmd']:
                self.jobquery.delete(safemode = False)
                self.job_id = None
                self.unregister_safe_exit()
            else:
                self.set_job_status(job_status = 'COMPLETED')

        if not self.job_id is None: # set the complete time
            dd = dict(job_id = self.job_id,
                      task_endtime = datetime.now())
            self.schema.ComputeTask().update1(dd)

    def set_job_status(self, job_status=None, job_log=None, job_waiting=0):
        '''Update the status of this compute job in ``ComputeTask``.

        Parameters
        ----------
        job_status : str or None, optional
            New status string (e.g. ``'WORKING'``, ``'COMPLETED'``,
            ``'FAILED'``).  ``None`` leaves the current status unchanged.
        job_log : str or None, optional
            Log message stored in ``ComputeTask.task_log``.
        job_waiting : int, optional
            Whether the job should be picked up again (``1``) or is taken
            (``0``).  Default is ``0``.
        '''
        from ..schema import ComputeTask
        if not self.job_id is None:
            dd = dict(job_id = self.job_id,
                      task_waiting = job_waiting,
                      task_host = prefs['hostname']) # so we know where it failed.)
            if not job_status is None:
                dd['task_status'] = job_status
            if not job_log is None:
                dd['task_log'] = job_log  
                if type(dd['task_log']) is str:
                    #prevent error due to unsupported characters
                    dd['task_log'] = dd['task_log'].encode('utf-8')
            self.schema.ComputeTask.update1(dd)
            if not job_status is None:
                if not 'WORK' in job_status: # display the message
                    print(f'Check job_id {self.job_id} : {job_status}')

    def _post_compute(self):
        '''
        Inserts the data to the database
        '''
        return

    def _compute(self):
        '''
        Runs the compute job on a scratch folder.
        '''
        return

class PopulateCompute(BaseCompute):
    '''Compute class that calls DataJoint ``populate()`` on an arbitrary table.

    Useful for running DataJoint computed tables (``dj.Computed`` /
    ``dj.Imported``) inside a container or on a remote worker without writing
    a custom compute class.  The target table, import module, number of
    processes, and optional restriction are configured via secondary CLI
    arguments.

    Secondary arguments (passed after ``--``)::

        populate <table> [-i IMPORTS] [-p PROCESSES] [-r RESTRICTIONS] [-s]

    Examples
    --------
    Queue a populate job for ``UnitMetrics`` on all sessions::

        parse_analysis('populate', secondary_args=['populate', 'UnitMetrics'])
    '''
    container = 'labdata-base'
    cuda = False
    name = 'populate'
    url = 'http://github.com/jcouto/labdata'
    def __init__(self, job_id, project=None, allow_s3=None, **kwargs):
        super(PopulateCompute,self).__init__(job_id, project = project, allow_s3 = allow_s3)
        self.file_filters = None
        # default parameters
        self.parameters = dict(imports = 'labdata.schema',
                               table = 'UnitMetrics',
                               processes = 10)

        self._init_job() # gets the parameters

    def _secondary_parse(self,arguments,parameter_number):
        '''
        Handles parsing the command line interface
        '''
        import argparse
        parser = argparse.ArgumentParser(
            description = 'Populate arbitrary tables',
            usage = 'populate -- -t <TABLE> -i <IMPORTS>')

        parser.add_argument('table',action='store',type = str,
                            help = 'Table to populate')
        parser.add_argument('-s','--stop-on-errors',action='store_true',default= False,
                            help = 'Stop on errors (negates suppress_errors)')
        parser.add_argument('-r','--restrictions',action='store',default= '',
                            help = 'Restrictions to the populate table (dict(X = "x")) or completed_today (to run for sessions that were completed less than 24h ago)')
        parser.add_argument('-i','--imports',action='store',default= 'labdata.schema',type = str,
                            help = 'import modules to load the table')
        parser.add_argument('-p','--processes',
                            action='store', default=1, type = int,
                            help = "Required imports.")
        parser.add_argument

        args = parser.parse_args(arguments[1:])
        self.parameters = dict(table = args.table,
                               imports = args.imports,
                               processes = args.processes,
                               suppress_errors = not args.stop_on_errors,
                               restrictions = args.restrictions)
        # try the import and check if the default container exists.
        if not self.parameters["imports"] in ['none','']:
            to_import = self.parameters["table"]
            if '.' in to_import: # to import plugins
                to_import = to_import.split('.')[0]
            exec(f'from {self.parameters["imports"]} import {to_import}')
        table = eval(f'{self.parameters["table"]}')
        if hasattr(table,'default_container'):
            self.container = table.default_container
        print(self.parameters)

    def find_datasets(self):
        '''Return ``None`` — ``PopulateCompute`` does not filter by dataset.'''
        return

    def _compute(self):
        # import
        if not self.parameters["imports"] in ['none','']:
            to_import = self.parameters["table"]
            if '.' in to_import: # to import plugins
                to_import = to_import.split('.')[0]
            exec(f'from {self.parameters["imports"]} import {to_import}')
        processes = 1
        # check nprocesses
        if 'processes' in self.parameters.keys():
            processes = int(self.parameters['processes'])
        # submit populate
        suppress_errors = self.parameters['suppress_errors']
        if self.parameters['restrictions'] == '':
            exec(f'{self.parameters["table"]}.populate(suppress_errors={suppress_errors}, processes = {processes}, display_progress = True)')
        else:
            if self.parameters['restrictions'] == 'completed_today': # this will look for uploads that happened less than 24h ago
                restrictions = (self.schema.Session() & (self.schema.UploadJob() & 'job_status = "COMPLETED"') & 'session_datetime > DATE_SUB(CURDATE(), INTERVAL 24 HOUR)').proj().fetch(as_dict = True)
            else:
                restrictions = eval(self.parameters['restrictions'])
            exec(f'{self.parameters["table"]}.populate({restrictions}, suppress_errors={suppress_errors}, processes = {processes}, display_progress = True)')

class SpksCompute(BaseCompute):
    '''Spike sorting compute class using the ``spks`` / Kilosort pipeline.

    Runs Kilosort (via the ``spks`` wrapper) on each probe in turn.  For
    each probe the steps are:

    1. Locate ``.ap.cbin`` (or ``.ap.bin``) files assigned to the job.
    2. Copy the file to the scratch directory.
    3. Run spike sorting (Kilosort 4 by default).
    4. Delete the raw file from scratch.
    5. Upload results (features, waveforms) to S3 and insert into
       ``SpikeSorting`` and ``SpikeSorting.Unit``.

    Parameters used for sorting are stored in ``SpikeSortingParams`` and
    matched across jobs so that identical parameter sets share a row.

    Requires the ``labdata-spks`` Apptainer container and a CUDA GPU.
    '''
    container = 'labdata-spks'
    cuda = True
    ec2 = dict(small=dict(instance_type='g4dn.2xlarge'),
               large=dict(instance_type='g6.4xlarge',
                          availability_zone='us-west-2b'))
    name = 'spks'
    url = 'http://github.com/spkware/spks'
    def __init__(self, job_id, project=None, allow_s3=None, **kwargs):
        super(SpksCompute,self).__init__(job_id, project = project, allow_s3 = allow_s3)
        self.file_filters = ['.ap.']
        # default parameters
        self.parameters = dict(algorithm_name = 'spks_kilosort4.0',
                               motion_correction = 1,
                               low_pass = 300.,
                               high_pass = 13000.)
        self.parameter_keys = ['motion_correction','low_pass','high_pass','thresholds','remove_cross_duplicates',
                               'waveforms_from_input','dredge']
        # the parameters that go on the SpikeSortingParams
        self.use_hdf5 = True  # flag to use h5py or zarr format for the waveforms.
        self.parameter_set_num = None # identifier in SpikeSortingParams
        self._init_job()
        if type(self.dataset_key) is dict:
            self.dataset_key = [self.dataset_key] # make it a list
        if not self.job_id is None:
            self.add_parameter_key()

    def _get_parameter_number(self):
        parameter_set_num = None
        # check if in spike sorting
        parameters = pd.DataFrame(self.schema.SpikeSortingParams().fetch())
        filtered_par = {k:self.parameters[k] for k in self.parameter_keys if k in self.parameters.keys()}  
        for i,r in parameters.iterrows():
            # go through every algo
            if filtered_par == json.loads(r.parameters_dict) and self.parameters['algorithm_name'] == r['algorithm_name']:
                parameter_set_num = r.parameter_set_num
        if parameter_set_num is None:
            if len(parameters) == 0:
                parameter_set_num = 1
            else:
                parameter_set_num = np.max(parameters.parameter_set_num.values)+1
            print(f'  --> Using parameter set num {parameter_set_num}')
        return parameter_set_num,parameters

    def add_parameter_key(self):
        '''Insert a ``SpikeSortingParams`` row for the current parameters if it does not exist.

        Resolves or creates a ``parameter_set_num``, inserts into
        ``SpikeSortingParams``, and raises ``ValueError`` if the dataset has
        already been sorted with these parameters.
        '''
        parameter_set_num, parameters = self._get_parameter_number()
        if not parameter_set_num in parameters.parameter_set_num.values:
            filtered_par = {k:self.parameters[k] for k in self.parameter_keys if k in self.parameters.keys()}
            self.schema.SpikeSortingParams().insert1(dict(parameter_set_num = parameter_set_num,
                                               algorithm_name = self.parameters['algorithm_name'],
                                               parameters_dict = json.dumps(filtered_par),
                                               code_link = self.url),
                                          skip_duplicates=True)
        self.parameter_set_num = parameter_set_num
        recordings = self.schema.EphysRecording.ProbeSetting() & self.dataset_key
        sortings = self.schema.SpikeSorting() & [dict(d, parameter_set_num = self.parameter_set_num) for d in self.dataset_key]
        if len(recordings) == len(sortings):
            self.set_job_status(
                job_status = 'FAILED',
                job_waiting = 0,
                job_log = f'{self.dataset_key[0]} was already sorted with parameters {self.parameter_set_num}.')    
            raise(ValueError(f'{self.dataset_key[0]} was already sorted with parameters {self.parameter_set_num}.'))

    def _secondary_parse(self,arguments,parameter_number = None):
        '''
        Handles parsing the command line interface
        '''
        if not parameter_number is None:
            self.parameters = json.loads((self.schema.SpikeSortingParams() & f'parameter_set_num = {parameter_number}').fetch('parameters_dict'))
            self.parameters['algorithm_name'] = ((self.schema.SpikeSortingParams() & f'parameter_set_num = {parameter_number}')).fetch1('algorith_name')
            if not len(self.parameters):
                raise(f'Could not find parameter {parameter_number} in SpikeSortingParams.')
            self.parameters = self.parameters[0]
        else:
            import argparse
            parser = argparse.ArgumentParser(
                description = 'Spike sorting using kilosort through the spks package.',
                usage = 'spks -a <SUBJECT> -s <SESSION> -- <PARAMETERS>')

            parser.add_argument('-p','--probe',
                                action='store', default=None, type = int,
                                help = "THIS DOES NOTHING NOW. WILL BE FOR OPENING PHY")
            parser.add_argument('-m','--method',action='store',default = 'ks4.0',type = str,
                                help = 'Method for spike sorting [Kilosort] ks2.5, ks3.0, ks4.0 or [MountainSort] ms5 (default ks4.0)')
            parser.add_argument('-l','--low-pass',
                                action='store', default=self.parameters['low_pass'], type = float,
                                help = "Lowpass filter (default 300.Hz)")
            parser.add_argument('-i','--high-pass',
                                action='store', default=self.parameters['high_pass'], type = float,
                                help = "Highpass filter (default 13000.Hz)")
            parser.add_argument('-t','--thresholds',
                                action='store', default=None, type = float, nargs = 2,
                                help = "Thresholds for spike detection default depends on method.")
            parser.add_argument('-n','--motion-correction',
                                action='store', default = 1, type = int,
                                help = "Motion correction (0 to disable, 1 for rigid, 2+ for blockwise)")
            parser.add_argument('-c','--remove_cross-unit-duplicates',
                                action='store_true', default = False,
                                help = "Skip removing duplicates across units.")
            parser.add_argument('--waveforms-from-sorter',
                                action='store_true', default = False,
                                help = "Extract the waveforms from the file processed by the sorter.")
            parser.add_argument('--interval',
                                action='store', default = None, nargs=2, type = float,
                                help = "Interval to sort in seconds: start end")
            parser.add_argument('--dredge',
                                action='store_true', default = False, 
                                help = "Motion correction using dredge")

            args = parser.parse_args(arguments[1:])
            if 'ks2.5' in  args.method: # defaults for ks2.5
                self.parameters = dict(algorithm_name = 'spks_kilosort2.5',
                                    motion_correction = args.motion_correction>0,
                                    low_pass = args.low_pass,
                                    high_pass = args.high_pass,
                                    thresholds = [9.,3.],
                                    remove_cross_duplicates = args.remove_cross_unit_duplicates)
            elif 'ks3.0' in  args.method: # defaults for ks3.0
                self.parameters = dict(algorithm_name = 'spks_kilosort3.0',
                                    motion_correction = args.no_motion_correction>0,
                                    low_pass = args.low_pass,
                                    high_pass = args.high_pass,
                                    thresholds = [9.,9.],
                                    remove_cross_duplicates = args.remove_cross_unit_duplicates)
            elif 'ks4.0' in  args.method: # defaults for ks4.0
                self.parameters = dict(algorithm_name = 'spks_kilosort4.0',
                                    motion_correction = args.motion_correction,
                                    low_pass = args.low_pass,
                                    high_pass = args.high_pass,
                                    thresholds = [9.,8.],
                                    remove_cross_duplicates = args.remove_cross_unit_duplicates)
            else:
                raise(NotImplemented(f'{args.method} not implemented.'))
            if args.waveforms_from_sorter:
                # default is to extract from the original input file
                self.parameters['waveforms_from_input'] = False 
            if args.dredge:
                # default is to extract from the original input file
                self.parameters['dredge'] = True 
            print(args)
            if not args.thresholds is None:
                self.parameters['thresholds'] = args.thresholds
            if self.parameters['motion_correction'] < 2:
                self.parameters['motion_correction'] = self.parameters['motion_correction']
        self.probe = args.probe
        if not self.probe is None:
            self.parameters['probe'] = int(self.probe) # submit a single probe

    def find_datasets(self, subject_name = None, session_name = None):
        '''
        Searches for subjects and sessions in EphysRecording
        '''
        if subject_name is None and session_name is None:
            print("\n\nPlease specify a 'subject_name' and a 'session_name' to perform spike-sorting.\n\n")

        parameter_set_num, parameters = self._get_parameter_number()
        keys = []
        if not subject_name is None:
            if len(subject_name) > 1:
                raise ValueError(f'Please submit one subject at a time {subject_name}.')
            if not subject_name[0] == '':
                subject_name = subject_name[0]
        if not session_name is None:
            for s in session_name:
                if not s == '':
                    keys.append(dict(subject_name = subject_name,
                                     session_name = s))
        else:
            # find all sessions that can be spike sorted
            sessions = np.unique(((
                (self.schema.EphysRecording() & f'subject_name = "{subject_name}"') -
                (self.schema.SpikeSorting() & f'parameter_set_num = {parameter_set_num}'))).fetch('session_name'))
            for ses in sessions:
                keys.append(dict(subject_name = subject_name,
                                 session_name = ses))
        datasets = []
        for k in keys:
            datasets += (self.schema.EphysRecording()& k).proj('subject_name','session_name','dataset_name').fetch(as_dict = True)

        if not parameter_set_num is None:
            datasets = (self.schema.EphysRecording & ((self.schema.EphysRecording.ProbeSetting() & datasets) -
                        (self.schema.SpikeSorting() & f'parameter_set_num = {parameter_set_num}'))).proj(
                            'subject_name',
                            'session_name',
                            'dataset_name').fetch(as_dict = True)
        print(keys,parameter_set_num)
        return datasets

    def place_tasks_in_queue(self, datasets, task_cmd=None, force_submit=False, multisession=False):
        '''Submit one ``ComputeTask`` per probe per dataset.

        Overrides ``BaseCompute.place_tasks_in_queue`` to split multi-probe
        datasets into separate jobs (one job per ``probe_num``) so each probe
        can be sorted independently and in parallel.

        Parameters
        ----------
        datasets : list of dict
            Dataset keys to sort.
        task_cmd : str or None, optional
            Full shell command stored in ``ComputeTask.task_cmd``.
        force_submit : bool, optional
            Submit even if an identical job exists.  Default is False.
        multisession : bool, optional
            Combine datasets into a single multi-session job.  Default is False.

        Returns
        -------
        list of int
            ``ComputeTask.job_id`` values for all submitted jobs.
        '''
        # this is a special place tasks so we can submit a compute task per probe (when there are multiple probes per dataset)
        if not 'probe' in self.parameters.keys():
            probes = np.unique((self.schema.EphysRecording.ProbeSetting & datasets).fetch('probe_num'))
            jobids = []
            for probe in probes:
                p = dict(self.parameters,probe = int(probe))
                jobids.extend(self._place_tasks_in_queue(datasets,
                                           task_cmd = task_cmd, 
                                           force_submit = force_submit, 
                                           multisession = multisession, 
                                           parameters = p))
            return jobids
        else:
            return self._place_tasks_in_queue(datasets,
                                              task_cmd = task_cmd, 
                                              force_submit = force_submit, 
                                              multisession = multisession, 
                                              parameters = self.parameters)
    def _compute(self):
        print(self.parameters)
        try:
            import matplotlib
            matplotlib.use('Agg') # launch with Agg background because tkinter can cause crashes in strange situations
        except:
            print('Could not set matplotlib Agg engine.')
            pass
        # this performs the actual spike sorting.
        datasets = pd.DataFrame((self.schema.EphysRecording.ProbeFile() & self.dataset_key).fetch())
        # check if a probe was selected
        if 'probe' in self.parameters.keys():
            datasets = datasets[datasets.probe_num.values == self.parameters['probe']]
        for probe_num in np.unique(datasets.probe_num):
            self.set_job_status(job_log = f'Sorting {probe_num}')
            files = datasets[datasets.probe_num.values == probe_num]
            dset = []
            for i,f in files.iterrows():
                if 'ap.cbin' in f.file_path or 'ap.ch' in f.file_path:
                    dset.append(i)
                elif 'ap.meta' in f.file_path: # requires a metadata file (spikeglx)
                    dset.append(i)
            dset = files.loc[dset]
            if not len(dset):
                print(files)
                raise(ValueError(f'Could not find ap.cbin files for probe {probe_num}'))
            localfiles = self.get_files(dset, allowed_extensions = ['.ap.bin'])
            probepath = list(filter(lambda x: str(x).endswith('bin'),localfiles))
            if 'kilosort' in self.parameters['algorithm_name']:
                from spks.sorting import run_kilosort
            dredge_motion_correction = False
            if 'dredge' in self.parameters.keys():
                dredge_motion_correction = self.parameters['dredge']
            # run kilosort using spks for preprocessing.
            if self.parameters['algorithm_name'] == 'spks_kilosort2.5':      
                results_folder = run_kilosort(version = '2.5',sessionfiles = probepath,
                                              temporary_folder = prefs['scratch_path'],
                                              do_post_processing = False,
                                              motion_correction = self.parameters['motion_correction'],
                                              thresholds = self.parameters['thresholds'],
                                              lowpass = self.parameters['low_pass'],
                                              highpass = self.parameters['high_pass'],
                                              dredge_motion_correction = dredge_motion_correction)
            elif self.parameters['algorithm_name'] == 'spks_kilosort3.0':      
                results_folder = run_kilosort(version = '3.0',
                                              sessionfiles = probepath,
                                              temporary_folder = prefs['scratch_path'],
                                              do_post_processing = False,
                                              motion_correction = self.parameters['motion_correction'],
                                              thresholds = self.parameters['thresholds'],
                                              lowpass = self.parameters['low_pass'],
                                              highpass = self.parameters['high_pass'],
                                              dredge_motion_correction = dredge_motion_correction)
            elif self.parameters['algorithm_name'] == 'spks_kilosort4.0':      
                results_folder = run_kilosort(version = '4.0',
                                              sessionfiles = probepath,
                                              temporary_folder = prefs['scratch_path'],
                                              do_post_processing = False,
                                              motion_correction = self.parameters['motion_correction'],
                                              thresholds = self.parameters['thresholds'],
                                              lowpass = self.parameters['low_pass'],
                                              highpass = self.parameters['high_pass'],
                                              dredge_motion_correction = dredge_motion_correction)
            elif self.parameters['algorithm_name'] == 'spks_mountainsort5':
                raise(NotImplemented(f"[{self.name} job] - Algorithm {self.parameters['algorithm_name']} not implemented."))
            else:
                raise(NotImplemented(f"[{self.name} job] - Algorithm {self.parameters['algorithm_name']} not implemented."))
            self.set_job_status(job_log = f'Probe {probe_num} sorted, running post-processing.')
            try: # attempt to close all figures before using joblib 
                import pylab as plt
                plt.close('all')
            except:
                pass
            self.postprocess_and_insert(results_folder,
                                        probe_num = probe_num,
                                        remove_duplicates = True,
                                        n_pre_samples = 45)
            self.unregister_safe_exit() # in case these get triggered by shutdown
            try:
                from joblib.externals.loky import get_reusable_executor
                get_reusable_executor().shutdown(wait=True)

            except:
                print(f'[{self.name} job] Tried to clear joblib Loky executers and failed.')
            self.register_safe_exit() # put it back..

            if not self.keep_intermediate:
                # delete results_folder
                print(f'[{self.name} job] Removing the results folder {results_folder}.')
                import shutil
                shutil.rmtree(results_folder)
                # delete local files if they did not exist
                if not self.files_existed:
                    for f in localfiles:
                        os.unlink(f)
            else:
                # send the files to AWS # this will skip the raw data binary files
                dataset = dict(**self.dataset_key[0])
                dataset['dataset_name'] = f'spike_sorting/imec{probe_num}/{self.parameter_set_num}/{self.parameters["algorithm_name"]}'
                src = [a for a in Path(results_folder).rglob('*') if not (str(a).endswith('.bin') or 
                                                                          str(a).endswith('.dat') or
                                                                          str(a).endswith('.hdf5'))] # get the files minus the .bin/.dat
                new_src = [Path(prefs['local_paths'][0])/a for a in self.schema.AnalysisFile().generate_filepaths(src,dataset)]
                from tqdm.auto import tqdm
                import shutil
                def move(s, d):
                    '''Move a single file, creating parent directories as needed.'''
                    Path(d).parent.mkdir(exist_ok=True, parents=True)
                    shutil.move(s, d)
                res = Parallel(n_jobs=DEFAULT_N_JOBS)(delayed(move)(s,n) for s,n in tqdm(zip(src,new_src),total = len(src),desc = "Copying to intermediate files to local_path"))
                # copy to a local path
                #print(f'[{self.name} job] Moving the temporary files from {results_folder} to {kept_results_folder}')
                filekeys = self.schema.AnalysisFile().upload_files(new_src,dataset)

                base_key = dict((self.schema.EphysRecording & self.dataset_key[0]).proj().fetch1(),
                                probe_num = probe_num,
                                parameter_set_num = self.parameter_set_num)
                self.schema.SpikeSorting.IntermediateFiles.insert([dict(base_key,**k) for k in filekeys])
                shutil.rmtree(results_folder)
                print(f'[{self.name} job] Removing the results folder {results_folder}. Intermediate results kept and uploaded to AnalysisFile.')


    def prepare_results(self, results_folder, probe_num, remove_duplicates, n_pre_samples):
        '''Parse Kilosort output and prepare structured dicts for database insertion.

        Optionally removes duplicate spikes across shanks, then collects per-unit
        spike times, positions, amplitudes, and template features.

        Parameters
        ----------
        results_folder : Path or str
            Kilosort output directory.
        probe_num : int
            Probe index (used to build keys and stream names).
        remove_duplicates : bool
            Run ``Clusters.remove_duplicate_spikes()`` before parsing.
        n_pre_samples : int
            Number of samples before the spike peak (stored in ``SpikeSorting``).

        Returns
        -------
        tuple
            ``(clu, base_keys, ssdict, udict, featurestosave, events)``
            where ``clu`` is the ``spks.Clusters`` object and the rest are
            lists of dicts ready for DataJoint insertion.
        '''
        from spks import Clusters
        if remove_duplicates:
            clu = Clusters(results_folder, get_waveforms = False, get_metrics = False)
            clu.remove_duplicate_spikes(
                overwrite_phy = True,
                remove_cross_duplicates = self.parameters['remove_cross_duplicates']) 
            del clu
        clu = Clusters(results_folder, get_waveforms = False, get_metrics = False)
        clu.compute_template_amplitudes_and_depths()
        base_keys = []
        ssdict = [] # sortings
        udict = [] # unit
        featurestosave = [] # features
        events = [] # events for sync

        for ifile,(dataset_key,(o,e)) in enumerate(zip(self.dataset_key,clu.metadata['file_offsets'])):
            base_keys.append(dict((self.schema.EphysRecording & dataset_key).proj().fetch1(),
                                  probe_num = probe_num,
                                  parameter_set_num = self.parameter_set_num))
            # this handles multi-session results
            idx = np.where((clu.spike_times>=o) & (clu.spike_times<e))[0]
            spikes = clu.spike_times[idx].flatten().astype(np.uint64) - o # subtract start of session
            clusters = clu.spike_clusters[idx]
            amplitudes = clu.spike_amplitudes[idx].flatten().astype(np.float32)
            positions = clu.spike_positions[idx,:].astype(np.float32)

            ssdict.append(dict(base_keys[-1],
                               n_pre_samples = n_pre_samples,
                               n_sorted_units = len(np.unique(clusters)),
                               n_detected_spikes = len(spikes),
                               sorting_datetime = datetime.fromtimestamp(
                                   Path(results_folder).stat().st_ctime),
                               sorting_channel_indices = clu.channel_map.flatten(),
                               sorting_channel_coords = clu.channel_positions))
            udict.append([]) # list of lists
            for iclu in np.unique(clusters):
                subidx = np.where(clusters == iclu)[0]
                udict[-1].append(dict(
                    base_keys[-1],unit_id = iclu,
                    spike_positions = positions[subidx],
                    spike_times = spikes[subidx],
                    spike_amplitudes = amplitudes[subidx]))

            featurestosave.append(dict(template_features = clu.spike_pc_features[idx].astype(np.float32),
                                       spike_templates = clu.spike_templates[idx],
                                       cluster_indices = clusters,
                                       whitening_matrix = clu.whitening_matrix,
                                       templates = clu.templates, # maybe should save only existing templates?
                                       template_feature_ind = clu.template_pc_features_ind))
            k = f'file{ifile}_sync_onsets'
            stream_name = f'imec{probe_num}' # to save the events and files
            events.append([])
            if k in clu.metadata.keys():
                for ik in clu.metadata[k].keys():
                    ev = clu.metadata[k][ik].astype(np.uint64)
                    events[-1].append(dict((self.schema.Dataset & base_keys[-1]).proj().fetch1(), 
                                       stream_name = stream_name,
                                       event_name = str(k),
                                       event_timestamps = ev))
        return clu,base_keys,ssdict, udict, featurestosave, events

    def postprocess_and_insert(self,
                               results_folder,
                               probe_num,
                               remove_duplicates = True,
                               n_pre_samples = 45):
        '''Does the preprocessing for a spike sorting and inserts'''
        # get the results in a dictionary and remove duplicates
        clu,base_keys,ssdicts, udicts, featurestosave, events = self.prepare_results(results_folder,
                                                                                     probe_num,
                                                                                     remove_duplicates,
                                                                                     n_pre_samples)
        # save the features to a file, will take around 2 min for a 1h recording.
        if not featurestosave[0]['template_features'] is None:
            for idset,features in enumerate(featurestosave):
                save_dict_to_h5(Path(results_folder)/f'features{idset}.hdf5',features)
        n_jobs = DEFAULT_N_JOBS  # gets the default number of jobs from labdata
        # extract the waveforms from the binary file
        n_jobs_wave = n_jobs
        for idset,(base_key,dataset_key,udict,ssdict,offsets,event) in enumerate(zip(
            base_keys,
            self.dataset_key,
            udicts,
            ssdicts,
            clu.metadata['file_offsets'],
            events)):
            if len(udict) > 800:
                n_jobs_wave = 3 # to prevent running out of memory when collecting waveforms
            udict, binaryfile, nchannels,res = self.extract_waveforms(udict,
                                                                      clu.metadata['nchannels'],
                                                                      results_folder,
                                                                      n_pre_samples,
                                                                      n_jobs_wave,
                                                                      offsets = offsets)
            def median_waves(r, gains):
                '''Compute the gain-corrected median waveform across extracted snippets.'''
                if not r is None:
                    return np.median(r.astype(np.float32), axis=0) * gains
                else:
                    return None
            waves_dict = []
            extras = dict(compression = 'gzip',
                          compression_opts = 1,
                          chunks = True, 
                          shuffle = True)
            from tqdm.auto import tqdm
            # print(f'Collecting waveforms and saving for dataset:{idset}.')
            # save these to zarr to be compressed faster?
            if self.use_hdf5: # zarr not implemented yet.
                import h5py as h5
                with h5.File(Path(results_folder)/f'waveforms{idset}.hdf5','w') as wavefile:
                    for u,w in tqdm(zip(udict,res), total = len(udict),
                                    desc = f'Saving waveforms to file [dataset {idset}]'):
                        m = median_waves(w, gains = clu.channel_gains)
                        if not w is None:
                            waves_dict.append(dict(base_key,
                                                    unit_id = u['unit_id'],
                                                    waveform_median = m))
                            # save to the file
                            wavefile.create_dataset(str(u['unit_id'])+'/waveforms',data = w,**extras)
                            wavefile.create_dataset(str(u['unit_id'])+'/indices',data = u['waveform_indices'],**extras)
                        else:
                            print(f"Unit {u['unit_id']} had no spikes extracted (dataset {idset})")
            stream_name = f'imec{probe_num}' # to save the events and files, also defined elsewhere (fix that on re-factor)
            src = [Path(results_folder)/f'waveforms{idset}.hdf5',Path(results_folder)/f'features{idset}.hdf5']
            dataset = dict(**dataset_key)
            dataset['dataset_name'] = f'spike_sorting/{stream_name}/{self.parameter_set_num}'

            filekeys = self.schema.AnalysisFile().upload_files(src,dataset)
            ssdict['waveforms_file'] = filekeys[0]['file_path']
            ssdict['waveforms_storage'] = filekeys[0]['storage']
            if not featurestosave[idset]['template_features'] is None:
                ssdict['features_file'] = filekeys[1]['file_path']
                ssdict['features_storage'] = filekeys[1]['storage']
            # insert the syncs
            if len(event):
                # Add stream
                print(f'Inserting {len(event)} events.')
                self.schema.DatasetEvents.insert1(dict(dataset_key,
                                                       stream_name = stream_name),
                                                  skip_duplicates = True,
                                                  allow_direct_insert = True)
                self.schema.DatasetEvents.Digital.insert(event,
                                                         skip_duplicates = True,
                                                         allow_direct_insert = True)

            # inserts
            import logging
            logging.getLogger('datajoint').setLevel(logging.WARNING)
            # these can't be done in a safe way quickly so if they fail we have delete SpikeSorting
            self.schema.SpikeSorting.insert1(ssdict,skip_duplicates = True)
            # Insert datajoint in parallel.
            parallel_insert(self.schema.schema_project,'SpikeSorting.Unit',udict, n_jobs = DEFAULT_N_JOBS,
                            skip_duplicates = True, ignore_extra_fields = True)
            parallel_insert(self.schema.schema_project,'SpikeSorting.Waveforms',waves_dict, n_jobs = DEFAULT_N_JOBS,
                            skip_duplicates = True, ignore_extra_fields = True)

            # Add a segment from a random location.
            from spks.io import map_binary
            dat = map_binary(binaryfile, nchannels = nchannels)
            nsamples = int(clu.sampling_rate*2)
            offset_samples = int(np.random.uniform(offsets[0]+nsamples, offsets[1]-nsamples-1))
            self.schema.SpikeSorting.Segment.insert1(dict(base_key,
                                            segment_num = 1,
                                            offset_samples = offset_samples - offsets[0],
                                            segment = np.array(dat[offset_samples : offset_samples + nsamples])))
            del dat
        self.set_job_status(job_log = f'Completed {base_key}')
        from labdata.schema import UnitMetrics
        # limit number of jobs because of memory constraints
        self.schema.UnitMetrics.populate(base_key, display_progress = True, processes = int(max(1,np.ceil(n_jobs/2))))

    def extract_waveforms(self, udict, nchannels, results_folder, n_pre_samples, n_jobs, offsets):
        '''Select spike indices and extract waveform snippets from the filtered binary.

        Parameters
        ----------
        udict : list of dict
            Unit dicts (from ``prepare_results``); will have ``waveform_indices``
            added in place.
        nchannels : int
            Number of channels in the binary file.
        results_folder : Path
            Kilosort output directory (must contain ``filtered_recording*.bin``).
        n_pre_samples : int
            Samples before spike peak.
        n_jobs : int
            Parallel workers for waveform extraction.
        offsets : tuple of int
            ``(start_sample, end_sample)`` of this dataset within the binary.

        Returns
        -------
        tuple
            ``(udict, binaryfile, nchannels, res)`` where ``res`` is a
            generator of waveform arrays (one per unit).
        '''
        # extract the waveforms
        from spks.io import map_binary
        # if not 'waveforms_from_sorter' in self.parameters.keys(): 
        #     self.parameters['waveforms_from_sorter'] = False
        # if not self.parameters['waveforms_from_sorter']:    
        binaryfile = list(Path(results_folder).glob("filtered_recording*.bin"))[0]
        # else: # not implemented
        #     binaryfile = list(Path(results_folder).glob("temp_wh.dat"))[0]
        #     nchannels = clu.metadata['nchannels']
        dat = map_binary(binaryfile, nchannels = nchannels) # to get the duration

        udict = select_random_waveforms(udict, 
                                        wpre = n_pre_samples, 
                                        wpost = n_pre_samples,
                                        duration = offsets[-1]-offsets[0])
        del dat
        res = get_waveforms_from_binary(binaryfile, nchannels,
                                        [u['waveform_indices']+offsets[0] for u in udict],
                                        wpre = n_pre_samples,
                                        wpost = n_pre_samples,
                                        n_jobs = n_jobs)
        return udict, binaryfile, nchannels, res