import h5py
import math

class LabNotebookReader(object):
    def __init__(self):
        self.register_enabled_names()

    # mapping of notebook keys to keys representing if that value is
    #   enabled
    # move this to subclasses if/when key names diverge
    def register_enabled_names(self):
        self.enabled = {}
        self.enabled["V-Clamp Holding Level"] = "V-Clamp Holding Enable"
        self.enabled["RsComp Bandwidth"] = "RsComp Enable"
        self.enabled["RsComp Correction"] = "RsComp Enable"
        self.enabled["RsComp Prediction"] = "RsComp Enable"
        self.enabled["Whole Cell Comp Cap"] = "Whole Cell Comp Enable"
        self.enabled["Whole Cell Comp Resist"] = "Whole Cell Comp Enable"
        self.enabled["I-Clamp Holding Level"] = "I-Clamp Holding Enable"
        self.enabled["Neut Cap Value"] = "Neut Cap Enable"
        self.enabled["Bridge Bal Value"] = "Bridge Bal Enable"


    # lab notebook has two sections, one for numeric data and the other
    #   for text data. this is an internal function to fetch data from
    #   the numeric part of the notebook
    def get_numeric_value(self, name, data_col, sweep_col, enable_col, sweep_num, default_val):
        data = self.val_number
        # val_number has 3 dimensions -- the first has a shape of
        #   (#fields * 9). there are many hundreds of elements in this
        #   dimension. they look to represent the full array of values
        #   (for each field for each multipatch) for a given point in
        #   time, and thus given sweep
        # according to Thomas Braun (igor nwb dev), the first 8 pages are
        #   for headstage data, and the 9th is for headstage-independent
        #   data
        # return value is last non-empty entry in specified column
        #   for specified sweep number
        return_val = default_val
        for sample in data:
            swp = sample[sweep_col][0]
            if math.isnan(swp):
                continue
            if int(swp) == sweep_num:
                if enable_col is not None and sample[enable_col][0] != 1.0:
                    continue # 'enable' flag present and it's turned off
                val = sample[data_col][0]
                if not math.isnan(val):
                    return_val = val
        return return_val

    # internal function for fetching data from the text part of the notebook
    def get_text_value(self, name, data_col, sweep_col, enable_col, sweep_num, default_val):
        data = self.val_text
        # algorithm mirrors get_numeric_value
        # return value is last non-empty entry in specified column
        #   for specified sweep number
        return_val = default_val
        for sample in data:
            swp = sample[sweep_col][0]
            if len(swp) == 0:
                continue
            if int(swp) == int(sweep_num):
                if enable_col is not None: # and sample[enable_col][0] != 1.0:
                    # this shouldn't happen, but if it does then bitch
                    #   as this situation hasn't been tested (eg, is 
                    #   enabled indicated by 1.0, or "1.0" or "true" or ??)
                    Exception("Enable flag not expected for text values")
                    #continue # 'enable' flag present and it's turned off
                val = sample[data_col][0]
                if len(val) > 0:
                    return_val = val
        return return_val

    # looks for key in lab notebook and returns the value associated with
    #   the specified sweep, or the default value if no value is found
    #   (NaN and empty strings are considered to be non-values)
    def get_value(self, name, sweep_num, default_val):
        # name_number has 3 dimensions -- the first has shape
        #   (#fields * 9) and stores the key names. the second looks
        #   to store units for those keys. The third is numeric text
        #   but it's role isn't clear
        numeric_fields = self.colname_number[0]
        text_fields = self.colname_text[0]
        # val_number has 3 dimensions -- the first has a shape of
        #   (#fields * 9). there are many hundreds of elements in this
        #   dimension. they look to represent the full array of values
        #   (for each field for each multipatch) for a given point in
        #   time, and thus given sweep
        if name in numeric_fields:
            sweep_idx = numeric_fields.tolist().index("SweepNum")
            enable_idx = None
            if name in self.enabled:
                enable_col = self.enabled[name]
                enable_idx = numeric_fields.tolist().index(enable_col)
            field_idx = numeric_fields.tolist().index(name)
            return self.get_numeric_value(name, field_idx, sweep_idx, enable_idx, sweep_num, default_val)
        elif name in text_fields:
            # first check to see if file includes old version of column name
            if "Sweep #" in text_fields:
                sweep_idx = text_fields.tolist().index("Sweep #")
            else:
                sweep_idx = text_fields.tolist().index("SweepNum")
            enable_idx = None
            if name in self.enabled:
                enable_col = self.enabled[name]
                enable_idx = text_fields.tolist().index(enable_col)
            field_idx = text_fields.tolist().index(name)
            return self.get_text_value(name, field_idx, sweep_idx, enable_idx, sweep_num, default_val)
        else:
            return default_val
            
        

""" Loads lab notebook data out of a first-generation IVSCC NWB file,
    that was manually translated from the IGOR h5 dump.
    Notebook data can be read through get_value() function
"""
class LabNotebookReaderIvscc(LabNotebookReader):
    def __init__(self, nwb_file, h5_file):
        LabNotebookReader.__init__(self)
        # for lab notebook, select first group 
        h5 = h5py.File(h5_file, "r")
        #
        # TODO FIXME check notebook version... but how?
        #
        notebook = h5["MIES/LabNoteBook/ITC18USB/Device0"]
        # load column data into memory
        self.colname_number = notebook["KeyWave/keyWave"].value
        self.val_number = notebook["settingsHistory/settingsHistory"].value
        self.colname_text = notebook["TextDocKeyWave/txtDocKeyWave"].value
        self.val_text = notebook["textDocumentation/txtDocWave"].value
        h5.close()



########################################################################
########################################################################
""" Loads lab notebook data out of an Igor-generated NWB file.
    Module input is the name of the nwb file.
    Notebook data can be read through get_value() function
"""
class LabNotebookReaderIgorNwb(LabNotebookReader):
    def __init__(self, nwb_file):
        LabNotebookReader.__init__(self)
        # for lab notebook, select first group 
        # NOTE this probably won't work for multipatch
        h5 = h5py.File(nwb_file, "r")
        #
        # TODO FIXME check notebook version
        #
        for k in h5["general/labnotebook"]:
            notebook = h5["general/labnotebook"][k]
            break
        # load column data into memory
        self.val_text = notebook["textualValues"].value
        self.colname_text = notebook["textualKeys"].value
        self.val_number = notebook["numericalValues"].value
        self.colname_number = notebook["numericalKeys"].value
        h5.close()
        #
        self.register_enabled_names()

        

# creates LabNotebookReader appropriate to ivscc-NWB file version
def create_lab_notebook_reader(nwb_file, h5_file=None):
    pass
    h5 = h5py.File(nwb_file, "r")
    if "general/labnotebook" in h5:
        version = "IgorNwb"
    else:
        version = "IgorH5"
    h5.close()
    if version == "IgorNwb":
        return LabNotebookReaderIgorNwb(nwb_file)
    elif version == "IgorH5":
        return LabNotebookReaderIvscc(nwb_file, h5_file)
    else:
        Exception("Unable to determine NWB input type")

