"""
This file was automatically generated by the DendroTweaks toolbox (https://dendrotweaks.dendrites.gr/).
It provides a plain (Python) NEURON code representation of a biophysical model of a single neuron
built using DendroTweaks, which can be run independently of DendroTweaks.

Classes:
    Cell: Represents a biophysical NEURON model of a single neuron with methods to:
        - Load morphology from an SWC file.
        - Distribute passive properties.
        - Insert ion channel mechanisms.
        - Distribute parameters across the neuron.
        - Set the number of segments for each section.
        - Add stimuli and recordings.

Functions:
    load_mechanisms: Load NEURON mechanisms from a specified directory.
    init_simulation: Initialize a NEURON simulation with specified parameters.
    run: Run a NEURON simulation for a given duration.
    get_domain: Retrieve the domain of a given segment.
    linear: Compute a linear function based on distance.
    sinusoidal: Compute a sinusoidal function based on distance.
    exponential: Compute an exponential function based on distance.
    sigmoid: Compute a sigmoid function based on distance.
    gaussian: Compute a Gaussian function based on distance.
    step: Compute a step function based on distance.
    inherit: Inherit a parameter value from the parent segment.
"""

import os

import numpy as np
from numpy import polyval

import neuron
from neuron import h
h.load_file('stdrun.hoc')


class Cell():
    """
    A class representing a biophysical NEURON model of a single neuron.
    """

    def __init__(self, path_to_swc_file: str):
        self.name = os.path.basename(path_to_swc_file).replace('.swc', '')
        self.load_morphology(path_to_swc_file)
        self.distribute_passive()
        self.set_geom_nseg()
        self.insert_mechanisms()
        self.distribute_parameters()

    @property
    def all_segments(self):
        return [seg for sec in self.all for seg in sec]
        
    ### Morphology methods ###

    def load_morphology(self, path_to_swc_file: str) -> None:
        if path_to_swc_file.endswith('.swc'):
            self._load_swc(path_to_swc_file)
        elif path_to_swc_file.endswith('.asc'):
            self._load_asc(path_to_swc_file)
        else:
            raise ValueError(f"File type not supported: {path_to_swc_file}")

    def _load_swc(self, path_to_swc_file: str) -> None:
        h.load_file('import3d.hoc')
        swc_importer = h.Import3d_SWC_read()
        swc_importer.input(path_to_swc_file)
        imported_cell = h.Import3d_GUI(swc_importer, False)
        imported_cell.instantiate(self)

    def _load_asc(self, path_to_asc_file: str) -> None:
        h.load_file('import3d.hoc')
        asc_importer = h.Import3d_Neurolucida3()
        asc_importer.input(path_to_asc_file)
        imported_cell = h.Import3d_GUI(asc_importer, False)
        imported_cell.instantiate(self)

    def distance(self, seg, from_seg=None):
        if from_seg is None:
            from_seg = self.soma[0](0.5)
        return h.distance(from_seg, seg)

    def domain_distance(self, seg):
        parent = self._find_parent_with_different_domain(seg.sec)
        if parent:
            return h.distance(parent(1), seg)
        return 0
    
    def _find_parent_with_different_domain(self, sec):
        parentseg = sec.parentseg()
        if not parentseg:
            return None
        parent = parentseg.sec
        while parent:
            if get_domain(parent(0.5)) != get_domain(sec(0.5)):
                return parent
            parentseg = parent.parentseg()
            if not parentseg:
                return None
            parent = parentseg.sec
        return None

    ### Segmentation methods ###

    def set_geom_nseg(self, d_lambda:float=0.1, f:float=100):
        for sec in self.all:
            sec.nseg = int((sec.L/(d_lambda*h.lambda_f(f, sec=sec)) + 0.9)/2)*2 + 1
        
    ### Mechanism methods ###

    def insert_mechanisms(self):
        {% for domain, mechanisms in domains_to_mechs.items() %}
        for sec in self.{{ domains_to_NEURON[domain] }}:
            {% for mechanism in mechanisms %}
            sec.insert('{{ mechanism }}')
            {%- endfor %}
        {% endfor %}

    ### Parameter distribution methods ###

    def set_param(self, seg, param: str, mech: str, value: float) -> None:
        if param == 'Ra':
            setattr(seg.sec, param, value)
        if param == 'cm':
            setattr(seg, param, value)
        else:
            if seg.sec.has_membrane(mech):
                setattr(seg, param, value)
            else:
                if param in ['ena', 'ek', 'eca']:
                    if hasattr(seg, param):
                        setattr(seg, param, value)
                
    def distribute_passive(self):

        for seg in self.all_segments:

            domain = get_domain(seg)
            distance = self.distance(seg)
            domain_distance = self.domain_distance(seg)
            diam = seg.diam
            section_diam = seg.sec.diam

            {% for param, mech in params_to_mechs.items() -%}
            {% if param in ['cm', 'Ra']%}
            {% set groups = param_dict[param] -%}
            {% for group_name, distribution in groups.items() -%}
            {% set group = groups_dict[group_name] -%}
            if domain in {{ params_to_valid_domains[param][group_name] }}:
                {% if group.select_by -%}
                {% set min_val = group.min_value if group.min_value is not none else 0 -%}
                {% if group.max_value is not none -%}
                if {{ min_val }} < {{ group.select_by }} <= {{ group.max_value }}:
                {% else -%}
                if {{ min_val }} < {{ group.select_by }}:
                {% endif %}
                    {% if distribution.function_name == "constant" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", {{ distribution.parameters.value }})
                    {% elif distribution.function_name == "linear" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", linear(distance, slope={{ distribution.parameters.slope }}, intercept={{ distribution.parameters.intercept }}))
                    {% elif distribution.function_name == "polynomial" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", polyval([{% for coeff in distribution.parameters.coeffs %}{{ coeff }}{% if not loop.last %}, {% endif %}{% endfor %}], distance))
                    {% else -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", {{ distribution.function_name }}(distance, **{{ distribution.parameters }}))
                    {% endif %}
                {% else -%}
                    {% if distribution.function_name == "constant" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", {{ distribution.parameters.value }})
                    {% elif distribution.function_name == "linear" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", linear(distance, slope={{ distribution.parameters.slope }}, intercept={{ distribution.parameters.intercept }}))
                    {% elif distribution.function_name == "polynomial" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", polyval([{% for coeff in distribution.parameters.coeffs %}{{ coeff }}{% if not loop.last %}, {% endif %}{% endfor %}], distance))
                    {% else -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", {{ distribution.function_name }}(distance, **{{ distribution.parameters }}))
                    {% endif %}
                {% endif %}
            {% endfor -%}
            {% endif -%}
            {% endfor %}

    def distribute_parameters(self):

        for seg in self.all_segments:

            domain = get_domain(seg)
            distance = self.distance(seg)
            domain_distance = self.domain_distance(seg)
            diam = seg.diam
            section_diam = seg.sec.diam

            {% for param, mech in params_to_mechs.items() -%}
            {% if param not in ['cm', 'Ra']%}
            {% set groups = param_dict[param] -%}
            {% for group_name, distribution in groups.items() -%}
            {% set group = groups_dict[group_name] -%}
            if domain in {{ params_to_valid_domains[param][group_name] }}:
                {% if group.select_by -%}
                {% set min_val = group.min_value if group.min_value is not none else 0 -%}
                {% if group.max_value is not none -%}
                if {{ min_val }} < {{ group.select_by }} <= {{ group.max_value }}:
                {% else -%}
                if {{ min_val }} < {{ group.select_by }}:
                {% endif %}
                    {% if distribution.function_name == "constant" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", {{ distribution.parameters.value }})
                    {% elif distribution.function_name == "linear" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", linear(distance, slope={{ distribution.parameters.slope }}, intercept={{ distribution.parameters.intercept }}))
                    {% elif distribution.function_name == "polynomial" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", polyval([{% for coeff in distribution.parameters.coeffs %}{{ coeff }}{% if not loop.last %}, {% endif %}{% endfor %}], distance))
                    {% else -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", {{ distribution.function_name }}(distance, **{{ distribution.parameters }}))
                    {% endif %}
                {% else -%}
                    {% if distribution.function_name == "constant" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", {{ distribution.parameters.value }})
                    {% elif distribution.function_name == "linear" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", linear(distance, slope={{ distribution.parameters.slope }}, intercept={{ distribution.parameters.intercept }}))
                    {% elif distribution.function_name == "polynomial" -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", polyval([{% for coeff in distribution.parameters.coeffs %}{{ coeff }}{% if not loop.last %}, {% endif %}{% endfor %}], distance))
                    {% else -%}
                    self.set_param(seg, "{{ param }}", "{{ mech }}", {{ distribution.function_name }}(distance, **{{ distribution.parameters }}))
                    {% endif %}
                {% endif %}
            {% endfor -%}
            {% endif -%}
            {% endfor %}

    ### Stimulation methods ###

    def add_stimuli(self):
        self.add_iclamps()
        self.add_synapses()

    def add_recordings(self):
        recordings = []
        {% for var, recs in recordings.items() %}
        {% for seg, rec in recs.items() %}
        rec = h.Vector()
        rec.record(self.{{seg._section.domain}}[{{seg._section.domain_idx}}]({{seg.x}})._ref_{{ var }})
        recordings.append(rec)
        {%- endfor -%}
        {% endfor %}
        return recordings

    def add_iclamps(self):
        iclamps = []
        {% for seg, iclamp in iclamps.items() %}
        iclamp = h.IClamp(self.{{seg._section.domain}}[{{seg._section.domain_idx}}]({{seg.x}}))
        iclamp.delay = {{ iclamp.delay }}
        iclamp.dur = {{ iclamp.dur }}
        iclamp.amp = {{ iclamp.amp }}
        iclamps.append(iclamp)
        {% endfor %}
        return iclamps


# =================================================================================================
# SIMULATION FUNCTIONS
# =================================================================================================

def load_mechanisms(path_to_mod: str, recompile:bool=False) -> None:
    """
    Load NEURON mechanisms from a directory.

    Parameters
    ----------
    path_to_mod : str
        The path to the directory containing the mod files.
    recompile : bool
        Whether to recompile the mod files before loading
    """

    if recompile:
        cwd = os.getcwd()
        os.chdir(path_to_mod)
        os.system('nrnivmodl')
        os.chdir(cwd)
        print(f'Compiled mod files in "{path_to_mod}"')

    neuron.load_mechanisms(path_to_mod)
    print(f'Loaded mod files from "{path_to_mod}"')

def init_simulation(temperature:float=37, dt:float=0.025, v_init:float=-65, cvode:bool=False) -> None:
    """
    Initialize a NEURON simulation.

    Parameters
    ----------
    temperature : float
        The temperature in degrees Celsius.
    dt : float
        The time step of the simulation in ms.
    v_init : float
        The initial membrane potential in mV.
    cvode : bool
        Whether to use the CVode variable time step integrator.
    """
    h.CVode().active(cvode)
    h.celsius = temperature
    h.dt = dt
    h.stdinit()
    h.init()
    h.finitialize(v_init)
    if h.cvode.active():
        h.cvode.re_init()
    else:
        h.fcurrent()
    h.frecord_init()

def run(duration:float=300, **kwargs) -> None:
    """
    Run a NEURON simulation.

    Parameters
    ----------
    duration : float
        The duration of the simulation in ms.
    """
    init_simulation(**kwargs)
    h.continuerun(duration)

# =================================================================================================
# HELPER FUNCTIONS
# =================================================================================================

def get_domain(seg) -> str:
    sec = seg.sec
    sec_name = sec.name()
    domain = sec_name.split('.')[-1].split('[')[0]
    return domain

def linear(distance: float, slope:float = 0, intercept: float = 0) -> float:
    return slope * distance + intercept

def sinusoidal(distance: float, amplitude: float = 0, frequency: float = 1, phase: float = 0) -> float:
    return amplitude * np.sin(frequency * distance + phase)

def exponential(distance: float, vertical_shift:float = 0, scale_factor: float =1, growth_rate: float=1, horizontal_shift: float = 0) -> float:
    return vertical_shift + scale_factor * np.exp(growth_rate * (distance - horizontal_shift))

def sigmoid(distance: float, vertical_shift: float=0, scale_factor: float=1, growth_rate: float=1, horizontal_shift: float=0) -> float:
    return vertical_shift + scale_factor / (1 + np.exp(-growth_rate*(distance - horizontal_shift)))

def gaussian(distance: float, amplitude: float, mean: float, std: float) -> float:
    return amplitude * np.exp(-((distance - mean) ** 2) / (2 * std ** 2))

def step(distance: float, max_value: float,  min_value: float, start: float, end: float) -> float:
    if start < distance < end:
        return max_value
    else:
        return min_value

def inherit(seg, param: str) -> float:
    return getattr(seg.sec.parentseg(), param, np.nan)