text: from mountaintools import client as mt
  from PIL import Image
  import json
  import pandas as pd
  from spikeforest import SFMdaRecordingExtractor, SFMdaSortingExtractor
  import spikeextractors as si
  
  
  def kb_read_text_file(fname):
      return mt.loadText(path=fname)
  
  
  def kb_read_json_file(fname):
      return mt.loadObject(path=fname)
  
  
  class SFSortingResult():
      def __init__(self, obj, recording):
          self._obj = obj
          self._recording = recording
  
      def getObject(self):
          return self._obj
  
      def recording(self):
          return self._recording
  
      def sorterName(self):
          return self._obj['sorter']['name']
  
      def plotNames(self):
          plots = self._obj['summary'].get('plots', dict())
          return list(plots.keys())
  
      def sorting(self):
          return SFMdaSortingExtractor(firings_file=self._obj['firings'])
  
      def plot(self, name, format='image'):
          plots = self._obj['summary'].get('plots', dict())
          url = plots[name]
          if format == 'url':
              return url
          else:
              path = mt.realizeFile(url)
              if format == 'image':
                  return Image.open(path)
              elif format == 'path':
                  return path
              else:
                  raise Exception('Invalid format: ' + format)
  
      def comparisonWithTruth(self, *, format='dataframe'):
          A = self._obj['comparison_with_truth']
          if not A:
              return None
          if format == 'html':
              return kb_read_text_file(A['html'])
          else:
              B = kb_read_json_file(A['json'])
              if format == 'json':
                  return B
              elif format == 'dataframe':
                  return pd.DataFrame(B).transpose()
              else:
                  raise Exception('Invalid format: ' + format)
  
  
  class SFRecording():
      def __init__(self, obj, study):
          self._obj = obj
          self._sorting_result_names = []
          self._sorting_results_by_name = dict()
          self._summary_result = None
          if 'summary' in obj:
              self._summary_result = obj['summary']
          self._study = study
  
      def getObject(self):
          return self._obj
  
      def getSummaryObject(self):
          return self._summary_result
  
      def study(self):
          return self._study
  
      def name(self):
          return self._obj.get('recording_name', self._obj.get('name'))
  
      def description(self):
          return self._obj['description']
  
      def directory(self):
          return self._obj['directory']
  
      def recordingFileIsLocal(self):
          fname = self.directory() + '/raw.mda'
          fname2 = mt.findFile(fname, local_only=True)
          if fname2 and (not _is_url(fname2)):
              return True
          return False
  
      def realizeRecordingFile(self):
          fname = self.directory() + '/raw.mda'
          return mt.realizeFile(fname)
  
      def firingsTrueFileIsLocal(self):
          fname = self.directory() + '/firings_true.mda'
          fname2 = mt.findFile(fname, local_only=True)
          if fname2 and (not _is_url(fname2)):
              return True
          return False
  
      def realizeFiringsTrueFile(self):
          fname = self.directory() + '/firings_true.mda'
          return mt.realizeFile(fname)
  
      def recordingExtractor(self, download=False):
          X = SFMdaRecordingExtractor(dataset_directory=self.directory(), download=download)
          if 'channels' in self._obj:
              if self._obj['channels']:
                  X = si.SubRecordingExtractor(parent_recording=X, channel_ids=self._obj['channels'])
          return X
  
      def sortingTrue(self):
          return SFMdaSortingExtractor(firings_file=self.directory() + '/firings_true.mda')
  
      def plotNames(self):
          if not self._summary_result:
              return []
          plots = self._summary_result.get('plots', dict())
          return list(plots.keys())
  
      def plot(self, name, format='image'):
          plots = self._summary_result.get('plots', dict())
          url = plots[name]
          if format == 'url':
              return url
          else:
              path = mt.realizeFile(url)
              if format == 'image':
                  return Image.open(path)
              elif format == 'path':
                  return path
              else:
                  raise Exception('Invalid format: ' + format)
  
      def trueUnitsInfo(self, format='dataframe'):
          if not self._summary_result:
              return None
          B = kb_read_json_file(self._summary_result['true_units_info'])
          if format == 'json':
              return B
          elif format == 'dataframe':
              return pd.DataFrame(B)
          else:
              raise Exception('Invalid format: ' + format)
  
      def setSummaryResult(self, obj):
          self._summary_result = obj
  
      def addSortingResult(self, obj):
          sorter_name = obj['sorter']['name']
          if sorter_name in self._sorting_results_by_name:
              print('Sorting result already in recording: {}'.format(sorter_name))
          else:
              R = SFSortingResult(obj, self)
              self._sorting_result_names.append(sorter_name)
              self._sorting_results_by_name[sorter_name] = R
  
      def sortingResultNames(self):
          return self._sorting_result_names
  
      def sortingResult(self, name):
          return self._sorting_results_by_name.get(name, None)
  
  
  class SFStudy():
      def __init__(self, obj):
          self._obj = obj
          self._recordings_by_name = dict()
          self._recording_names = []
  
      def getObject(self):
          return self._obj
  
      def name(self):
          return self._obj.get('recording_name', self._obj.get('name'))
  
      def description(self):
          return self._obj['description']
  
      def addRecording(self, obj):
          name = obj.get('recording_name', obj.get('name'))
          if name in self._recordings_by_name:
              print('Recording already in study: ' + name)
          else:
              self._recording_names.append(name)
              D = SFRecording(obj, self)
              self._recordings_by_name[name] = D
  
      def recordingNames(self):
          return self._recording_names
  
      def recording(self, name):
          return self._recordings_by_name.get(name, None)
  
  
  class SFData():
      def __init__(self):
          self._studies_by_name = dict()
          self._study_names = []
  
      def loadStudy(self, study):
          name = study['name']
          if name in self._studies_by_name:
              print('Study already loaded: ' + name)
          else:
              self._study_names.append(study['name'])
              S = SFStudy(study)
              self._studies_by_name[name] = S
  
      def loadStudies(self, studies):
          for study in studies:
              self.loadStudy(study)
  
      def loadRecording(self, recording):
          study = recording.get('study_name', recording.get('study_name', recording.get('study')))
          self._studies_by_name[study].addRecording(recording)
  
      def loadRecordings2(self, recordings):
          for recording in recordings:
              self.loadRecording(recording)
  
      def loadRecordings(self, *, key=None, verbose=False):
          # old
          if key is None:
              key = dict(name='spikeforest_studies_processed')
          obj = mt.loadObject(key=key)
          studies = obj['studies']
          for study in studies:
              name = study['name']
              if name in self._studies_by_name:
                  print('Study already loaded: ' + name)
              else:
                  self._study_names.append(study['name'])
                  S = SFStudy(study)
                  self._studies_by_name[name] = S
          recordings = obj['recordings']
          print('recordings ===================================================================')
          print(recordings)
          for ds in recordings:
              study = ds.get('study_name', ds.get('study'))
              self._studies_by_name[study].addRecording(ds)
          if verbose:
              print('Loaded {} recordings'.format(len(recordings)))
  
      def loadSortingResults(self, sorting_results):
          for result in sorting_results:
              self.loadSortingResult(result)
  
      def loadSortingResult(self, X):
          study_name = X['recording'].get('study_name', X['recording'].get('study'))
          recording_name = X['recording'].get('recording_name', X['recording'].get('name'))
          # sorter_name=X['sorter']['name']
          S = self.study(study_name)
          if S:
              D = S.recording(recording_name)
              if D:
                  D.addSortingResult(X)
              else:
                  print('Warning: recording not found: ' + recording_name)
          else:
              print('Warning: study not found: ' + study_name)
  
      def studyNames(self):
          return self._study_names
  
      def study(self, name):
          return self._studies_by_name.get(name, None)
  
  
  def _is_url(path):
      return (path.startswith('http://') or path.startswith('https://'))
  