from allensdk.api.queries.biophysical_api import \
    BiophysicalApi
from allensdk.api.queries.cell_types_api import CellTypesApi
from allensdk.api.queries.rma_api import RmaApi
import os, sys, shutil


#bp = BiophysicalApi('http://api.brain-map.org')
#bp.cache_stimulus = True # change to False to not download the large stimulus NWB file
# neuronal_model_id = 472451419    # get this from the web site as above
# bp.cache_data(neuronal_model_id, working_directory='neuronal_model')

# According to this there are 49 biophysical models.
# http://api.brain-map.org/api/v2/data/query.json?q=model::NeuronalModel,rma::include,neuronal_model_template[name$eq%27Biophysical%20-%20perisomatic%27],rma::options[num_rows$eqall]

# Note, am I supposed to be only archiving Biophysical models or also GLIFs?

class BiophysicalArchiver(object):
    def __init__(self, archive_dir=None):
        self.bp = BiophysicalApi('http://api.brain-map.org')
        self.bp.cache_stimulus = True # change to False to not download the large stimulus NWB file
        self.cta = CellTypesApi()
        self.rma = RmaApi()
        self.neuronal_model_download_endpoint = 'http://celltypes.brain-map.org/neuronal_model/download/'
        self.template_names = {}
        self.nwb_list = []
        
        if archive_dir == None:
            archive_dir = '.'
        self.archive_dir = archive_dir
    
    def get_template_names(self):
        template_response = self.rma.model_query('NeuronalModelTemplate')
        self.template_names = { t['id']: str(t['name']).replace(' ', '_') for t in template_response}
    
    def get_cells(self):
        return self.cta.list_cells(True, True)
    
    def get_neuronal_models(self, specimen_ids):
        return self.rma.model_query('NeuronalModel',
                                    criteria='specimen[id$in%s]' % ','.join(str(i) for i in specimen_ids),
                                    include='specimen',
                                    num_rows='all')
    
    def get_stimulus_file(self, neuronal_model_id):
        result = self.rma.model_query('NeuronalModel',
                                      criteria='[id$eq%d]' % (neuronal_model_id),
                                      include="specimen(ephys_result(well_known_files(well_known_file_type[name$il'NWB*'])))",
                                      tabular=['path'])
        
        stimulus_filename = result[0]['path']
        
        return stimulus_filename
        
        
        stimulus_filename = os.path.basename(result[0]['path'])
        
        return stimulus_filename
    
    def archive_cell(self, ephys_result_id, specimen_id, template, neuronal_model_id):
        url = self.neuronal_model_download_endpoint + "/%d" % (neuronal_model_id)
        file_name = os.path.join(self.archive_dir, 'ephys_result_%d_specimen_%d_%s_neuronal_model_%d.zip' % (ephys_result_id,
                                                                                                             specimen_id,
                                                                                                             template,
                                                                                                             neuronal_model_id))
        self.rma.retrieve_file_over_http(url, file_name)
        nwb_file = self.get_stimulus_file(neuronal_model_id)
        shutil.copy(nwb_file, self.archive_dir) 
        self.nwb_list.append("%s\t%s" % (os.path.basename(nwb_file),
                                         file_name))
    
if __name__ == '__main__':
    archive_dir = sys.argv[-1] # /data/informatics/mousecelltypes/model_cache_may_2015
    ba = BiophysicalArchiver(archive_dir)
    ba.get_template_names()
    cells = ba.get_cells()
    
    specimen_ids = (cell['id'] for cell in cells)
    neuronal_models = ba.get_neuronal_models(specimen_ids)
    for nm in neuronal_models:
        ephys_result_id = nm['specimen']['ephys_result_id']
        template_id = nm['neuronal_model_template_id'] 
        if template_id in ba.template_names:
            template = ba.template_names[template_id]
        else:
            template = 'unknown'
        ba.archive_cell(ephys_result_id, nm['specimen_id'], template, nm['id'])
    with open(os.path.join(ba.archive_dir, 'STIMULUS.csv'), 'w') as f:
        f.write("\n".join(ba.nwb_list))