Rules

class UploadRule():
    def __init__(self,job_id,prefs = None):
        '''

Rule to apply on upload. 

        1) Checksum on the files; compare with provided (reserve job if use_db)
        2) Apply function
        3) Checksum on the output - the files that changed
        4) Submit upload
        5) Update tables 

Can submit job on slurm, some of these can be long or take resources.

        '''
        self.rule_name = 'default'

        self.job_id = job_id
        self.src_paths = None
        self.processed_paths = None
        self.dst_paths = None
        self.inserted_files = []
        if prefs is None:
            from ..utils import prefs
        self.prefs = prefs
        self.local_path = self.prefs['local_paths'][0]
        self.dataset_key = None # will get written on upload, use in _post_upload
        self.max_concurrent = -1  # maximum number of concurrent UploadJobs that can run [-1 is infinite].
        self.schema = None

    def apply(self):
        # parse inputs
        from ..schema import UploadJob, File, dj
        if not self.job_id is None:
            with dj.conn().transaction:
                self.jobquery = (UploadJob() & dict(job_id = self.job_id))
                job_status = self.jobquery.fetch(as_dict = True)
                if len(job_status):
                    if job_status[0]['job_waiting']:
                        # check here if there are other jobs running for the same rule on this host
                        if self.max_concurrent > 0:
                            kk = dict(job_host = self.prefs['hostname'],
                                      job_waiting = 0,
                                      job_status = 'WORKING',
                                      job_rule = None)
                            if not self.rule_name == 'default':
                                kk['job_rule'] = self.rule_name
                            number_of_running_jobs = len(UploadJob() & kk)
                            if number_of_running_jobs >= self.max_concurrent:
                                print(f"Job {self.job_id} can not run because there are {number_of_running_jobs} for the rule {self.rule_name} [limit:{self.max_concurrent}].")
                                self.set_job_status(job_status = 'WAITING', job_waiting = 1,
                                                    job_log = f'Waiting for {number_of_running_jobs} running {self.rule_name} to complete [limit:{self.max_concurrent}].')
                                return
                        self.set_job_status(job_status = 'WORKING', job_starttime = datetime.now(), job_waiting = 0) # take the job
                        self.schema = load_project_schema(job_status[0]['project_name'])
                    else:
                        print(f"Job {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((self.schema.UploadJob.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 Upload.AssignedFiles.')
            raise ValueError(f'Could not find files for {self.job_id} in Upload.AssignedFiles.')
        self.upload_storage = self.jobquery.fetch('upload_storage')[0]
        if self.upload_storage is None:
            if 'upload_storage' in self.prefs.keys():
                self.upload_storage = self.prefs['upload_storage']

        # this should not fail because we have to keep track of errors, should update the table
        src = [Path(self.local_path) / p for p in self.src_paths.src_path.values]
        try:
            comparison = compare_md5s(src,self.src_paths.src_md5.values)
        except Exception as err:
            print('File not found for {0}?'.format(Path(self.src_paths.src_path.iloc[0]).parent))
            self.set_job_status(job_status = 'FAILED', job_log = f'FILE NOT FOUND; check file transfer {err}.')
            return

        if not comparison:
            print('CHECKSUM FAILED for {0}'.format(Path(self.src_paths.src_path.iloc[0]).parent))
            self.set_job_status(job_status = 'FAILED',job_log = 'MD5 CHECKSUM failed; check file transfer.')
            return # exit.
        import traceback
        try:
            paths = self._apply_rule() # can use the src_paths
            self._upload()             # compare the hashes after
        except Exception as err:
            # log the error
            print('There was an error processing or uploading this dataset.')
            print(err)
            self.set_job_status(job_status = 'FAILED',job_log = f'ERROR uploading {traceback.format_exc()}')
            return
        try:
            self._post_upload()        # so the rules can insert tables and all.
        except Exception as err:
            # log the error
            print('There was an error with the post-upload this dataset.')
            print(err)
            self.set_job_status(job_status = 'FAILED',job_log = f'POST-UPLOAD: {traceback.format_exc()}')
            return 
        return paths # so apply can output paths.

    def set_job_status(self, job_status = 'FAILED',job_waiting = 0, **kwargs):
        from ..schema import UploadJob
        if not self.job_id is None:
            kk = dict(job_id = self.job_id,
                      job_waiting = job_waiting,
                      job_status = job_status,
                      **kwargs)
            if 'job_log' in kk.keys():
                if len(kk['job_log']) > 500:
                    kk['job_log'] = kk['job_log'][:500-1]
            if job_status == 'FAILED':
                print(f'Check job_id {self.job_id} : {job_status}')
                kk['job_host'] = self.prefs['hostname'], # write the hostname so we know where it failed.
            UploadJob.update1(kk)

    def _handle_processed_and_src_paths(self, processed_files,new_files):
        '''
        Put the files in the proper place and compute checksums for new files.
        Call this from the apply method.
        '''
        n_jobs = DEFAULT_N_JOBS
        self.processed_paths = []
        for f in processed_files:
            i = np.where(self.src_paths.src_path == f)[0][0]
            self.processed_paths.append(self.src_paths.iloc[i])
            self.src_paths.drop(self.src_paths.iloc[i].name,axis = 0,inplace = True)
            self.src_paths.reset_index(drop=True,inplace = True)
        self.processed_paths = pd.DataFrame(self.processed_paths).reset_index(drop=True)        

        res = Parallel(n_jobs = n_jobs)(delayed(_checksum_files)(
            path,
            local_path = self.local_path) for path in new_files)
        for r in res:
            r['job_id'] = self.job_id
        self.src_paths = pd.concat([self.src_paths,pd.DataFrame(res)], ignore_index=True)
        # drop duplicate paths in case there are any
        self.src_paths = self.src_paths.drop_duplicates(subset=['src_path'], keep='last')

    def _post_upload(self):
        return

    def _upload(self):
        # this reads the attributes and uploads
        # It also puts the files in the Tables

        # destination in the bucket is actually the path
        dst = [k for k in self.src_paths.src_path.values]
        # source is the place where data are
        src = [Path(self.local_path) / p for p in self.src_paths.src_path.values] # same as md5
        # s3 copy in parallel hashes were compared before so no need to do it now.
        self.set_job_status(job_status = 'WORKING',
                            job_log = datetime.now().strftime(f'Uploading {len(src)} files %Y %m %d %H:%M:%S'),
                            job_waiting = 0)
        copy_to_s3(src, dst, md5_checksum=None,storage_name=self.upload_storage)
        self.set_job_status(job_status = 'WORKING',
                            job_log = datetime.now().strftime('%Y %m %d %H:%M:%S'),
                            job_endtime = datetime.now(),
                            job_waiting = 0)
        import traceback
        with self.schema.dj.conn().transaction:  # make it all update at the same time
            # insert to Files so we know where to get the data
            insfiles = []
            for i,f in self.src_paths.iterrows():
                insfiles.append(dict(file_path = f.src_path,
                                     storage = self.upload_storage,
                                     file_datetime = f.src_datetime,
                                     file_size = f.src_size,
                                     file_md5 = f.src_md5))
            if len(insfiles):
                try:
                    self.schema.File.insert(insfiles)
                except Exception as err:
                    print(f'There was an error inserting the files to the File table. {insfiles}')
                    print(err)
                    self.set_job_status(job_status = 'FAILED',
                                        job_log = f'DUPLICATES? {[k["file_path"] for k in insfiles]}')
                    return
            # Add to dataset?
            job = self.jobquery.fetch(as_dict=True)[0]
            # check if it has a dataset
            if all([not job[a] is None for a in ['subject_name','session_name','dataset_name']]):
                for i,p in enumerate(insfiles):
                    insfiles[i] = dict(subject_name = job['subject_name'],
                                       session_name = job['session_name'],
                                       dataset_name = job['dataset_name'],
                                       file_path = p['file_path'],
                                       storage = self.upload_storage)
                if len(insfiles):
                    seskey = dict(subject_name = job['subject_name'],
                                  session_name = job['session_name'])
                    if not len(self.schema.Session & seskey): # check if the session needs to be added
                        # TODO: This needs to parse the session datetime from the path otherwise it won't be able to insert.
                        self.schema.Session.insert1(seskey,skip_duplicates = True)
                    dsetkey=dict(subject_name = job['subject_name'],
                                 session_name = job['session_name'], 
                                 dataset_name = job['dataset_name'])
                    if not len(self.schema.Dataset & dsetkey):# check if the dataset needs to be added
                        self.schema.Dataset.insert1(dsetkey,
                                                    skip_duplicates = True)
                    self.schema.Dataset.DataFiles.insert(insfiles)
                self.dataset_key = dict(subject_name = job['subject_name'],
                                        session_name = job['session_name'],
                                        dataset_name = job['dataset_name'])
            self.inserted_files += insfiles
            # Insert the processed files so the deletions are safe
            if not self.processed_paths is None:
                ins = []
                for i,f in self.processed_paths.iterrows():
                    ins.append(dict(file_path = f.src_path,
                                    file_datetime = f.src_datetime,
                                    file_size = f.src_size,
                                    file_md5 = f.src_md5))
                if len(ins):
                    self.schema.ProcessedFile.insert(ins)
            if len(self.inserted_files): # keep the job in the queue 
                # completed
                self.set_job_status(job_status = 'COMPLETED',
                                    job_log = f'UPLOADED {len(self.inserted_files)}',
                                    job_endtime = datetime.now(),
                                    job_waiting = 0)
            else:
                self.set_job_status(job_status = 'FAILED',
                                    job_log = 'No files inserted',
                                    job_endtime = datetime.now(),
                                    job_waiting = 0)

    def _apply_rule(self):
        # this rule does nothing, so the src_paths are going to be empty,
        # and the "paths" are going to be the src_paths
        self.processed_paths = None # processed paths are just the same, no file changed, so no need to do anything.
        # needs to compute the checksum on all the new files
        return self.src_paths.src_path.values

class FixedBrainRule(UploadRule):
    def __init__(self, job_id, prefs = None):
        super(FixedBrainRule,self).__init__(job_id = job_id, prefs = prefs)
        self.rule_name = 'fixed_brain'
        self.metadata = dict(num_channels = None,
                             um_per_pixel = None,
                             hardware = 'unknown',
                             width = None,
                             height = None)
        self.max_concurrent = 2

    def parse_metadata(self,data):
        '''
        Reads metadata from a MultifolderTiffStack
        '''
        from tifffile import TiffFile
        tf = TiffFile(data.filenames[0][0])
        metadata = tf.ome_metadata
        um_per_pixel = []
        if not metadata is None: # to work with non-OME files 
            for t in metadata.split('\n'):
                if 'Pixels' in t:
                    for p in t.strip('>').split(' '):
                        if 'PhysicalSize' in p:
                            um_per_pixel.append(float(p.split('=')[1].split('"')[1]))
            self.metadata['um_per_pixel'] = um_per_pixel[::-1]

            if 'UltraII' in metadata:
                self.metadata['hardware'] = 'lightsheet UltraII'

        self.metadata['num_channels'] = len(data.filenames)
        self.metadata['height'] = data.shape[-2]
        self.metadata['width'] = data.shape[-1]

    def _apply_rule(self):
        extensions = ['.ome.tif','.tif','.TIFF','.TIF']
        for extension in extensions:
            files_to_compress = list(filter(lambda x: (extension in x),
                                            self.src_paths.src_path.values))
            if len(files_to_compress):
                break # found files
        local_path = self.prefs['local_paths'][0]
        local_path = Path(local_path)

        new_files = []
        full_file_paths = [local_path/f for f in files_to_compress]
        channel_folders = np.sort(np.unique([p.parent for p in full_file_paths]))
        print(f'Found folders: {channel_folders}')
        data = MultifolderTiffStack(channel_folders = channel_folders,
                                    extensions = extensions)
        compressed_file = channel_folders[0].parent.with_suffix('.zarr.zip')
        # read the first file and get metadata
        self.parse_metadata(data)
        # compression to zarr.zip
        self.set_job_status(job_status = 'WORKING',
                            job_log = datetime.now().strftime('Compressing data %Y %m %d %H:%M:%S'),
                            job_waiting = 0)
        z1 = compress_imaging_stack(data,
                                    compressed_file,
                                    chunksize = 4,
                                    compression = 'zstd',
                                    clevel = 6,
                                    shuffle = 1,
                                    filters = [])
        new_files.append(str(compressed_file).replace(str(local_path),'').strip(os.sep))
        if len(new_files):
            self._handle_processed_and_src_paths(files_to_compress, new_files)
        return self.src_paths
    # need to implement post-upload for lightsheet.
    def _post_upload(self):
        if not self.dataset_key is None:
            if len(self.inserted_files) > 0:
                storage = self.inserted_files[0]['storage']
                filenames = [f['file_path'] for f in self.inserted_files if f['file_path'].endswith('.zarr.zip')]
                if not len(filenames):
                    print('Could not find zarr compressed stack.')
                self.schema.FixedBrain().insert1(dict(self.dataset_key,
                                         file_path = filenames[0],
                                         storage = storage,
                                         **self.metadata),allow_direct_insert=True)

class EphysRule(UploadRule):
    def __init__(self, job_id,prefs = None):
        super(EphysRule,self).__init__(job_id = job_id,prefs = prefs)
        self.rule_name = 'ephys'
        self.max_concurrent = 2

    def _apply_rule(self):
        # compress ap files or lf files.
        files_to_compress = list(filter(lambda x: ('.ap.bin' in x) or
                                        ('.lf.bin' in x),
                                        self.src_paths.src_path.values))
        n_jobs = DEFAULT_N_JOBS
        # compress these in parallel, will work for multiprobe sessions faster?
        if len(files_to_compress): # in some cases data might have already been compressed
            self.set_job_status(job_status = 'WORKING',
                                job_log = datetime.now().strftime(f'Compressing {len(files_to_compress)} files [%Y %m %d %H:%M:%S]'))
            res = Parallel(n_jobs = n_jobs)(delayed(compress_ephys_file)(
                filename,
                local_path = self.local_path,
                n_jobs = n_jobs,
                prefs = self.prefs) for filename in files_to_compress)
            new_files = np.stack(res).flatten() # stack the resulting files and add them to the path
            self.set_job_status(job_status = 'WORKING',
                                job_log = datetime.now().strftime(f'Compressed {len(files_to_compress)} files [%Y %m %d %H:%M:%S]'))
            self._handle_processed_and_src_paths(files_to_compress, new_files)
        return self.src_paths

    def _post_upload(self):
        if not self.dataset_key is None:
            self.schema.EphysRecording().add_spikeglx_recording(self.dataset_key)
            self.schema.EphysRecording().add_nidq_events(self.dataset_key)

class TwoPhotonRule(UploadRule):
    def __init__(self, job_id,prefs = None):
        super(TwoPhotonRule,self).__init__(job_id = job_id, prefs = prefs)
        self.rule_name = 'two_photon'
        self.recording_metadata = None
        self.planes_metadata = []
        self.max_concurrent = 2    

    def _apply_rule(self):
        # compress ap files or lf files.
        files_to_compress = list(filter(lambda x: ('.sbx' in x),
                                        self.src_paths.src_path.values))
        if not len(files_to_compress):
            print('There are no sbx files.')
            files_to_compress = list(filter(lambda x: ('.tif' in x),
                                        self.src_paths.src_path.values))
            if len(files_to_compress):
                raise(ValueError('Tiff processing not implemented for this rule.'))

        local_path = self.prefs['local_paths'][0]
        local_path = Path(local_path)
        # TODO: make this work for tiff files also
        if len(files_to_compress):
            if len(files_to_compress)>1:
                raise(ValueError(f'Can only handle one file {files_to_compress}'))
            new_files = []
            for f in files_to_compress:
                compressed_file = self.process_sbx(local_path/f)
                # pass the file path without the "local_path"!
                new_files.append(str(compressed_file).replace(str(local_path),'').strip(os.sep))
            if len(new_files):
                self._handle_processed_and_src_paths(files_to_compress, new_files)
        return self.src_paths

    def process_sbx(self,sbxfile):
        try:
            from sbxreader import sbx_memmap
        except:
            msg = f'Could not import SBXREADER: "pip install sbxreader" on {pref["hostname"]}'
            print(msg)
            self.set_job_status(job_status = 'FAILED',job_log = msg,job_waiting = 0)
        sbx = sbx_memmap(sbxfile)
        compressed_file = sbxfile.with_suffix('.zarr.zip')
        z1 = compress_imaging_stack(sbx,
                                    compressed_file,
                                    chunksize = 128,
                                    compression = 'zstd',#'blosc2',
                                    clevel = 6,
                                    shuffle = 1)
        nframes,nplanes,nchannels,H,W = z1.shape

        pmt_gain = [sbx.metadata['pmt0_gain']]
        if nchannels>1:
            pmt_gain+= [sbx.metadata['pmt1_gain']],

        self.recording_metadata = dict(n_planes = nplanes,
                                       n_channels = nchannels,
                                       n_frames = nframes,
                                       width = W,
                                       height = H,
                                       magnification = sbx.metadata['magnification'],
                                       objective_angle = sbx.metadata['stage_angle'],
                                       objective = sbx.metadata['objective'],
                                       um_per_pixel = np.array([sbx.metadata['um_per_pixel_x'],
                                                                sbx.metadata['um_per_pixel_y']]),
                                       frame_rate = sbx.metadata['frame_rate']/nchannels,
                                       scanning_mode = sbx.metadata['scanning_mode'],
                                       pmt_gain = np.array(pmt_gain),
                                       imaging_software = f"scanbox_{sbx.metadata['scanbox_version']}")
        for iplane in range(nplanes):
            depth = sbx.metadata['stage_pos'][-1]
            if nplanes > 1:
                depth += sbx.metadata['etl_pos'][iplane] - sbx.metadata['etl_pos'][0]

            self.planes_metadata.append(dict(plane_num = iplane,
                                             plane_depth = depth))
        return compressed_file

    def _post_upload(self):
        if not self.dataset_key is None:
            if len(self.inserted_files) > 0:
                storage = self.inserted_files[0]['storage']
                filenames = [f['file_path'] for f in self.inserted_files if f['file_path'].endswith('.zarr.zip')]
                if not len(filenames):
                    print('Could not find zarr compressed stack.')
                self.schema.TwoPhoton().insert1(dict(self.dataset_key,
                                         file_path = filenames[0],
                                         storage = storage,
                                         **self.recording_metadata),allow_direct_insert=True)
                self.schema.TwoPhoton.Plane().insert([dict(self.dataset_key,**d) for d in self.planes_metadata],allow_direct_insert=True)

class OnePhotonRule(UploadRule):
    def __init__(self, job_id, prefs = None):
        super(OnePhotonRule,self).__init__(job_id = job_id, prefs = prefs)
        self.rule_name = 'one_photon'
        self.imaging_metadata = None
        self.max_concurrent = 2    

    def _apply_rule(self):
        files_to_compress = list(filter(lambda x: ('.dat' in x),
                                        self.src_paths.src_path.values))
        # TODO: make this work for tiff files also
        local_path = self.prefs['local_paths'][0]
        local_path = Path(local_path)
        if len(files_to_compress):
            if len(files_to_compress)>1:
                raise(ValueError(f'Can only handle one file {files_to_compress}'))
            new_files = []
            for f in files_to_compress:
                filename = local_path/f
                data = mmap_wfield_binary(filename)
                compressed_file = filename.with_suffix('.zarr.zip')
                # compression to zarr.zip
                self.set_job_status(job_status = 'WORKING',
                                    job_log = datetime.now().strftime('Compressing data %Y %m %d %H:%M:%S'),
                                    job_waiting = 0)
                z1 = compress_imaging_stack(data,
                                            compressed_file,
                                            chunksize = 16,
                                            compression = 'zstd',
                                            clevel = 6,
                                            shuffle = 1,
                                            filters = [])
                # pass the file path without the "local_path"!
                new_files.append(str(compressed_file).replace(str(local_path),'').strip(os.sep))
            if len(new_files):
                self._handle_processed_and_src_paths(files_to_compress, new_files)
        return self.src_paths
    # need to implement post-upload for widefield imaging.
    def _post_upload(self):
        if not self.dataset_key is None:
            insert_widefield_dataset(self.schema,self.dataset_key,local_paths=[self.local_path])

class MiniscopeRule(UploadRule):
    def __init__(self, job_id, prefs = None):
        super(MiniscopeRule,self).__init__(job_id = job_id, prefs = prefs)
        self.rule_name = 'miniscope'
        self.imaging_metadata = None
        self.max_concurrent = 4    

    def _apply_rule(self):
        from natsort import natsorted
        files_to_compress = natsorted([str(s) for s in list(filter(lambda x: ('.avi' in x),
                                                         self.src_paths.src_path.values))])
        local_path = self.prefs['local_paths'][0]
        local_path = Path(local_path)

        new_files = []
        files_to_process = [local_path/f for f in files_to_compress]
        # open the video stack
        from ..stacks import VideoStack
        data = VideoStack(files_to_process)  # this does not support multiple channels at the moment.
        compressed_file = (Path(files_to_process[0]).parent/'miniscope_stack').with_suffix('.zarr.zip')
        self.set_job_status(job_status = 'WORKING',
                            job_log = datetime.now().strftime('Compressing data %Y %m %d %H:%M:%S'),
                            job_waiting = 0)
        z1 = compress_imaging_stack(data,
                                    compressed_file,
                                    chunksize = 512,
                                    compression = 'zstd',
                                    clevel = 6,
                                    shuffle = 1,
                                    filters = [])
        new_files.append(str(compressed_file).replace(str(local_path),'').strip(os.sep))
        if len(new_files):
            self._handle_processed_and_src_paths(files_to_compress, new_files)
        return self.src_paths

    def _post_upload(self):
        if not self.dataset_key is None:
            insert_miniscope_dataset(self.schema,self.dataset_key, local_paths = [self.local_path])

class MultifolderTiffStack(object):
    def __init__(self,channel_folders, extensions = ['.ome.tif','.tif','.TIFF']):
        '''
        Simple class to access tiff files that are organized in a folders
        Each folder is a channel and contains multiple TIFF files.

        This is the format of the lightsheet microscope for example. 
        It is a place-holder class that should be modified to work for scanimage files also.
        '''
        self.nchannels = len(channel_folders)
        self.filenames = []
        for folder in channel_folders:
            for extension in extensions:
                tiffs = list(Path(folder).rglob('*' + extension))
                if len(tiffs) > 0:
                    self.filenames.append(natsorted(tiffs))
                    break # found extension
        # read the first file
        im = read_ome_tif(self.filenames[0][0])
        if len(im.shape) == 2:
            # then the number frames is the number of tiff files
            self.nframes = len(self.filenames[0])
            self.frame_dims = im.shape
            self.dtype = im.dtype
        else:
            raise NotImplementedError('This needs to be changed to work with multiple-frame TIFF files.')
        self.shape = (self.nframes,self.nchannels,*self.frame_dims)
    def __len__(self):
        return self.nframes

    def __getitem__(self,*args):
        index = args[0]
        if type(index) is slice:
            idx1 = range(*index.indices(self.nframes))
        elif type(index) in [int,np.int32, np.int64]: # just one frame
            idx1 = [index]
        else: # np.array?
            idx1 = index
        img = []
        for i in idx1:
            img.append(self.get(i))
        return np.stack(img).squeeze().reshape(len(img),self.nchannels, *self.frame_dims)

    def get(self,idx):
        return np.stack([read_ome_tif(files[idx]) for files in self.filenames])