ALLEN INSTITUTE

# Source code for bmtk.simulator.bionet.sonata_adaptors

    # 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 types
    import numpy as np

    from bmtk.simulator.core.sonata_reader import NodeAdaptor, SonataBaseNode, EdgeAdaptor, SonataBaseEdge
    from bmtk.simulator.bionet import nrn

    class BioNode(SonataBaseNode):
        @property
        def position(self):
            return self._prop_adaptor.position(self._node)

        @property
        def morphology_file(self):
            if 'morphology' not in self._node:
                return None
            return self._node['morphology']

        @property
        def rotation_angle_xaxis(self):
            return self._prop_adaptor.rotation_angle_xaxis(self._node)

        @property
        def rotation_angle_yaxis(self):
            # TODO: Combine rotation alnges into a single property
            return self._prop_adaptor.rotation_angle_yaxis(self._node)

        @property
        def rotation_angle_zaxis(self):
            return self._prop_adaptor.rotation_angle_zaxis(self._node)

        @property
        def rotations(self):
            return self._prop_adaptor.rotations(self._node)

        @property
        def rotations_quaternion(self):
            return self._prop_adaptor.rotations(self._node)

        def load_cell(self):
            model_template = self.model_template
            template_name = model_template[1]
            model_type = self.model_type
            if nrn.py_modules.has_cell_model(self['model_template'], model_type):
                cell_fnc = nrn.py_modules.cell_model(self['model_template'], model_type)
            else:
                cell_fnc = nrn.py_modules.cell_model(model_template[0], model_type)

            dynamics_params = self.dynamics_params
            hobj = cell_fnc(self, template_name, dynamics_params)

            for model_processing_str in self.model_processing:
                processing_fnc = nrn.py_modules.cell_processor(model_processing_str)
                hobj = processing_fnc(hobj, self, dynamics_params)

            return hobj

    class BioNodeAdaptor(NodeAdaptor):

        def get_node(self, sonata_node):
            return BioNode(sonata_node, self)

        @classmethod
        def patch_adaptor(cls, adaptor, node_group, network):
            node_adaptor = NodeAdaptor.patch_adaptor(adaptor, node_group, network)

            # Position
            if 'positions' in node_group.all_columns:
                node_adaptor.position = types.MethodType(positions, adaptor)
            elif 'position' in node_group.all_columns:
                node_adaptor.position = types.MethodType(position, adaptor)
            elif 'x' in node_group.all_columns:
                if 'z' in node_group.all_columns and 'y' in node_group.all_columns:
                    node_adaptor.position = types.MethodType(position_xyz, adaptor)
                elif 'y' in node_group.all_columns:
                    node_adaptor.position = types.MethodType(position_xy, adaptor)
                else:
                    node_adaptor.position = types.MethodType(position_x, adaptor)
            else:
                node_adaptor.position = types.MethodType(positions_default, adaptor)

            # Rotation angles
            if 'rotation_angle_xaxis' in node_group.all_columns:
                node_adaptor.rotation_angle_xaxis = types.MethodType(rotation_angle_x, node_adaptor)
            else:
                node_adaptor.rotation_angle_xaxis = types.MethodType(rotation_angle_default, node_adaptor)

            if 'rotation_angle_yaxis' in node_group.all_columns:
                node_adaptor.rotation_angle_yaxis = types.MethodType(rotation_angle_y, node_adaptor)
            else:
                node_adaptor.rotation_angle_yaxis = types.MethodType(rotation_angle_default, node_adaptor)

            if 'rotation_angle_zaxis' in node_group.all_columns:
                node_adaptor.rotation_angle_zaxis = types.MethodType(rotation_angle_z, node_adaptor)
            else:
                node_adaptor.rotation_angle_zaxis = types.MethodType(rotation_angle_default, node_adaptor)

            if 'rotations' in node_group.all_columns:
                node_adaptor.rotations = types.MethodType(rotations, node_adaptor)
            else:
                node_adaptor.rotations = types.MethodType(value_none, node_adaptor)

            return node_adaptor

    def rotations(self, node):
        return node['rotations']

    def value_none(self, node):
        return None

    def positions_default(self, node):
        return np.array([0.0, 0.0, 0.0])

    def positions(self, node):
        return node['positions']

    def position(self, node):
        return node['position']

    def rotation_angle_default(self, node):
        return 0.0

    def rotation_angle_x(self, node):
        return node['rotation_angle_xaxis']

    def rotation_angle_y(self, node):
        return node['rotation_angle_yaxis']

    def rotation_angle_z(self, node):
        return node['rotation_angle_zaxis']

    def position_xyz(self, node):
        return np.array([node['x'], node['y'], node['z']])

    def position_xy(self, node):
        return np.array([node['x'], node['y'], 0.0])

    def position_x(self, node):
        return np.array([node['x'], 0.0, 0.0])

    class BioEdge(SonataBaseEdge):

        def load_synapses(self, section_x, section_id):
            synapse_fnc = nrn.py_modules.synapse_model(self.model_template)
            return synapse_fnc(self.dynamics_params, section_x, section_id)

        @property
        def afferent_section_id(self):
            return self._prop_adaptor.aff_sec_id(self._edge)

        @property
        def afferent_section_pos(self):
            return self._prop_adaptor.aff_sec_pos(self._edge)

    class BioEdgeAdaptor(EdgeAdaptor):

        def get_edge(self, sonata_edge):
            return BioEdge(sonata_edge, self)

        @classmethod
        def patch_adaptor(cls, adaptor, edge_group):
            edge_adaptor = EdgeAdaptor.patch_adaptor(adaptor, edge_group)

            if 'sec_id' in edge_group.all_columns:
                edge_adaptor.aff_sec_id = types.MethodType(aff_sec_id_old, adaptor)
            elif 'afferent_section_id' in edge_group.all_columns:
                edge_adaptor.aff_sec_id = types.MethodType(aff_sec_id, adaptor)
            else:
                edge_adaptor.aff_sec_id = types.MethodType(ret_none_function, adaptor)

            if 'sec_x' in edge_group.all_columns:
                edge_adaptor.aff_sec_pos = types.MethodType(aff_sec_pos_old, adaptor)
            elif 'afferent_section_pos' in edge_group.all_columns:
                edge_adaptor.aff_sec_pos = types.MethodType(aff_sec_pos, adaptor)
            else:
                edge_adaptor.aff_sec_pos = types.MethodType(ret_none_function, adaptor)

            return edge_adaptor

    def ret_none_function(self, edge):
        return None

    def aff_sec_id(self, edge):
        return edge['afferent_section_id']

    def aff_sec_id_old(self, edge):
        return edge['sec_id']

    def aff_sec_pos(self, edge):
        return edge['afferent_section_pos']

    def aff_sec_pos_old(self, edge):
        return edge['sec_x']