ALLEN INSTITUTE

# Source code for bmtk.simulator.bionet.default_setters.cell_models

    # Copyright 2017. Allen Institute. All rights reserved
    #
    # Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
    # following conditions are met:
    #
    # 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
    # disclaimer.
    #
    # 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
    # disclaimer in the documentation and/or other materials provided with the distribution.
    #
    # 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
    # products derived from this software without specific prior written permission.
    #
    # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
    # INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
    # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
    # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
    # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
    # WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
    # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
    #
    import os
    import numpy as np
    from neuron import h
    import inspect
    try:
        from sklearn.decomposition import PCA
    except Exception as e:
        pass

    from bmtk.simulator.bionet.pyfunction_cache import add_cell_model, add_cell_processor
    from bmtk.simulator.bionet.io_tools import io
    from bmtk.simulator.bionet.nml_reader import NMLTree

    """
    Functions for loading NEURON cell objects.

    Functions will be loaded by bionetwork and called when a new cell object is created. These are for standard models
    loaded with Cell-Types json files or their NeuroML equivelent, but may be overridden by the users.
    """

    def loadHOC(cell, template_name, dynamics_params):
        """A Generic function for creating a cell object from a NEURON HOC Template (eg. a *.hoc file with
        `begintemplate template_name` in header). It essentially tries to guess the correct parameters that need to be
        called so may not work the majority of the times and require to be overloaded.

        :param cell: A SONATA node object, can be used as a dict to get individual properties of current cell.
        :param template_name: name of HOCTemplate as stored in "model_template" attribute (hoc:<template_name>).
        :param dynamics_params: Dictionary containing contents of cell['dynamics_params'] as loaded from a json file or hdf5.
            If cell does not have "dynamics_params" attributes then will be set to None.
        """
        try:
            # Get template to instantiate
            template_call = getattr(h, template_name)
        except AttributeError as ae:
            io.log_error(
                f'loadHOC was unable to load in Neuron HOC Template "{template_name}, '
                'Make sure appropiate .hoc file is stored in templates_dir.'
            )
            raise ae

        try:
            if dynamics_params is not None and 'params' in dynamics_params:
                template_params = dynamics_params['params']
                if isinstance(template_params, list):
                    # pass in a list of parameters
                    hobj = template_call(*template_params)
                else:
                    # only a single parameter
                    hobj = template_call(template_params)
            elif cell.morphology_file is not None:
                # instantiate template with no parameters
                hobj = template_call(cell.morphology_file)
            else:
                hobj = template_call()
        except RuntimeError as rte:
            io.log_error(
                f'bmtk.simualtor.bionet.default_setters.cell_models.loadHOC function failed to load HOC template "{template_call}". '
                'If Hoc Templates requires special call to be initialized consider using `bmtk.simulator.bionet.add_cell_model()` '
                'to overwrite this function.'
            )
            raise rte

        # TODO: All "all" section if it doesn't exist
        # hobj.all = h.SectionList()
        # hobj.all.wholetree(sec=hobj.soma[0])
        return hobj

    def IntFire1(cell, template_name, dynamics_params):
        """Loads a point integrate and fire neuron"""
        hobj = h.IntFire1()
        hobj.tau = dynamics_params['tau']*1000.0  # Convert from seconds to ms.
        hobj.refrac = dynamics_params['refrac']*1000.0  # Convert from seconds to ms.
        return hobj

    def Biophys1(cell, template_name, dynamic_params):
        """Loads a biophysical NEURON hoc object using Cell-Types database objects."""
        morphology_file = cell.morphology_file
        hobj = h.Biophys1(str(morphology_file))
        #fix_axon(hobj)
        #set_params_peri(hobj, dynamic_params)
        return hobj

    def Biophys1_nml(json_file):
        # TODO: look at examples to see how to convert .nml files
        raise NotImplementedError()

    def Biophys1_dict(cell):
        """ Set parameters for cells from the Allen Cell Types database Prior to setting parameters will replace the
        axon with the stub
        """
        morphology_file = cell['morphology']
        hobj = h.Biophys1(str(morphology_file))
        return hobj

    def aibs_perisomatic(hobj, cell, dynamics_params):
        if dynamics_params is not None:
            fix_axon_peri(hobj)
            set_params_peri(hobj, dynamics_params)

        return hobj

    def fix_axon_peri(hobj):
        """Replace reconstructed axon with a stub

        :param hobj: hoc object
        """
        for sec in hobj.axon:
            h.delete_section(sec=sec)

        h.execute('create axon[2]', hobj)

        for sec in hobj.axon:
            sec.L = 30
            sec.diam = 1
            hobj.axonal.append(sec=sec)
            hobj.all.append(sec=sec)  # need to remove this comment

        hobj.axon[0].connect(hobj.soma[0], 0.5, 0)
        hobj.axon[1].connect(hobj.axon[0], 1, 0)

        h.define_shape()

    def set_params_peri(hobj, biophys_params):
        """Set biophysical parameters for the cell

        :param hobj: NEURON's cell object
        :param biophys_params: name of json file with biophys params for cell's model which determine spiking behavior
        :return:
        """
        passive = biophys_params['passive'][0]
        conditions = biophys_params['conditions'][0]
        genome = biophys_params['genome']

        # Set passive properties
        cm_dict = dict([(c['section'], c['cm']) for c in passive['cm']])
        for sec in hobj.all:
            sec.Ra = passive['ra']
            sec.cm = cm_dict[sec.name().split(".")[1][:4]]
            sec.insert('pas')

            for seg in sec:
                seg.pas.e = passive["e_pas"]

        # Insert channels and set parameters
        for p in genome:
            sections = [s for s in hobj.all if s.name().split(".")[1][:4] == p["section"]]

            for sec in sections:
                if p["mechanism"] != "":
                    sec.insert(p["mechanism"])
                setattr(sec, p["name"], p["value"])

        # Set reversal potentials
        for erev in conditions['erev']:
            sections = [s for s in hobj.all if s.name().split(".")[1][:4] == erev["section"]]
            for sec in sections:
                sec.ena = erev["ena"]
                sec.ek = erev["ek"]

    def aibs_allactive(hobj, cell, dynamics_params):
        fix_axon_allactive(hobj)
        set_params_allactive(hobj, dynamics_params)
        return hobj

    def fix_axon_allactive(hobj):
        """Replace reconstructed axon with a stub

        Parameters
        ----------
        hobj: instance of a Biophysical template
            NEURON's cell object
        """
        # find the start and end diameter of the original axon, this is different from the perisomatic cell model
        # where diameter == 1.
        axon_diams = [hobj.axon[0].diam, hobj.axon[0].diam]
        for sec in hobj.all:
            section_name = sec.name().split(".")[1][:4]
            if section_name == 'axon':
                axon_diams[1] = sec.diam

        for sec in hobj.axon:
            h.delete_section(sec=sec)

        h.execute('create axon[2]', hobj)
        for index, sec in enumerate(hobj.axon):
            sec.L = 30
            sec.diam = axon_diams[index]  # 1

            hobj.axonal.append(sec=sec)
            hobj.all.append(sec=sec)  # need to remove this comment

        hobj.axon[0].connect(hobj.soma[0], 1.0, 0)
        hobj.axon[1].connect(hobj.axon[0], 1.0, 0)

        h.define_shape()

    def set_params_allactive(hobj, params_dict):
        # params_dict = json.load(open(params_file_name, 'r'))
        passive = params_dict['passive'][0]
        genome = params_dict['genome']
        conditions = params_dict['conditions'][0]

        section_map = {}
        for sec in hobj.all:
            section_name = sec.name().split(".")[1][:4]
            if section_name in section_map:
                section_map[section_name].append(sec)
            else:
                section_map[section_name] = [sec]

        for sec in hobj.all:
            sec.insert('pas')
            # sec.insert('extracellular')

        if 'e_pas' in passive:
            e_pas_val = passive['e_pas']
            for sec in hobj.all:
                for seg in sec:
                    seg.pas.e = e_pas_val

        if 'ra' in passive:
            ra_val = passive['ra']
            for sec in hobj.all:
                sec.Ra = ra_val

        if 'cm' in passive:
            # print('Setting cm')
            for cm_dict in passive['cm']:
                cm = cm_dict['cm']
                for sec in section_map.get(cm_dict['section'], []):
                    sec.cm = cm

        for genome_dict in genome:
            g_section = genome_dict['section']
            if genome_dict['section'] == 'glob':
                io.log_warning("There is a section called glob, probably old json file")
                continue

            g_value = float(genome_dict['value'])
            g_name = genome_dict['name']
            g_mechanism = genome_dict.get("mechanism", "")
            for sec in section_map.get(g_section, []):
                if g_mechanism != "":
                    sec.insert(g_mechanism)
                setattr(sec, g_name, g_value)

        for erev in conditions['erev']:
            erev_section = erev['section']
            erev_ena = erev['ena']
            erev_ek = erev['ek']

            if erev_section in section_map:
                for sec in section_map.get(erev_section, []):
                    if h.ismembrane('k_ion', sec=sec) == 1:
                        setattr(sec, 'ek', erev_ek)
                    if h.ismembrane('na_ion', sec=sec) == 1:
                        setattr(sec, 'ena', erev_ena)
            else:
                io.log_warning("Can't set erev for {}, section array doesn't exist".format(erev_section))

    def aibs_perisomatic_directed(hobj, cell, dynamics_params):
        fix_axon_perisomatic_directed(hobj)
        set_params_peri(hobj, dynamics_params)
        return hobj

    def aibs_allactive_directed(hobj, cell, dynamics_params):
        fix_axon_allactive_directed(hobj)
        set_params_allactive(hobj, dynamics_params)
        return hobj

    def fix_axon_perisomatic_directed(hobj):
        # io.log_info('Fixing Axon like perisomatic')
        all_sec_names = []
        for sec in hobj.all:
            all_sec_names.append(sec.name().split(".")[1][:4])

        if 'axon' not in all_sec_names:
            io.log_exception('There is no axonal recostruction in swc file.')
        else:
            beg1, end1, beg2, end2 = get_axon_direction(hobj)

        for sec in hobj.axon:
            h.delete_section(sec=sec)
        h.execute('create axon[2]', hobj)

        h.pt3dadd(beg1[0], beg1[1], beg1[2], 1, sec=hobj.axon[0])
        h.pt3dadd(end1[0], end1[1], end1[2], 1, sec=hobj.axon[0])
        hobj.all.append(sec=hobj.axon[0])
        h.pt3dadd(beg2[0], beg2[1], beg2[2], 1, sec=hobj.axon[1])
        h.pt3dadd(end2[0], end2[1], end2[2], 1, sec=hobj.axon[1])
        hobj.all.append(sec=hobj.axon[1])

        hobj.axon[0].connect(hobj.soma[0], 0.5, 0)
        hobj.axon[1].connect(hobj.axon[0], 1.0, 0)

        hobj.axon[0].L = 30.0
        hobj.axon[1].L = 30.0

        h.define_shape()

        for sec in hobj.axon:
            # print "sec.L:", sec.L
            if np.abs(30-sec.L) > 0.0001:
                io.log_exception('Axon stub L is less than 30')

    def fix_axon_allactive_directed(hobj):
        all_sec_names = []
        for sec in hobj.all:
            all_sec_names.append(sec.name().split(".")[1][:4])

        if 'axon' not in all_sec_names:
            io.log_exception('There is no axonal recostruction in swc file.')
        else:
            beg1, end1, beg2, end2 = get_axon_direction(hobj)

        axon_diams = [hobj.axon[0].diam, hobj.axon[0].diam]
        for sec in hobj.all:
            section_name = sec.name().split(".")[1][:4]
            if section_name == 'axon':
                axon_diams[1] = sec.diam

        for sec in hobj.axon:
            h.delete_section(sec=sec)
        h.execute('create axon[2]', hobj)
        hobj.axon[0].connect(hobj.soma[0], 1.0, 0)
        hobj.axon[1].connect(hobj.axon[0], 1.0, 0)

        h.pt3dadd(beg1[0], beg1[1], beg1[2], axon_diams[0], sec=hobj.axon[0])
        h.pt3dadd(end1[0], end1[1], end1[2], axon_diams[0], sec=hobj.axon[0])
        hobj.all.append(sec=hobj.axon[0])
        h.pt3dadd(beg2[0], beg2[1], beg2[2], axon_diams[1], sec=hobj.axon[1])
        h.pt3dadd(end2[0], end2[1], end2[2], axon_diams[1], sec=hobj.axon[1])
        hobj.all.append(sec=hobj.axon[1])

        hobj.axon[0].L = 30.0
        hobj.axon[1].L = 30.0

        h.define_shape()

        for sec in hobj.axon:
            # io.log_info('sec.L: {}'.format(sec.L))
            if np.abs(30 - sec.L) > 0.0001:
                io.log_exception('Axon stub L is less than 30')

    def get_axon_direction(hobj):
        for sec in hobj.somatic:
            n3d = int(h.n3d(sec=sec))  # get number of n3d points in each section
            soma_end = np.asarray([h.x3d(n3d - 1, sec=sec), h.y3d(n3d - 1, sec=sec), h.z3d(n3d - 1, sec=sec)])
            mid_point = int(n3d / 2)
            soma_mid = np.asarray([h.x3d(mid_point, sec=sec), h.y3d(mid_point, sec=sec), h.z3d(mid_point, sec=sec)])

        for sec in hobj.all:
            section_name = sec.name().split(".")[1][:4]
            if section_name == 'axon':
                n3d = int(h.n3d(sec=sec))  # get number of n3d points in each section
                axon_p3d = np.zeros((n3d, 3))  # to hold locations of 3D morphology for the current section
                for i in range(n3d):
                    axon_p3d[i, 0] = h.x3d(i, sec=sec)
                    axon_p3d[i, 1] = h.y3d(i, sec=sec)  # shift coordinates such to place soma at the origin.
                    axon_p3d[i, 2] = h.z3d(i, sec=sec)

        # Add soma coordinates to the list
        p3d = np.concatenate(([soma_mid], axon_p3d), axis=0)

        # Compute PCA
        pca = PCA(n_components=3)
        pca.fit(p3d)
        unit_v = pca.components_[0]

        mag_v = np.sqrt(pow(unit_v[0], 2) + pow(unit_v[1], 2) + pow(unit_v[2], 2))
        unit_v[0] = unit_v[0] / mag_v
        unit_v[1] = unit_v[1] / mag_v
        unit_v[2] = unit_v[2] / mag_v

        # Find the direction
        axon_end = axon_p3d[-1] - soma_mid
        if np.dot(unit_v, axon_end) < 0:
            unit_v *= -1

        axon_seg_coor = np.zeros((4, 3))
        # unit_v = np.asarray([0,1,0])
        axon_seg_coor[0] = soma_end
        axon_seg_coor[1] = soma_end + (unit_v * 30.)
        axon_seg_coor[2] = soma_end + (unit_v * 30.)
        axon_seg_coor[3] = soma_end + (unit_v * 60.)

        return axon_seg_coor

    nml_files = {}  # For caching neuroml file trees

    def NMLLoad(cell, template_name, dynamic_params):
        """Convert a NEUROML file to a NEURON hoc cell object.

        Current limitations:
        * Ignores nml morphology section. You must pass in a swc file
        * Only for biophysically detailed cell biophysical components. All properties must be assigned to a segment group.

        :param cell:
        :param template_name:
        :param dynamic_params:
        :return:
        """
        # Last I checked there is no built in way to load a NML file directly into NEURON through the API, instead we have
        # to manually parse the nml file and build the NEUROM cell object section-by-section.
        morphology_file = cell.morphology_file
        hobj = h.Biophys1(str(morphology_file))
        # Depending on if the axon is cut before or after setting cell channels and mechanism can create drastically
        # different results. Currently NML files doesn't produce the same results if you use model_processing directives.
        # TODO: Find a way to specify model_processing directive with NML file
        fix_axon_peri(hobj)

        # Load the hoc template containing a swc initialized NEURON cell
        if template_name in nml_files:
            nml_params = nml_files[template_name]
        else:
            # Parse the NML parameters file xml tree and cache.
            biophys_dirs = cell.network.get_component('biophysical_neuron_models_dir')
            nml_path = os.path.join(biophys_dirs, template_name)
            nml_params = NMLTree(nml_path)
            nml_files[template_name] = nml_params

        # Iterate through the NML tree by section and use the properties to manually create cell mechanisms
        section_lists = [(sec, sec.name().split(".")[1][:4]) for sec in hobj.all]
        for sec, sec_name in section_lists:
            for prop_name, prop_obj in nml_params[sec_name].items():
                if prop_obj.element_tag() == 'resistivity':
                    sec.Ra = prop_obj.value

                elif prop_obj.element_tag() == 'specificCapacitance':
                    sec.cm = prop_obj.value

                elif prop_obj.element_tag() == 'channelDensity' and prop_obj.ion_channel == 'pas':
                    sec.insert('pas')
                    setattr(sec, 'g_pas', prop_obj.cond_density)
                    for seg in sec:
                        seg.pas.e = prop_obj.erev

                elif prop_obj.element_tag() == 'channelDensity' or prop_obj.element_tag() == 'channelDensityNernst':
                    sec.insert(prop_obj.ion_channel)
                    setattr(sec, prop_obj.id, prop_obj.cond_density)
                    if prop_obj.ion == 'na' and prop_obj:
                        sec.ena = prop_obj.erev
                    elif prop_obj.ion == 'k':
                        sec.ek = prop_obj.erev

                elif prop_obj.element_tag() == 'concentrationModel':
                    sec.insert(prop_obj.id)
                    setattr(sec, 'gamma_' + prop_obj.type, prop_obj.gamma)
                    setattr(sec, 'decay_' + prop_obj.type, prop_obj.decay)

        return hobj

    def set_extracellular(hobj, cell, dynamics_params):
        for sec in hobj.all:
            sec.insert('extracellular')

        return hobj

    add_cell_model(loadHOC, directive='hoc', model_type='biophysical')
    add_cell_model(NMLLoad, directive='nml', model_type='biophysical')
    add_cell_model(Biophys1, directive='ctdb:Biophys1', model_type='biophysical', overwrite=False)
    add_cell_model(Biophys1, directive='ctdb:Biophys1.hoc', model_type='biophysical', overwrite=False)
    add_cell_model(IntFire1, directive='nrn:IntFire1', model_type='point_process', overwrite=False)
    add_cell_model(IntFire1, directive='nrn:IntFire1', model_type='point_neuron', overwrite=False)

    add_cell_processor(aibs_perisomatic, overwrite=False)
    add_cell_processor(aibs_allactive, overwrite=False)
    add_cell_processor(aibs_perisomatic_directed, overwrite=False)
    add_cell_processor(aibs_allactive_directed, overwrite=False)
    add_cell_processor(set_extracellular, overwrite=False)
    add_cell_processor(set_extracellular, 'extracellular', overwrite=False)