#!/usr/bin/python
import os
import h5py
import sys
import shutil
import traceback
import subprocess

import nwb
from allensdk.internal.core.lims_pipeline_module import PipelineModule

# development/debugging code
#infile = "Ndnf-IRES2-dgCre_Ai14-256189.05.01-compressed.nwb"
#outfile = "foo.nwb"
#if len(sys.argv) == 1:
#    sys.argv.append(infile)
#    sys.argv.append(outfile)

# this script is meant to clone the core functionality of the 
#   existing (Igor) Hdf5->Nwb converter.
# the previous converter performed two distinct tasks. In this iteration,
#   those tasks will be split into separate modules. This module will
#   perform the file conversion. A second module will analyze the file
#   and extract sweep data

# window for leading test pulse, in seconds
PULSE_LEN = 0.1
EXPERIMENT_START_TIME = 0.75


def main():
    module = PipelineModule()
    jin = module.input_data()

    infile = jin["input_nwb"]
    outfile = jin["output_nwb"]

    # a temporary nwb file must be created. this is that file's name
    tmpfile = outfile + ".tmp"

    # create temp file and make modifications to it using h5py
    shutil.copy2(infile, tmpfile)
    f = h5py.File(tmpfile, "a")
    # change dataset names in acquisition time series to match that
    #   of existing ephys NWB files
    # also rescale the contents of 'data' fields to match the scaling
    #   in original files
    acq = f["acquisition/timeseries"]
    sweep_nums = []
    for k, v in acq.iteritems():
        # parse out sweep number
        try:
            num = int(k[5:10])
        except:
            print("Error - unexpected sweep name encountered in IGOR nwb file")
            print("Sweep called: '%s'" % k)
            print("Expecting 5-digit sweep number between chars 5 and 9")
            sys.exit(1)
        swp = "Sweep_%d" % num
        # rename objects
        try:
            acq.move(k, swp)
            ts = acq[swp]
            ts.move("stimulus_description", "aibs_stimulus_description")
        except:
            print("*** Error renaming HDF5 object in %s" % swp)
            type_, value_, traceback_ = sys.exc_info()
            print traceback.print_tb(traceback_)
            sys.exit(1)
        # rescale contents of data so conversion is 1.0
        try:
            data = ts["data"]
            scale = float(data.attrs["conversion"])
            data[...] = data.value * scale
            data.attrs["conversion"] = 1.0
        except:
            print("*** Error rescaling data in %s" % swp)
            type_, value_, traceback_ = sys.exc_info()
            print traceback.print_tb(traceback_)
            sys.exit(1)
        # keep track of sweep numbers
        sweep_nums.append("%d"%num)
        
    ###################################
    #... ditto for stimulus time series
    stim = f["stimulus/presentation"]
    for k, v in stim.iteritems():
        # parse out sweep number
        try:
            num = int(k[5:10])
        except:
            print("Error - unexpected sweep name encountered in IGOR nwb file")
            print("Sweep called: '%s'" % k)
            print("Expecting 5-digit sweep number between chars 5 and 9")
            sys.exit(1)
        swp = "Sweep_%d" % num
        try:
            stim.move(k, swp)
        except:
            print("Error renaming HDF5 group from %s to %s" % (k, swp))
            sys.exit(1)
        # rescale contents of data so conversion is 1.0
        try:
            ts = stim[swp]
            data = ts["data"]
            scale = float(data.attrs["conversion"])
            data[...] = data.value * scale
            data.attrs["conversion"] = 1.0
        except:
            print("*** Error rescaling data in %s" % swp)
            type_, value_, traceback_ = sys.exc_info()
            print traceback.print_tb(traceback_)
            sys.exit(1)
        
    f.close()

    ####################################################################
    # re-open file w/ nwb library and add indexing (epochs)
    nd = nwb.NWB(filename=tmpfile, modify=True)
    for num in sweep_nums:
        ts = nd.file_pointer["acquisition/timeseries/Sweep_" + num]
        # sweep epoch
        t0 = ts["starting_time"].value
        rate = float(ts["starting_time"].attrs["rate"])
        n = float(ts["num_samples"].value)
        t1 = t0 + (n-1) * rate
        ep = nd.create_epoch("Sweep_" + num, t0, t1)
        ep.add_timeseries("stimulus", "stimulus/presentation/Sweep_"+num)
        ep.add_timeseries("response", "acquisition/timeseries/Sweep_"+num)
        ep.finalize()
        if "CurrentClampSeries" in ts.attrs["ancestry"]:
            # test pulse epoch
            t0 = ts["starting_time"].value
            t1 = t0 + PULSE_LEN
            ep = nd.create_epoch("TestPulse_" + num, t0, t1)
            ep.add_timeseries("stimulus", "stimulus/presentation/Sweep_"+num)
            ep.add_timeseries("response", "acquisition/timeseries/Sweep_"+num)
            ep.finalize()
            # experiment epoch
            t0 = ts["starting_time"].value
            t1 = t0 + (n-1) * rate
            t0 += EXPERIMENT_START_TIME
            ep = nd.create_epoch("Experiment_" + num, t0, t1)
            ep.add_timeseries("stimulus", "stimulus/presentation/Sweep_"+num)
            ep.add_timeseries("response", "acquisition/timeseries/Sweep_"+num)
            ep.finalize()
    nd.close()

    # rescaling the contents of the data arrays causes the file to grow
    # execute hdf5-repack to get it back to its original size
    try:
        print("Repacking hdf5 file with compression")
        process = subprocess.Popen(["h5repack", "-f", "GZIP=4", tmpfile, outfile], stdout=subprocess.PIPE)
        process.wait()
    except:
        print("Unable to run h5repack on temporary nwb file")
        print("--------------------------------------------")
        raise

    try:
        print("Removing temporary file")
        os.remove(tmpfile)
    except:
        print("Unable to delete temporary file ('%s')" % tmpfile)
        raise

    # done (nothing to return)
    module.write_output_data({})



if __name__=='__main__': main()

