<commit_before>"""
Module related to saving

"""

from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals
from __future__ import absolute_import

from builtins import range
from builtins import open
from future import standard_library
standard_library.install_aliases()

# required to make json saving work in Python 2/3
try:
    to_unicode = unicode
except NameError:
    to_unicode = str

import os
from time import time
from datetime import datetime
import pickle as pk
from . import gather
from . import utils
from ..specs import Dict, ODict


#------------------------------------------------------------------------------
# Save JSON (Python 2/3 compatible)
#------------------------------------------------------------------------------
def saveJSON(fileName, data):
    """
    Function for/to <short description of `netpyne.sim.save.saveJSON`>

    Parameters
    ----------
    fileName : <type>
        <Short description of fileName>
        **Default:** *required*

    data : <type>
        <Short description of data>
        **Default:** *required*


    """


    import json, io
    from .utils import NpSerializer

    with io.open(fileName, 'w', encoding='utf-8') as fileObj:
        str_ = json.dumps(data,
                          indent=4, sort_keys=True,
                          separators=(',', ': '), ensure_ascii=False,
                          cls=NpSerializer)
        fileObj.write(to_unicode(str_))


#------------------------------------------------------------------------------
# Save data
#------------------------------------------------------------------------------
def saveData(include=None, filename=None, saveLFP=True):
    """
    Function to save simulation data to file

    Parameters
    ----------
    include : list
        What data to save
        **Default:** ``sim.cfg.saveDataInclude``
        **Options:** The list may include any combination of the following: ``'simData'``, ``'simConfig'``, ``'netParams'``, ``'net'``.

    filename : str
        Path and file name to save data to
        **Default:** ``None``

    """

    from .. import sim

    if sim.rank == 0 and not getattr(sim.net, 'allCells', None) and not getattr(sim, 'allSimData', None): needGather = True
    else: needGather = False
    if needGather: gather.gatherData()

    if filename: sim.cfg.filename = filename

    if sim.rank == 0:
        sim.timing('start', 'saveTime')
        import os

        # copy source files
        if isinstance(sim.cfg.backupCfgFile, list) and len(sim.cfg.backupCfgFile) == 2:
            simName = sim.cfg.simLabel if sim.cfg.simLabel else os.path.basename(sim.cfg.filename)
            print(('Copying cfg file %s ... ' % simName))
            source = sim.cfg.backupCfgFile[0]
            targetFolder = sim.cfg.backupCfgFile[1]
            # make dir
            try:
                os.mkdir(targetFolder)
            except OSError:
                if not os.path.exists(targetFolder):
                    print(' Could not create target folder: %s' % (targetFolder))
            # copy file
            targetFile = targetFolder + '/' + simName + '_cfg.py'
            if os.path.exists(targetFile):
                print(' Removing prior cfg file' , targetFile)
                os.system('rm ' + targetFile)
            os.system('cp ' + source + ' ' + targetFile)


        # create folder if missing
        targetFolder = os.path.dirname(sim.cfg.filename)
        if targetFolder and not os.path.exists(targetFolder):
            try:
                os.mkdir(targetFolder)
            except OSError:
                print(' Could not create target folder: %s' % (targetFolder))

        # saving data
        if not include: include = sim.cfg.saveDataInclude
        dataSave = {}
        net = {}

        dataSave['netpyne_version'] = sim.version(show=False)
        dataSave['netpyne_changeset'] = sim.gitChangeset(show=False)

        if getattr(sim.net.params, 'version', None): dataSave['netParams_version'] = sim.net.params.version
        if 'netParams' in include:
            sim.net.params.__dict__.pop('_labelid', None)
            net['params'] = utils.replaceFuncObj(sim.net.params.__dict__)
        if 'net' in include: include.extend(['netPops', 'netCells'])
        if 'netCells' in include and hasattr(sim.net, 'allCells'): net['cells'] = sim.net.allCells
        if 'netPops' in include and hasattr(sim.net, 'allPops'): net['pops'] = sim.net.allPops
        if net: dataSave['net'] = net
        if 'simConfig' in include: dataSave['simConfig'] = sim.cfg.__dict__
        if 'simData' in include:
            if saveLFP:
                if 'LFP' in sim.allSimData:
                    sim.allSimData['LFP'] = sim.allSimData['LFP'].tolist()
                #if hasattr(sim.net, 'recXElectrode'):
                #    dataSave['net']['recXElectrode'] = sim.net.recXElectrode
            dataSave['simData'] = sim.allSimData


        if dataSave:
            if sim.cfg.timestampFilename:
                timestamp = time()
                timestampStr = '-' + datetime.fromtimestamp(timestamp).strftime('%Y%m%d_%H%M%S')
            else:
                timestampStr = ''

            if hasattr(sim.cfg, 'saveFolder') and hasattr(sim.cfg, 'simLabel') and len(sim.cfg.simLabel):
                filePath = os.path.join(sim.cfg.saveFolder, sim.cfg.simLabel + '_data' + timestampStr)
            else:
                filePath = sim.cfg.filename + '_data' + timestampStr
            
            # Save to pickle file
            if sim.cfg.savePickle:
                import pickle
                dataSave = utils.replaceDictODict(dataSave)
                print(('Saving output as %s ... ' % (filePath + '.pkl')))
                with open(filePath+'.pkl', 'wb') as fileObj:
                    pickle.dump(dataSave, fileObj)
                print('Finished saving!')

            # Save to dpk file
            if sim.cfg.saveDpk:
                import gzip
                print(('Saving output as %s ... ' % (filePath+'.dpk')))
                #fn=filePath #.split('.')
                gzip.open(filePath, 'wb').write(pk.dumps(dataSave)) # write compressed string
                print('Finished saving!')

            # Save to json file
            if sim.cfg.saveJson:
                # Make it work for Python 2+3 and with Unicode
                print(('Saving output as %s ... ' % (filePath+'.json ')))
                #dataSave = utils.replaceDictODict(dataSave)  # not required since json saves as dict
                sim.saveJSON(filePath+'.json', dataSave)
                print('Finished saving!')

            # Save to mat file
            if sim.cfg.saveMat:
                from scipy.io import savemat
                print(('Saving output as %s ... ' % (filePath+'.mat')))
                savemat(filePath+'.mat', utils.tupleToList(utils.replaceNoneObj(dataSave)))  # replace None and {} with [] so can save in .mat format
                print('Finished saving!')

            # Save to HDF5 file (uses very inefficient hdf5storage module which supports dicts)
            if sim.cfg.saveHDF5:
                dataSaveUTF8 = utils._dict2utf8(utils.replaceNoneObj(dataSave)) # replace None and {} with [], and convert to utf
                import hdf5storage
                print(('Saving output as %s... ' % (filePath+'.hdf5')))
                hdf5storage.writes(dataSaveUTF8, filename=filePath+'.hdf5')
                print('Finished saving!')

            # Save to CSV file (currently only saves spikes)
            if sim.cfg.saveCSV:
                if 'simData' in dataSave:
                    import csv
                    print(('Saving output as %s ... ' % (filePath+'.csv')))
                    writer = csv.writer(open(filePath+'.csv', 'wb'))
                    for dic in dataSave['simData']:
                        for values in dic:
                            writer.writerow(values)
                    print('Finished saving!')

            # Save to Dat file(s)<commit_msg>wait until file exists when saving json<commit_after>"""
Module related to saving

"""

from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals
from __future__ import absolute_import

from builtins import range
from builtins import open
from future import standard_library
standard_library.install_aliases()

# required to make json saving work in Python 2/3
try:
    to_unicode = unicode
except NameError:
    to_unicode = str

import os
from time import time, sleep
from datetime import datetime
import pickle as pk
from . import gather
from . import utils
from ..specs import Dict, ODict


#------------------------------------------------------------------------------
# Save JSON (Python 2/3 compatible)
#------------------------------------------------------------------------------
def saveJSON(fileName, data, checkFileTimeout=0):
    """
    Function for/to <short description of `netpyne.sim.save.saveJSON`>

    Parameters
    ----------
    fileName : <type>
        <Short description of fileName>
        **Default:** *required*

    data : <type>
        <Short description of data>
        **Default:** *required*

    checkFileTimeout: timeout (sec)
        if >0 then will check if file exists before continuing 
        at 0.1 ms intervals for the timeout specified in secs


    """


    import json, io
    from .utils import NpSerializer

    with io.open(fileName, 'w', encoding='utf-8') as fileObj:
        str_ = json.dumps(data,
                          indent=4, sort_keys=True,
                          separators=(',', ': '), ensure_ascii=False,
                          cls=NpSerializer)
        fileObj.write(to_unicode(str_))


    if checkFile:
        sleepTime = 0.1
        timeoutCyles = checkFileTimeout / sleepTime 
        cycles = 0
        while not os.path.exists(fileName) and cycles <= timeOutCycles:
            sleep(sleepTime)
            


#------------------------------------------------------------------------------
# Save data
#------------------------------------------------------------------------------
def saveData(include=None, filename=None, saveLFP=True):
    """
    Function to save simulation data to file

    Parameters
    ----------
    include : list
        What data to save
        **Default:** ``sim.cfg.saveDataInclude``
        **Options:** The list may include any combination of the following: ``'simData'``, ``'simConfig'``, ``'netParams'``, ``'net'``.

    filename : str
        Path and file name to save data to
        **Default:** ``None``

    """

    from .. import sim

    if sim.rank == 0 and not getattr(sim.net, 'allCells', None) and not getattr(sim, 'allSimData', None): needGather = True
    else: needGather = False
    if needGather: gather.gatherData()

    if filename: sim.cfg.filename = filename

    if sim.rank == 0:
        sim.timing('start', 'saveTime')
        import os

        # copy source files
        if isinstance(sim.cfg.backupCfgFile, list) and len(sim.cfg.backupCfgFile) == 2:
            simName = sim.cfg.simLabel if sim.cfg.simLabel else os.path.basename(sim.cfg.filename)
            print(('Copying cfg file %s ... ' % simName))
            source = sim.cfg.backupCfgFile[0]
            targetFolder = sim.cfg.backupCfgFile[1]
            # make dir
            try:
                os.mkdir(targetFolder)
            except OSError:
                if not os.path.exists(targetFolder):
                    print(' Could not create target folder: %s' % (targetFolder))
            # copy file
            targetFile = targetFolder + '/' + simName + '_cfg.py'
            if os.path.exists(targetFile):
                print(' Removing prior cfg file' , targetFile)
                os.system('rm ' + targetFile)
            os.system('cp ' + source + ' ' + targetFile)


        # create folder if missing
        targetFolder = os.path.dirname(sim.cfg.filename)
        if targetFolder and not os.path.exists(targetFolder):
            try:
                os.mkdir(targetFolder)
            except OSError:
                print(' Could not create target folder: %s' % (targetFolder))

        # saving data
        if not include: include = sim.cfg.saveDataInclude
        dataSave = {}
        net = {}

        dataSave['netpyne_version'] = sim.version(show=False)
        dataSave['netpyne_changeset'] = sim.gitChangeset(show=False)

        if getattr(sim.net.params, 'version', None): dataSave['netParams_version'] = sim.net.params.version
        if 'netParams' in include:
            sim.net.params.__dict__.pop('_labelid', None)
            net['params'] = utils.replaceFuncObj(sim.net.params.__dict__)
        if 'net' in include: include.extend(['netPops', 'netCells'])
        if 'netCells' in include and hasattr(sim.net, 'allCells'): net['cells'] = sim.net.allCells
        if 'netPops' in include and hasattr(sim.net, 'allPops'): net['pops'] = sim.net.allPops
        if net: dataSave['net'] = net
        if 'simConfig' in include: dataSave['simConfig'] = sim.cfg.__dict__
        if 'simData' in include:
            if saveLFP:
                if 'LFP' in sim.allSimData:
                    sim.allSimData['LFP'] = sim.allSimData['LFP'].tolist()
                #if hasattr(sim.net, 'recXElectrode'):
                #    dataSave['net']['recXElectrode'] = sim.net.recXElectrode
            dataSave['simData'] = sim.allSimData


        if dataSave:
            if sim.cfg.timestampFilename:
                timestamp = time()
                timestampStr = '-' + datetime.fromtimestamp(timestamp).strftime('%Y%m%d_%H%M%S')
            else:
                timestampStr = ''

            if hasattr(sim.cfg, 'saveFolder') and hasattr(sim.cfg, 'simLabel') and len(sim.cfg.simLabel):
                filePath = os.path.join(sim.cfg.saveFolder, sim.cfg.simLabel + '_data' + timestampStr)
            else:
                filePath = sim.cfg.filename + '_data' + timestampStr
            
            # Save to pickle file
            if sim.cfg.savePickle:
                import pickle
                dataSave = utils.replaceDictODict(dataSave)
                print(('Saving output as %s ... ' % (filePath + '.pkl')))
                with open(filePath+'.pkl', 'wb') as fileObj:
                    pickle.dump(dataSave, fileObj)
                print('Finished saving!')

            # Save to dpk file
            if sim.cfg.saveDpk:
                import gzip
                print(('Saving output as %s ... ' % (filePath+'.dpk')))
                #fn=filePath #.split('.')
                gzip.open(filePath, 'wb').write(pk.dumps(dataSave)) # write compressed string
                print('Finished saving!')

            # Save to json file
            if sim.cfg.saveJson:
                # Make it work for Python 2+3 and with Unicode
                print(('Saving output as %s ... ' % (filePath+'.json ')))
                #dataSave = utils.replaceDictODict(dataSave)  # not required since json saves as dict
                sim.saveJSON(filePath+'.json', dataSave, checkFileTimeout=5)
                print('Finished saving!')

            # Save to mat file
            if sim.cfg.saveMat:
                from scipy.io import savemat
                print(('Saving output as %s ... ' % (filePath+'.mat')))
                savemat(filePath+'.mat', utils.tupleToList(utils.replaceNoneObj(dataSave)))  # replace None and {} with [] so can save in .mat format
                print('Finished saving!')

            # Save to HDF5 file (uses very inefficient hdf5storage module which supports dicts)
            if sim.cfg.saveHDF5:
                dataSaveUTF8 = utils._dict2utf8(utils.replaceNoneObj(dataSave)) # replace None and {} with [], and convert to utf
                import hdf5storage
                print(('Saving output as %s... ' % (filePath+'.hdf5')))
                hdf5storage.writes(dataSaveUTF8, filename=filePath+'.hdf5')
                print('Finished saving!')

            # Save to CSV file (currently only saves spikes)
            if sim.cfg.saveCSV:
                if 'simData' in dataSave:
                    import csv
                    print(('Saving output as %s ... ' % (filePath+'.csv')))
                    writer = csv.writer(open(filePath+'.csv', 'wb'))
                    for dic in dataSave['simData']:
                        for values in dic:
                            writer.writerow(values)
                    print('Finished saving!')

            # Save to Dat file(s)
{'fraction_non_alphanumeric': 0.14860960635608522, 'fraction_numerical': 0.0033706512579752016, 'mean_word_length': 1.741296815706979, 'pattern_counts': {'":': 0, '<': 14, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5588638', 'n_tokens_mistral': 4670, 'n_tokens_neox': 4452, 'n_words': 1215}