"""
Compartmental models.
This module defines the `SpatialNeuron` class, which defines multicompartmental
models.
"""
import weakref
import copy

import sympy as sp
import numpy as np

from brian2.core.variables import Variables
from brian2.equations.equations import (Equations, PARAMETER, SUBEXPRESSION,
                                        DIFFERENTIAL_EQUATION, SingleEquation,
                                        extract_constant_subexpressions)
from brian2.groups.group import Group, CodeRunner, create_runner_codeobj
from brian2.units.allunits import ohm, siemens, amp, meter, volt
from brian2.units.fundamentalunits import (Quantity,
                                           fail_for_dimension_mismatch,
                                           have_same_dimensions,
                                           DimensionMismatchError)
from brian2.units.stdunits import uF, cm
from brian2.parsing.sympytools import sympy_to_str, str_to_sympy
from brian2.utils.logger import get_logger
from brian2.groups.neurongroup import (NeuronGroup, SubexpressionUpdater,
                                       to_start_stop)
from brian2.groups.subgroup import Subgroup
from brian2.equations.codestrings import Expression

__all__ = ['SpatialNeuron']

logger = get_logger(__name__)


class FlatMorphology(object):
    """
    Container object to store the flattened representation of a morphology.
    Note that all values are stored as numpy arrays without unit information
    (i.e. in base units).
    """
    def __init__(self, morphology):
        self.n = n = morphology.total_compartments  # Total number of compartments
        # Per-compartment attributes
        self.length = np.zeros(n)
        self.distance = np.zeros(n)
        self.area = np.zeros(n)
        self.diameter = np.zeros(n)
        self.volume = np.zeros(n)
        self.r_length_1 = np.zeros(n)
        self.r_length_2 = np.zeros(n)
        self.start_x = np.zeros(n)
        self.start_y = np.zeros(n)
        self.start_z = np.zeros(n)
        self.x = np.zeros(n)
        self.y = np.zeros(n)
        self.z = np.zeros(n)
        self.end_x = np.zeros(n)
        self.end_y = np.zeros(n)
        self.end_z = np.zeros(n)
        self.depth = np.zeros(n, dtype=np.int32)
        self.sections = sections = morphology.total_sections
        self.end_distance = np.zeros(sections)
        # Index of the parent for each section (-1 for the root)
        self.morph_parent_i = np.zeros(sections, dtype=np.int32)
        # The children indices for each section (list of lists, will be later
        # transformed into an array representation)
        self.morph_children = []
        # each section is child of exactly one parent, this stores the index in
        # the parents list of children
        self.morph_idxchild = np.zeros(sections, dtype=np.int32)
        self.starts = np.zeros(sections, dtype=np.int32)
        self.ends = np.zeros(sections, dtype=np.int32)

        # recursively fill the data structures
        self._sections_without_coordinates = False
        self.has_coordinates = False
        self._offset = 0
        self._section_counter = 0
        self._insert_data(morphology)
        if self.has_coordinates and self._sections_without_coordinates:
            logger.info('The morphology has a mix of sections with and '
                        'without coordinates. The SpatialNeuron object '
                        'will store NaN values for the coordinates of '
                        'the sections that do not specify coordinates. '
                        'Call generate_coordinates on the morphology '
                        'before creating the SpatialNeuron object to fill '
                        'in the missing coordinates.')
        # Do not store coordinates for morphologies that don't define them
        if not self.has_coordinates:
            self.start_x = self.start_y = self.start_z = None
            self.x = self.y = self.z = None
            self.end_x = self.end_y = self.end_z = None

        # Transform the list of list of children into a 2D array (stored as
        # 1D) -- note that this wastes space if the number of children per
        # section is very different. In practice, this should not be much of a
        # problem since most sections have 0, 1, or 2 children (e.g. SWC files
        # on neuromorpho.org are all binary trees)
        self.morph_children_num = np.array([len(c)
                                            for c in self.morph_children] + [0])
        max_children = max(self.morph_children_num)
        morph_children = np.zeros((sections+1, max_children), dtype=np.int32)
        for idx, section_children in enumerate(self.morph_children):
            morph_children[idx, :len(section_children)] = section_children
        self.morph_children = morph_children.reshape(-1)

    def _insert_data(self, section, parent_idx=-1, depth=0):
        n = section.n
        start = self._offset
        end = self._offset + n
        # Compartment attributes
        self.depth[start:end] = depth
        self.length[start:end] = np.asarray(section.length)
        self.distance[start:end] = np.asarray(section.distance)
        self.area[start:end] = np.asarray(section.area)
        self.diameter[start:end] = np.asarray(section.diameter)
        self.volume[start:end] = np.asarray(section.volume)
        self.r_length_1[start:end] = np.asarray(section.r_length_1)
        self.r_length_2[start:end] = np.asarray(section.r_length_2)
        if section.x is None:
            self._sections_without_coordinates = True
            self.start_x[start:end] = np.ones(n)*np.nan
            self.start_y[start:end] = np.ones(n)*np.nan
            self.start_z[start:end] = np.ones(n)*np.nan
            self.x[start:end] = np.ones(n)*np.nan
            self.y[start:end] = np.ones(n)*np.nan
            self.z[start:end] = np.ones(n)*np.nan
            self.end_x[start:end] = np.ones(n)*np.nan
            self.end_y[start:end] = np.ones(n)*np.nan
            self.end_z[start:end] = np.ones(n)*np.nan
        else:
            self.has_coordinates = True
            self.start_x[start:end] = np.asarray(section.start_x)
            self.start_y[start:end] = np.asarray(section.start_y)
            self.start_z[start:end] = np.asarray(section.start_z)
            self.x[start:end] = np.asarray(section.x)
            self.y[start:end] = np.asarray(section.y)
            self.z[start:end] = np.asarray(section.z)
            self.end_x[start:end] = np.asarray(section.end_x)
            self.end_y[start:end] = np.asarray(section.end_y)
            self.end_z[start:end] = np.asarray(section.end_z)

        # Section attributes
        idx = self._section_counter
        # We start counting from 1 for the parent indices, since the index 0
        # is used for the (virtual) root compartment
        self.morph_parent_i[idx] = parent_idx + 1
        self.morph_children.append([])
        self.starts[idx] = start
        self.ends[idx] = end
        # Append ourselves to the children list of our parent
        self.morph_idxchild[idx] = len(self.morph_children[parent_idx+1])
        self.morph_children[parent_idx + 1].append(idx + 1)
        self.end_distance[idx] = section.end_distance
        # Recurse down the tree
        self._offset += n
        self._section_counter += 1
        for child in section.children:
            self._insert_data(child, parent_idx=idx, depth=depth+1)


class SpatialNeuron(NeuronGroup):
    """
    A single neuron with a morphology and possibly many compartments.

    Parameters
    ----------
    morphology : `Morphology`
        The morphology of the neuron.
    model : str, `Equations`
        The equations defining the group.
    method : str, function, optional
        The numerical integration method. Either a string with the name of a
        registered method (e.g. "euler") or a function that receives an
        `Equations` object and returns the corresponding abstract code. If no
        method is specified, a suitable method will be chosen automatically.
    threshold : str, optional
        The condition which produces spikes. Should be a single line boolean
        expression.
    threshold_location : (int, `Morphology`), optional
        Compartment where the threshold condition applies, specified as an
        integer (compartment index) or a `Morphology` object corresponding to
        the compartment (e.g. ``morpho.axon[10*um]``).
        If unspecified, the threshold condition applies at all compartments.
    Cm : `Quantity`, optional
        Specific capacitance in uF/cm**2 (default 0.9). It can be accessed and
        modified later as a state variable. In particular, its value can differ
        in different compartments.
    Ri : `Quantity`, optional
        Intracellular resistivity in ohm.cm (default 150). It can be accessed
        as a shared state variable, but modified only before the first run.
        It is uniform across the neuron.
    reset : str, optional
        The (possibly multi-line) string with the code to execute on reset.
    events : dict, optional
        User-defined events in addition to the "spike" event defined by the
        ``threshold``. Has to be a mapping of strings (the event name) to
        strings (the condition) that will be checked.
    refractory : {str, `Quantity`}, optional
        Either the length of the refractory period (e.g. ``2*ms``), a string
        expression that evaluates to the length of the refractory period
        after each spike (e.g. ``'(1 + rand())*ms'``), or a string expression
        evaluating to a boolean value, given the condition under which the
        neuron stays refractory after a spike (e.g. ``'v > -20*mV'``)
    namespace : dict, optional
        A dictionary mapping identifier names to objects. If not given, the
        namespace will be filled in at the time of the call of `Network.run`,
        with either the values from the ``namespace`` argument of the
        `Network.run` method or from the local context, if no such argument is
        given.
    dtype : (`dtype`, `dict`), optional
        The `numpy.dtype` that will be used to store the values, or a
        dictionary specifying the type for variable names. If a value is not
        provided for a variable (or no value is provided at all), the preference
        setting `core.default_float_dtype` is used.
    dt : `Quantity`, optional
        The time step to be used for the simulation. Cannot be combined with
        the `clock` argument.
    clock : `Clock`, optional
        The update clock to be used. If neither a clock, nor the `dt` argument
        is specified, the `defaultclock` will be used.
    order : int, optional
        The priority of of this group for operations occurring at the same time
        step and in the same scheduling slot. Defaults to 0.
    name : str, optional
        A unique name for the group, otherwise use ``spatialneuron_0``, etc.
    """

    def __init__(self, morphology=None, model=None, threshold=None,
                 refractory=False, reset=None, events=None,
                 threshold_location=None,
                 dt=None, clock=None, order=0, Cm=0.9 * uF / cm ** 2, Ri=150 * ohm * cm,
                 name='spatialneuron*', dtype=None, namespace=None,
                 method=('exact', 'exponential_euler', 'rk2', 'heun'),
                 method_options=None):

        # #### Prepare and validate equations
        if isinstance(model, str):
            model = Equations(model)
        if not isinstance(model, Equations):
            raise TypeError(f"model has to be a string or an Equations "
                            f"object, is '{type(model)}' instead.")

        # Insert the threshold mechanism at the specified location
        if threshold_location is not None:
            if hasattr(threshold_location,
                       '_indices'):  # assuming this is a method
                threshold_location = threshold_location._indices()
                # for now, only a single compartment allowed
                if len(threshold_location) == 1:
                    threshold_location = threshold_location[0]
                else:
                    raise AttributeError("Threshold can only be applied on a "
                                         "single location")
            threshold = f"({threshold}) and (i == {str(threshold_location)})"

        # Check flags (we have point currents)
        model.check_flags({DIFFERENTIAL_EQUATION: ('point current',),
                           PARAMETER: ('constant', 'shared', 'linked', 'point current'),
                           SUBEXPRESSION: ('shared', 'point current',
                                           'constant over dt')})
        #: The original equations as specified by the user (i.e. before
        #: inserting point-currents into the membrane equation, before adding
        #: all the internally used variables and constants, etc.).
        self.user_equations = model

        # Separate subexpressions depending whether they are considered to be
        # constant over a time step or not (this would also be done by the
        # NeuronGroup initializer later, but this would give incorrect results
        # for the linearity check)
        model, constant_over_dt = extract_constant_subexpressions(model)

        # Extract membrane equation
        if 'Im' in model:
            if len(model['Im'].flags):
                raise TypeError("Cannot specify any flags for the transmembrane "
                                "current 'Im'.")
            membrane_expr = model['Im'].expr  # the membrane equation
        else:
            raise TypeError("The transmembrane current 'Im' must be defined")

        model_equations = []
        # Insert point currents in the membrane equation
        for eq in model.values():
            if eq.varname == 'Im':
                continue  # ignore -- handled separately
            if 'point current' in eq.flags:
                fail_for_dimension_mismatch(eq.dim, amp,
                                            f"Point current {eq.varname} should be in amp")
                membrane_expr = Expression(
                    f"{str(membrane_expr.code)}+{eq.varname}/area")
                eq = SingleEquation(eq.type, eq.varname, eq.dim, expr=eq.expr,
                                    flags=list(set(eq.flags)-{'point current'}))
            model_equations.append(eq)

        model_equations.append(SingleEquation(SUBEXPRESSION, 'Im',
                                              dimensions=(amp/meter**2).dim,
                                              expr=membrane_expr))
        model_equations.append(SingleEquation(PARAMETER, 'v', volt.dim))
        model = Equations(model_equations)

        ###### Process model equations (Im) to extract total conductance and the remaining current
        # Expand expressions in the membrane equation
        for var, expr in model.get_substituted_expressions(include_subexpressions=True):
            if var == 'Im':
                Im_expr = expr
                break
        else:
            raise AssertionError("Model equations did not contain Im!")

        # Differentiate Im with respect to v
        Im_sympy_exp = str_to_sympy(Im_expr.code)
        v_sympy = sp.Symbol('v', real=True)
        diffed = sp.diff(Im_sympy_exp, v_sympy)

        unevaled_derivatives = diffed.atoms(sp.Derivative)
        if len(unevaled_derivatives):
            raise TypeError(
                f"Cannot take the derivative of '{Im_expr.code}' with respect to v.")

        gtot_str = sympy_to_str(sp.simplify(-diffed))
        I0_str = sympy_to_str(sp.simplify(Im_sympy_exp - diffed*v_sympy))

        if gtot_str == '0':
            gtot_str += '*siemens/meter**2'
        if I0_str == '0':
            I0_str += '*amp/meter**2'
        gtot_str = f"gtot__private={gtot_str}: siemens/meter**2"
        I0_str = f"I0__private={I0_str}: amp/meter**2"

        model += Equations(f"{gtot_str}\n{I0_str}")

        # Insert morphology (store a copy)
        self.morphology = copy.deepcopy(morphology)

        # Flatten the morphology
        self.flat_morphology = FlatMorphology(morphology)

        # Equations for morphology
        # TODO: check whether Cm and Ri are already in the equations
        #       no: should be shared instead of constant
        #       yes: should be constant (check)
        eqs_constants = Equations("""
        length : meter (constant)
        distance : meter (constant)
        area : meter**2 (constant)
        volume : meter**3
        Ic : amp/meter**2
        diameter : meter (constant)
        Cm : farad/meter**2 (constant)
        Ri : ohm*meter (constant, shared)
        r_length_1 : meter (constant)
        r_length_2 : meter (constant)
        time_constant = Cm/gtot__private : second
        space_constant = (2/pi)**(1.0/3.0) * (area/(1/r_length_1 + 1/r_length_2))**(1.0/6.0) /
                         (2*(Ri*gtot__private)**(1.0/2.0)) : meter
        """)
        if self.flat_morphology.has_coordinates:
            eqs_constants += Equations("""
            x : meter (constant)
            y : meter (constant)
            z : meter (constant)
            """)

        NeuronGroup.__init__(self, morphology.total_compartments,
                             model=model + eqs_constants,
                             method_options=method_options,
                             threshold=threshold, refractory=refractory,
                             reset=reset, events=events,
                             method=method, dt=dt, clock=clock, order=order,
                             namespace=namespace, dtype=dtype, name=name)
        # Parameters and intermediate variables for solving the cable equations
        # Note that some of these variables could have meaningful physical
        # units (e.g. _v_star is in volt, _I0_all is in amp/meter**2 etc.) but
        # since these variables should never be used in user code, we don't
        # assign them any units
        self.variables.add_arrays(['_ab_star0', '_ab_star1', '_ab_star2',
                                   '_b_plus', '_b_minus',
                                   '_v_star', '_u_plus', '_u_minus',
                                   '_v_previous', '_c',
                                   # The following two are only necessary for
                                   # C code where we cannot deal with scalars
                                   # and arrays interchangeably:
                                   '_I0_all', '_gtot_all'],
                                  size=self.N, read_only=True)

        self.Cm = Cm
        self.Ri = Ri
        # These explict assignments will load the morphology values from disk
        # in standalone mode
        self.distance_ = self.flat_morphology.distance
        self.length_ = self.flat_morphology.length
        self.area_ = self.flat_morphology.area
        self.diameter_ = self.flat_morphology.diameter
        self.r_length_1_ = self.flat_morphology.r_length_1
        self.r_length_2_ = self.flat_morphology.r_length_2
        if self.flat_morphology.has_coordinates:
            self.x_ = self.flat_morphology.x
            self.y_ = self.flat_morphology.y
            self.z_ = self.flat_morphology.z

        # Performs numerical integration step
        self.add_attribute('diffusion_state_updater')
        self.diffusion_state_updater = SpatialStateUpdater(self, method,
                                                           clock=self.clock,
                                                           order=order)

        # Update v after the gating variables to obtain consistent Ic and Im
        self.diffusion_state_updater.order = 1

        # Creation of contained_objects that do the work
        self.contained_objects.extend([self.diffusion_state_updater])

        if len(constant_over_dt):
            self.subexpression_updater = SubexpressionUpdater(self,
                                                              constant_over_dt)
            self.contained_objects.append(self.subexpression_updater)

    def __getattr__(self, name):
        """
        Subtrees are accessed by attribute, e.g. neuron.axon.
        """
        return self.spatialneuron_attribute(self, name)

    def __getitem__(self, item):
        """
        Selects a segment, where x is a slice of either compartment
        indexes or distances.
        Note a: segment is not a SpatialNeuron, only a Group.
        """
        return self.spatialneuron_segment(self, item)

    @staticmethod
    def _find_subtree_end(morpho):
        """
        Go down a morphology recursively to find the (absolute) index of the
        "final" compartment (i.e. the one with the highest index) of the
        subtree.

        Parameters
        ----------
        morpho : `Morphology`
            The morphology for which to find the index.

        Returns
        -------
        index : int
            The highest index within the subtree.
        """
        indices = [morpho.indices[-1]]
        for child in morpho.children:
            indices.append(SpatialNeuron._find_subtree_end(child))
        return max(indices)

    @staticmethod
    def spatialneuron_attribute(neuron, name):
        """
        Selects a subtree from `SpatialNeuron` neuron and returns a `SpatialSubgroup`.
        If it does not exist, returns the `Group` attribute.
        """
        if name == 'main':  # Main section, without the subtrees
            indices = neuron.morphology.indices[:]
            start, stop = indices[0], indices[-1]
            return SpatialSubgroup(neuron, start, stop + 1,
                                   morphology=neuron.morphology)
        elif (name != 'morphology') and ((name in getattr(neuron.morphology, 'children', [])) or
                                      all([c in 'LR123456789' for c in name])):  # subtree
            morpho = neuron.morphology[name]
            start = morpho.indices[0]
            stop = SpatialNeuron._find_subtree_end(morpho)
            return SpatialSubgroup(neuron, start, stop + 1, morphology=morpho)
        else:
            return Group.__getattr__(neuron, name)

    @staticmethod
    def spatialneuron_segment(neuron, item):
        """
        Selects a segment from `SpatialNeuron` neuron, where item is a slice of
        either compartment indexes or distances.
        Note a: segment is not a `SpatialNeuron`, only a `Group`.
        """
        if isinstance(item, slice) and isinstance(item.start, Quantity):
            if item.step is not None:
                raise ValueError("Cannot specify a step size for slicing based"
                                 "on length.")
            start, stop = item.start, item.stop
            if (not have_same_dimensions(start, meter) or
                    not have_same_dimensions(stop, meter)):
                raise DimensionMismatchError("Start and stop should have units "
                                             "of meter", start, stop)
            # Convert to integers (compartment numbers)
            indices = neuron.morphology.indices[item]
            start, stop = indices[0], indices[-1] + 1
        elif not isinstance(item, slice) and hasattr(item, 'indices'):
            start, stop = to_start_stop(item.indices[:], neuron._N)
        else:
            start, stop = to_start_stop(item, neuron._N)
            if isinstance(neuron, SpatialSubgroup):
                start += neuron.start
                stop += neuron.start

        if start >= stop:
            raise IndexError(f'Illegal start/end values for subgroup, {int(start)}>={int(stop)}')
        if isinstance(neuron, SpatialSubgroup):
            # Note that the start/stop values calculated above are always
            # absolute values, even for subgroups
            neuron = neuron.source
        return Subgroup(neuron, start, stop)


class SpatialSubgroup(Subgroup):
    """
    A subgroup of a `SpatialNeuron`.

    Parameters
    ----------
    source : int
        First compartment.
    stop : int
        Ending compartment, not included (as in slices).
    morphology : `Morphology`
        Morphology corresponding to the subgroup (not the full
        morphology).
    name : str, optional
        Name of the subgroup.
    """

    def __init__(self, source, start, stop, morphology, name=None):
        self.morphology = morphology
        if isinstance(source, SpatialSubgroup):
            source = source.source
            start += source.start
            stop += source.start
        Subgroup.__init__(self, source, start, stop, name)

    def __getattr__(self, name):
        return SpatialNeuron.spatialneuron_attribute(self, name)

    def __getitem__(self, item):
        return SpatialNeuron.spatialneuron_segment(self, item)


class SpatialStateUpdater(CodeRunner, Group):
    """
    The `CodeRunner` that updates the state variables of a `SpatialNeuron`
    at every timestep.
    """

    def __init__(self, group, method, clock, order=0):
        # group is the neuron (a group of compartments)
        self.method_choice = method
        self.group = weakref.proxy(group)

        compartments = group.flat_morphology.n
        sections = group.flat_morphology.sections

        CodeRunner.__init__(self, group,
                            'spatialstateupdate',
                            code="""_gtot = gtot__private
                                    _I0 = I0__private""",
                            clock=clock,
                            when='groups',
                            order=order,
                            name=f"{group.name}_spatialstateupdater*",
                            check_units=False,
                            template_kwds={'number_sections': sections})

        self.variables = Variables(self, default_index='_section_idx')
        self.variables.add_reference('N', group)
        # One value per compartment
        self.variables.add_arange('_compartment_idx', size=compartments)
        self.variables.add_array('_invr', dimensions=siemens.dim,
                                 size=compartments, constant=True,
                                 index='_compartment_idx')
        # one value per section
        self.variables.add_arange('_section_idx', size=sections)
        self.variables.add_array('_P_parent', size=sections,
                                 constant=True)  # elements below diagonal
        self.variables.add_arrays(['_morph_idxchild', '_morph_parent_i',
                                   '_starts', '_ends'], size=sections,
                                  dtype=np.int32, constant=True)
        self.variables.add_arrays(['_invr0', '_invrn'], dimensions=siemens.dim,
                                  size=sections, constant=True)
        # one value per section + 1 value for the root
        self.variables.add_arange('_section_root_idx', size=sections+1)
        self.variables.add_array('_P_diag', size=sections+1,
                                 constant=True, index='_section_root_idx')
        self.variables.add_array('_B', size=sections+1,
                                 constant=True, index='_section_root_idx')
        self.variables.add_array('_morph_children_num',
                                 size=sections+1, dtype=np.int32,
                                 constant=True, index='_section_root_idx')
        # 2D matrices of size (sections + 1) x max children per section
        self.variables.add_arange('_morph_children_idx',
                                  size=len(group.flat_morphology.morph_children))
        self.variables.add_array('_P_children',
                                 size=len(group.flat_morphology.morph_children),
                                 index='_morph_children_idx',
                                 constant=True)  # elements above diagonal
        self.variables.add_array('_morph_children',
                                 size=len(group.flat_morphology.morph_children),
                                 dtype=np.int32, constant=True,
                                 index='_morph_children_idx')
        self._enable_group_attributes()

        self._morph_parent_i = group.flat_morphology.morph_parent_i
        self._morph_children_num = group.flat_morphology.morph_children_num
        self._morph_children = group.flat_morphology.morph_children
        self._morph_idxchild = group.flat_morphology.morph_idxchild
        self._starts = group.flat_morphology.starts
        self._ends = group.flat_morphology.ends

    def before_run(self, run_namespace):
        super(SpatialStateUpdater, self).before_run(run_namespace)
        # Raise a warning if the slow pure numpy version is used
        from brian2.codegen.runtime.numpy_rt.numpy_rt import NumpyCodeObject
        if type(self.code_objects[0]) == NumpyCodeObject:
            # If numpy is used, raise a warning if scipy is not present
            try:
                import scipy
            except ImportError:
                logger.info(('SpatialNeuron will use numpy to do the numerical '
                             'integration -- this will be very slow. Either '
                             'switch to a different code generation target '
                             '(e.g. cython) or install scipy.'),
                            once=True)
