import logging
import sys
import os
import h5py
import subprocess
import shutil
import numpy as np
import traceback
import nwb.nwb as nwb
import nwb.nwbco as nwbco
import resource_file
from collections import defaultdict

from allensdk.core.nwb_data_set import NwbDataSet
import allensdk.ephys.extract_cell_features as extract_cell_features

from allensdk.internal.core.lims_pipeline_module import PipelineModule

# changes
"""
added to main block:
    "nwb_file": "igor_converted_256189.05.01.nwb",
    "publish_nwb": "to_publish.nwb",
    "metadata_file": "nwb_metadata.yml",


"""

local_dir = os.path.dirname(os.path.realpath(__file__))
if not local_dir.endswith('/'):
    local_dir += '/'


ELECTRODE_NAME = "Electrode 1"
ELECTRODE_PATH = "/general/intracellular_ephys/" + ELECTRODE_NAME

PIPELINE_NAME = "IVSCC"
PIPELINE_VERSION = "1.0"

def copy_val(old_ts, new_ts, name):
    if name in old_ts:
        val = old_ts[name].value
        attrs = {}
        for x in old_ts[name].attrs:
            if x == "neurodata_type":
                continue
            if x == "unit":
                continue
            attrs[x] = old_ts[name].attrs[x]
        new_ts.set_value(name, val, **attrs)

def copy_timeseries(timeseries, old_file, new_file, folder, metadata):
    try:
        name = ""
        for name in timeseries:
            old_ts = old_file[folder][name]
            family = old_ts.attrs["ancestry"]
            if family[-1] == "VoltageClampSeries":
                family = family[-1]
                category = "acquisition"
            elif family[-1] == "CurrentClampSeries":
                family = family[-1]
                category = "acquisition"
            elif family[-1] == "VoltageClampStimulusSeries":
                family = family[-1]
                category = "stimulus"
            elif family[-1] == "CurrentClampStimulusSeries":
                family = family[-1]
                category = "stimulus"
            else:
                raise Exception("Time series '%s' is of unknown type" % name)
            new_ts = new_file.create_timeseries(family, name, category)
            # copy data
            num_samples = old_ts["num_samples"].value
            data = old_ts["data"].value
            conversion = old_ts["data"].attrs["conversion"]
            resolution = old_ts["data"].attrs["resolution"]
            unit = old_ts["data"].attrs["unit"]
            new_ts.set_data(data, conversion=conversion, resolution=resolution, unit=unit)

            start_time = old_ts["starting_time"].value
            sampling_rate = old_ts["starting_time"].attrs["rate"]
            new_ts.set_time_by_rate(start_time, sampling_rate)
            new_ts.set_value("num_samples", num_samples)

            description = old_ts.attrs["description"]
            try:
                comments = old_ts.attrs["comments"]
            except:
                comments = old_ts.attrs["comment"]
            source = old_ts.attrs["source"]
            new_ts.set_value("comments", comments)
            new_ts.set_value("description", description)
            new_ts.set_value("source", source)
            
            copy_val(old_ts, new_ts, "electrode_name")
            copy_val(old_ts, new_ts, "capacitance_fast")
            copy_val(old_ts, new_ts, "capacitance_slow")
            copy_val(old_ts, new_ts, "resistance_comp_bandwidth")
            copy_val(old_ts, new_ts, "resistance_comp_correction")
            copy_val(old_ts, new_ts, "resistance_comp_prediction")
            copy_val(old_ts, new_ts, "whole_cell_capaictance_comp")
            copy_val(old_ts, new_ts, "whole_cell_series_resistance_comp")
            copy_val(old_ts, new_ts, "bias_current")
            copy_val(old_ts, new_ts, "bridge_balance")
            copy_val(old_ts, new_ts, "capacitance_compensation")
            copy_val(old_ts, new_ts, "stimulus_description")
            # 
            new_ts.finalize()
    except:
        print("** Error copying timeseries data **")
        print("** Timeseries: " + str(name))
        print("** Folder: " + folder)
        print("-----------------------------------")
        raise

def copy_epochs(timeseries, old_file, new_file, folder):
    try:
        for name in timeseries:
            anc = old_file["acquisition/timeseries/"+name].attrs["ancestry"]
            if anc[-1] == "VoltageClampSeries":
                continue
            num = int(name.split('_')[-1])
            # experiment block
            epname = "Experiment_%d" % num
            ep = old_file["epochs/%s" % epname]
            start = ep["start_time"].value
            stop = ep["stop_time"].value
            desc = ep["description"].value
            ep = new_file.create_epoch(epname, start, stop)
            ep.set_value("description", desc)
            ep.add_timeseries("stimulus", "/stimulus/presentation/Sweep_%d" % num)
            ep.add_timeseries("response", "/acquisition/timeseries/Sweep_%d" % num)
            ep.finalize()
            # test-pulse block
            epname = "TestPulse_%d" % num
            ep = old_file["epochs/%s" % epname]
            start = ep["start_time"].value
            stop = ep["stop_time"].value
            desc = ep["description"].value
            ep = new_file.create_epoch(epname, start, stop)
            ep.set_value("description", desc)
            ep.add_timeseries("stimulus", "/stimulus/presentation/Sweep_%d" % num)
            ep.add_timeseries("response", "/acquisition/timeseries/Sweep_%d" % num)
            ep.finalize()
            # sweep block
            epname = name
            ep = old_file["epochs/%s" % epname]
            start = ep["start_time"].value
            stop = ep["stop_time"].value
            desc = ep["description"].value
            ep = new_file.create_epoch(epname, start, stop)
            ep.set_value("description", desc)
            ep.add_timeseries("stimulus", "/stimulus/presentation/Sweep_%d" % num)
            ep.add_timeseries("response", "/acquisition/timeseries/Sweep_%d" % num)
            ep.finalize()

    except:
        print("** Error copying epoch data **")
        print("------------------------------")
        raise


def copy_file(infile, outfile, passing_sweeps, rsrc, metadata):
    print("Opening '%s'" % infile)
    old = h5py.File(infile, 'r')
    # top-level data
    try:
        vargs = {}
        vargs["identifier"] = old["identifier"].value[0] + ".edit"
        vargs["start_time"] = old["session_start_time"].value[0]
        vargs["description"] = old["session_description"].value[0]
        vargs["overwrite"] = True
        vargs["filename"] = outfile
        vargs["auto_compress"] = True
    except:
        print("** Error extracting top-level metadata from input file **")
        print("---------------------------------------------------------")
        raise

    print("Creating '%s'" % outfile)
    try:
        out = nwb.NWB(**vargs)
    except:
        print("** Error creating output file '%s' **" % outfile)
        print("---------------------------------------------------------")
        raise

    # make list of time series
    timeseries = []
    try:
        acq = old["acquisition/timeseries"]
        for ts in acq:
            timeseries.append(ts)
    except:
        print("** Error extracting timeseries list **")
        print("--------------------------------------")
        raise

    ts_list = []
    # TODO remove items from list that are not to be copied
    for num in passing_sweeps:
        swp = "Sweep_%d" % num
        for name in timeseries:
            if name == swp:
                ts_list.append(swp)
                break
    timeseries = ts_list

    # copy acquisition time series from source to destination files
    copy_timeseries(timeseries, old, out, "acquisition/timeseries", metadata)
    copy_timeseries(timeseries, old, out, "stimulus/presentation", metadata)

    copy_epochs(timeseries, old, out, "stimulus/presentation")

    write_metadata(out, rsrc, metadata)
        
    out.close()

def organize_metadata(ephys_roi_result):
    metadata = { 'sweeps': {} }
    
    cell_specimen = ephys_roi_result['specimens'][0]
    slice_specimen = ephys_roi_result['specimen']

    metadata['donor_id'] = cell_specimen['donor_id']
    metadata['specimen_name'] = cell_specimen['name']
    metadata['specimen_id'] = cell_specimen['id']
    try:
        metadata['species'] = slice_specimen['donor']['organism']["name"]
    except Exception as e:
        logging.error("Unable to read organism name from input.json file")
        raise

    # structure
    try:
        structure = cell_specimen['structure']
    except Exception as e:
        logging.error("Cell has no structure association.")
        raise
    
    soma_location = {}

    db_cell_soma_location = cell_specimen['cell_soma_locations'][0]
    soma_location = {}
    try:
        soma_location['cell_soma_location_x'] = 1e-9 * db_cell_soma_location['x']
        soma_location['cell_soma_location_y'] = 1e-9 * db_cell_soma_location['y']
        soma_location['cell_soma_location_z'] = 1e-9 * db_cell_soma_location['z']
        nd = db_cell_soma_location['normalized_depth']
        if nd is not None:
            soma_location['cell_soma_location_normalized_depth'] = nd 
    except Exception as e:
        logging.error(e.message)
        raise

    structure_info = {}
    try:
        structure_info['structure_id'] = structure['id']
        structure_info['structure_name'] =  structure['name']
        structure_info['structure_acronym'] = structure['acronym']
    except Exception as e:
        logging.error("Structure information is missing from input.json")
        raise 

    structure_info.update(soma_location)
    metadata['location'] = structure_info


    tags = cell_specimen["specimen_tags"]

    dend_trunc = None
    dend_type = None
    for i in range(len(tags)):
        name = tags[i]["name"]
        toks = name.split(" - ")
        if len(toks) != 2:
            continue
        if name.startswith("apical"):
            dend_trunc = toks[1]
        elif name.startswith("dendrite type"):
            dend_type = toks[1]

    if dend_trunc is None:
        raise Exception("Cell has no dendrite truncation tag.")

    if dend_type is None:
        raise Exception("Cell has no dendrite type tag.")

    metadata['dendrite_type'] = dend_type
    metadata['dendrite_trunc'] = dend_trunc
    
    metadata['ephys_roi_result_id'] = ephys_roi_result['id']
    metadata['seal_gohm'] = ephys_roi_result['seal_gohm']
    metadata['initial_access_resistance_mohm'] = ephys_roi_result['initial_access_resistance_mohm']

    slice_specimen = ephys_roi_result['specimen']
    donor = slice_specimen['donor']

    # gender
    try:
        metadata['gender'] = donor['gender']['name']
    except Exception as e:
        logging.error("Donor requires gender association.")
        raise

    # age
    try:
        age = donor['age']
    except Exception as e:
        logging.error("Donor requires age association.")
        raise

    metadata['age'] = {
        'date_of_birth': donor['date_of_birth'],
        'name': age['name']
        }


    # cre line and genotype are mouse-only
    if metadata['species'] == 'Mus musculus':
        genotypes = donor['genotypes']

        try:
            reporter_genotype = next( g for g in genotypes if g['genotype_type_id'] == 177835595 )
            metadata['cre_line'] = reporter_genotype['name']
        except Exception, e:
            logging.error("Could not find reporter genotype for mouse cell")
            raise

        metadata['genotype'] = {
            'description': [ g['description'] for g in genotypes ],
            'type': [ g['name'] for g in genotypes ]
            }
    else:
        logging.info("non-mouse cells do not have cre line or genotypes")

    # subject
    metadata['subject'] = {
        'subject_id': cell_specimen['donor_id'],
        'comments': 'subject_id value here corresponds to Allen Institute cell specimen "donor_id"'
        }

    # sweeps
    sweeps = cell_specimen['ephys_sweeps']
    for sweep in sweeps:
        if "invalid" in sweep and sweep["invalid"]:
            continue
        wfs = sweep['workflow_state']
        if wfs not in [ 'manual_passed', 'auto_passed' ]:
            logging.debug("skipping sweep %d, not passed" % sweep['sweep_number'])
            continue
            
        stimulus = sweep['ephys_stimulus']
        stimulus_type = stimulus['ephys_stimulus_type']


        metadata['sweeps'][sweep['sweep_number']] = {
            'stimulus_name': stimulus['description'],
            'stimulus_interval': sweep['stimulus_interval'],
            'stimulus_amplitude': sweep['stimulus_amplitude'],
            'stimulus_type_name': stimulus_type['name']
            }

    # IT-12498 add additional metadata to NWB file
    url = "http://help.brain-map.org/display/celltypes/Documentation"
    metadata["data_collection"] = "please see " + url
    metadata["protocol"] = "please see " + url
    metadata["pharmacology"] = "please see " + url
    metadata["citation_policy"] = "please see " + url
    metadata["institution"] = "Allen Institute for Brain Science"
    metadata["generated_by"] = ["pipeline", PIPELINE_NAME, "version", PIPELINE_VERSION]

    return metadata


def write_metadata(nwb_file, resources, metadata):
    nwb_file.set_metadata(nwbco.SEX, metadata['gender'])
    if 'cre_line' in metadata:
        nwb_file.set_metadata("aibs_cre_line", metadata['cre_line'])

    if 'genotype' in metadata:
        genotype = metadata['genotype']
        genotype_name = '; '.join(genotype['type'])
        nwb_file.set_metadata(nwbco.GENOTYPE, genotype_name, **genotype)

    nwb_file.set_metadata('generated_by', metadata['generated_by'])

    subject = metadata['subject']
    nwb_file.set_metadata(nwbco.SUBJECT, resources.get("subject"), **subject)

    age = metadata['age']
    nwb_file.set_metadata(nwbco.AGE, age['name'], **age)

    trode = ELECTRODE_NAME
    nwb_file.set_metadata(nwbco.INTRA_ELECTRODE_DESCRIPTION(trode), resources.get("electrode_description"))
    nwb_file.set_metadata(nwbco.INTRA_ELECTRODE_FILTERING(trode), resources.get("electrode_filtering"))
    nwb_file.set_metadata(nwbco.INTRA_ELECTRODE_DEVICE(trode), resources.get("electrode_device"))

    location = metadata['location']
    nwb_file.set_metadata(nwbco.INTRA_ELECTRODE_LOCATION(trode), location['structure_name'], **location)

    nwb_file.set_metadata(nwbco.INTRA_ELECTRODE_RESISTANCE(trode), resources.get("electrode_resistance"))
    nwb_file.set_metadata(nwbco.INTRA_ELECTRODE_SLICE(trode), resources.get("electrode_slice"))

    seal_gohm = str(metadata['seal_gohm'])
    nwb_file.set_metadata(nwbco.INTRA_ELECTRODE_SEAL(trode), seal_gohm + " GOhm")

    acc = str(metadata["initial_access_resistance_mohm"])
    nwb_file.set_metadata(nwbco.INTRA_ELECTRODE_INIT_ACCESS_RESISTANCE(trode), acc + " MOhm")

    session = { 'comments': 'session_id value corresponds to ephys_result_id' }
    nwb_file.set_metadata(nwbco.SESSION_ID, str(metadata['ephys_roi_result_id']), **session)

    nwb_file.set_metadata("aibs_specimen_name", metadata['specimen_name'])
    nwb_file.set_metadata("aibs_specimen_id", str(metadata['specimen_id']))

    nwb_file.set_metadata("aibs_dendrite_type", metadata['dendrite_type'])
    nwb_file.set_metadata("aibs_dendrite_trunc", metadata['dendrite_trunc'])    

    # IT-12498 add additional metadata to NWB file
    nwb_file.set_metadata(nwbco.DATA_COLLECTION, metadata["data_collection"])
    nwb_file.set_metadata(nwbco.INSTITUTION, metadata["institution"])
    nwb_file.set_metadata(nwbco.PROTOCOL, metadata["protocol"])
    nwb_file.set_metadata(nwbco.PHARMACOLOGY, metadata["pharmacology"])
    nwb_file.set_metadata("citation_policy", metadata["citation_policy"])
    nwb_file.set_metadata(nwbco.SPECIES, metadata['species'])



def main(jin):
    infile = jin[0]["nwb_file"]
    outfile = jin[0]["publish_nwb"]
    tmpfile = outfile + ".working"
    metafile = local_dir + jin[0]["metadata_file"]
    # load metadata stored in YML file
    metadata_desc_file = os.path.join(os.path.dirname(__file__), metafile)
    rsrc = resource_file.ResourceFile()
    rsrc.load(metadata_desc_file)

    #
    metadata = organize_metadata(jin[0])

    # TODO dig deeper here
    # only fetching metadata for passing sweeps
    passing_sweeps = metadata['sweeps'].keys()

    copy_file(infile, outfile, passing_sweeps, rsrc, metadata)

#    try:
#        shutil.copyfile(infile, tmpfile)
#    except:
#        print("Unable to copy '%s' to %s" % (infile, tmpfile))
#        print("----------------------------")
#        raise

#    # open NWB file so the modification date is updated
#    # add metadata then close file and do remaining manipulations using 
#    #   HDF5 library (except DF's legacy code that interfaces w/ nwb file
#    #   using nwb library)
#    args = {}
#    args["filename"] = tmpfile
#    args["modify"] = True
#    try:
#        nwb_file = nwb.NWB(**args)
#    except:
#        print("Error opening NWB file '%s'" % args["filename"])
#        raise
#    write_metadata(nwb_file, rsrc, metadata)
#    nwb_file.close()



    # open publish file directlya using HDF5 library
    # 1) remove hdf5 groups corresponding to failed sweeps
    # 2) add sweep-specific metadata data to file to match original publish
    #   format. this includes (acquisition and stimulus):
    #     aibs_stimulus_amplitude_pa
    #     aibs_stimulus_interval
    #     aibs_stimulus_name
    #     initial_access_resistance
    #     seal
    hdf = h5py.File(outfile, "r+")
#    ################################
#    # delete epochs, stim, recordings for non-passed sweeps
#    epochs = hdf["epochs/"]
#    for grp in epochs:
#        try:
#            num = int(str(grp).split('_')[-1])
#        except:
#            continue
#        if num not in passing_sweeps:
#            del epochs[str(grp)]
#    stim = hdf["stimulus/presentation"]
#    for grp in stim:
#        try:
#            num = int(str(grp).split('_')[-1])
#        except:
#            continue
#        if num not in passing_sweeps:
#            del stim[str(grp)]
#    acq = hdf["acquisition/timeseries"]
#    for grp in acq:
#        try:
#            num = int(str(grp).split('_')[-1])
#        except:
#            continue
#        if num not in passing_sweeps:
#            del acq[str(grp)]
    ################################
    # add data
    acq = hdf["acquisition/timeseries"]
    stim = hdf["stimulus/presentation"]
    sweeps = jin[0]["specimens"][0]["ephys_sweeps"]
    for grp in acq:
        try:
            num = int(str(grp).split('_')[-1])
        except:
            continue
        try:
            for sweep in sweeps:
                if sweep["sweep_number"] == num:
                    break
            if sweep["sweep_number"] != num:
                print sweep
                print num
                raise Exception("WTF")
            # stim amplitude
            amp = sweep["stimulus_amplitude"]
            if amp is None:
                amp = float('nan')
            else:
                amp = float(amp)
            ds = acq["Sweep_%d" % num].create_dataset("aibs_stimulus_amplitude_pa", data=amp)
            ds.attrs["neurodata_type"] = "Custom"
            ds = stim["Sweep_%d" % num].create_dataset("aibs_stimulus_amplitude_pa", data=amp)
            ds.attrs["neurodata_type"] = "Custom"
            # stim interval
            interval = sweep["stimulus_interval"]
            if interval is None:
                interval = float('nan')
            else:
                interval = float(interval)
            ds = acq["Sweep_%d" % num].create_dataset("aibs_stimulus_interval", data=interval)
            ds.attrs["neurodata_type"] = "Custom"
            ds = stim["Sweep_%d" % num].create_dataset("aibs_stimulus_interval", data=interval)
            ds.attrs["neurodata_type"] = "Custom"
            # stim name
            name = sweep["ephys_stimulus"]["ephys_stimulus_type"]["name"]
            ds = acq["Sweep_%d" % num].create_dataset("aibs_stimulus_name", data=name)
            ds.attrs["neurodata_type"] = "Custom"
            ds = stim["Sweep_%d" % num].create_dataset("aibs_stimulus_name", data=name)
            ds.attrs["neurodata_type"] = "Custom"
            # seal
            seal = jin[0]["seal_gohm"]
            if seal is None:
                seal = float('nan')
            else:
                seal = float(seal)
            ds = acq["Sweep_%d" % num].create_dataset("seal", data=seal)
            ds.attrs["neurodata_type"] = "Custom"
            ds = stim["Sweep_%d" % num].create_dataset("seal", data=seal)
            ds.attrs["neurodata_type"] = "Custom"
            # initial access resistance
            res = jin[0]["initial_access_resistance_mohm"]
            if res is None:
                res = float('nan')
            else:
                res = float(res)
            ds = acq["Sweep_%d" % num].create_dataset("initial_access_resistance", data=res)
            ds.attrs["neurodata_type"] = "Custom"
            ds = stim["Sweep_%d" % num].create_dataset("initial_access_resistance", data=res)
            ds.attrs["neurodata_type"] = "Custom"
            # 
#            # recycle code from old publish module for custom sweep metadata
#            if num in metadata['sweeps']:
#                sweep_md = metadata['sweeps'][num]
#                stimulus_interval = sweep_md['stimulus_interval']
#                if stimulus_interval is None:
#                    stimulus_interval = float('nan')
#                ds = acq["Sweep_%d" % num].create_dataset("aibs_stimulus_interval", data=stimulus_interval)
#                ds.attrs["neurodata_type"] = "Custom"
#                #
#                ds = acq["Sweep_%d" % num].create_dataset("aibs_stimulus_name", data=sweep_md['stimulus_type_name'])
#                ds.attrs["neurodata_type"] = "Custom"
#                #
#                ds = acq["Sweep_%d" % num].create_dataset("aibs_stimulus_amplitude_%s" % stim_units, sweep_md['stimulus_amplitude'])
#                ds.attrs["neurodata_type"] = "Custom"
#                #
#                ds = acq["Sweep_%d" % num].create_dataset("seal", sweep_md['seal_gohm'])
#                ds.attrs["neurodata_type"] = "Custom"
        except:
            print("json parse error for sweep %d" % num)
            raise
    # all done
    hdf.close()


    # TODO describe what's happening here
    sweeps_by_type = defaultdict(list)
    for sweep_number, sweep_data in metadata['sweeps'].iteritems():
        sweeps_by_type[sweep_data['stimulus_type_name']].append(sweep_number)

    sweep_features = extract_cell_features.extract_sweep_features(NwbDataSet(outfile), sweeps_by_type)

    # TODO describe what's happening here
    for sweep_num in passing_sweeps:
        try:
            spikes = sweep_features[sweep_num]['spikes']
            spike_times = [ s['threshold_t'] for s in spikes ]
            NwbDataSet(outfile).set_spike_times(sweep_num, spike_times)
        except Exception, e:
            logging.info("sweep %d has no sweep features. %s" % (sweep_num, e.message) )
#    try:
#        # remove spike times for non-passing sweeps
#        spk = hdf["analysis/spike_times"]
#        for grp in spk:
#            try:
#                num = int(str(grp).split('_')[-1])
#            except:
#                continue
#            if num not in passing_sweeps:
#                del spk[str(grp)]
#    except:
#        

#    # 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

    empty = {}
    return empty


if __name__ == "__main__": 
    # read module input. PipelineModule object automatically parses the 
    #   command line to pull out input.json and output.json file names
    module = PipelineModule()
    jin = module.input_data()   # loads input.json
    jout = main(jin)
    module.write_output_data(jout)  # writes output.json

