"""
NRV-:class:`.nerve_results` handling.
"""


from ...backend.NRV_Results import sim_results
from .fascicles_results import fascicle_results
from .axons_results import axon_results
from ...backend.log_interface import rise_error, rise_warning, pass_info
from ...backend.MCore import MCH
from ...utils.units import nm, convert, from_nrv_unit
from ...utils.misc import membrane_capacitance_from_model


import matplotlib.pyplot as plt
import numpy as np


def number_in_str(s: str) -> bool:
    return any(i.isdigit() for i in s)

class nerve_results(sim_results):
    """

    """

    def __init__(self, context=None):
        super().__init__(context)


    @property
    def fascicle_keys(self)->list:
        """
        Number of fascicles in the fascicle
        """
        all_keys = self.keys()
        return [i for i in all_keys if ("fascicle" in i and number_in_str(i))]

    @property
    def n_fasc(self)->int:
        """
        Number of fascicles in the fascicle
        """
        return len(self.fascicles)

    @property
    def n_ax(self)->int:
        """
        Number of axons in the fascicle
        """
        fasc_keys = self.fascicle_keys
        _n_ax = 0
        for key in fasc_keys:
            _n_ax += self[key].n_ax
        return _n_ax


    @property
    def axons_type(self)->np.ndarray:
        """
        type of axons of each fascicles

        Returns
        -------
        np.ndarray (self.n_ax, 2)
            _description_
        """
        fasc_keys = self.fascicle_keys
        _mye = np.zeros((self.n_ax, 2))
        _offset = 0
        for key in fasc_keys:
            fasc_mye = np.vstack(
                (
                    self[key].ID*np.ones(self[key].n_ax),
                    self[key].axons_type, 
                )
            )
            _mye[:, _offset:_offset+self[key].n_ax] = fasc_mye.T
            _offset += self[key].n_ax
        return _mye

    def get_fascicle_results(self, ID: int) -> fascicle_results:
        if ID not in self.fascicles_IDs:
            rise_error(("Fascicle ID does not exists."))
        else:
            return self[f"fascicle{ID}"]

    def get_fascicle_results(self, fasc_ID: int, ax_ID:int) -> axon_results:
        if fasc_ID not in self.fascicles_IDs:
            rise_error(f"Fascicle ID: {fasc_ID} does not exists.")
        else:
            if ax_ID >= self[f"fascicle{fasc_ID}"].n_ax:
                rise_error(f"Axon ID: {ax_ID} does not exists in Fascicle {fasc_ID} .")
            else:
                return self[f"fascicle{fasc_ID}"][f"axon{ax_ID}"]

    def get_fascicle_key(self) ->list:
        all_keys = self.keys()
        fascicle_keys = [ i for i in all_keys if ("fascicle" in i and number_in_str(i))]
        return(fascicle_keys)


    def get_recruited_axons(self, ax_type:str = "all", normalize:bool=False) -> int|float:
        """
        Return the number or the ratio of recruited axons in the nerve

        Parameters
        ----------
        diam : float
            diameter above wich axon should be counted.
        ax_type : str, optional
            type of axon counted,possible options:
             - "all" (default)
             - "unmyelinated"
             - "myelinated"
        normalize : bool, optional
            if False the total number of recruited axons is returned, else the ratio is returned, by default False

        Returns
        -------
        float
            number of recruited axons
        """
        fasc_keys = self.fascicle_keys
        for key in fasc_keys:
            fasc_res = self[key]
            n_recr += fasc_res.get_recruited_axons(ax_type=ax_type, normalize=normalize)
        if normalize:
            n_recr /= self.n_fasc
        return n_recr


    # impeddance related methods
    def get_membrane_conductivity(self, x:float=0, t:float=0, unit:str="S/cm**2", mem_th:float=7*nm)->np.array:
        """
        get the membran conductivity of each axon at a position x and a time t

        Parameters
        ----------
        x : float, optional
            x-position in um where to get the conductivity, by default 0
        t : float, optional
            simulation time in ms when to get the conductivity, by default 0
        unit : str, optional
            unit of the returned conductivity see `units`, by default "S/cm**2"
        mem_th : float, optional
            membrane thickness in um, by default 7*nm

        Note
        ----
        depending of the unit parameter this function either return :

            - the surface conductivity in [S]/([m]*[m]): from neuron simulation
            - the conductivity in [S]/[m]:  by multiplying surface conductivity by membrane thickness
        """

        g = []
        f_keys = self.get_fascicle_key()
        for key in f_keys:
            g_ = self[key].get_membrane_conductivity(x=x, t=t, unit=unit, mem_th=mem_th)
            if g_ is not None:
                g = np.concatenate((g,g_))
            else:
                return None
        return g

    def get_membrane_capacitance(self, unit:str="uF/cm**2", mem_th:float=7*nm)->tuple[float]:
            """
            get the membrane capacitance or permitivity of unmyelinated and myelinated axons filling the ner

            Parameters
            ----------
            unit : str, optional
                unit of the returned conductivity see `units`, by default "S/cm**2"
            mem_th : float, optional
                membrane thickness in um, by default 7*nm

            Note
            ----
            depending of the unit parameter this function either return :

                - the surface conductivity in [S]/([m]*[m]): from neuron simulation
                - the conductivity in [S]/[m]:  by multiplying surface conductivity by membrane thickness
            """
            u_c_mem = membrane_capacitance_from_model(self.unmyelinated_param["model"])
            m_c_mem = membrane_capacitance_from_model(self.myelinated_param["model"])

            # Surface capacity in [F]/([m]*[m])
            if "2" in unit:
                return convert(u_c_mem, "S/cm**2", unit), convert(m_c_mem, "S/cm**2", unit)
            # permitivity in [F]/[m]
            else:
                u_c_mem *= from_nrv_unit(mem_th, "cm")
                m_c_mem *= from_nrv_unit(mem_th, "cm")
                return convert(u_c_mem, "S/cm", unit), convert(m_c_mem, "S/cm", unit)


    def get_membrane_complexe_admitance(self, f:float=1., x:float=0, t:float=0, unit:str="S/m", mem_th:float=7*nm)->np.array:
        """
        get the membran complexe admitance of each axon at a position x and a time t for a given frequency

        Parameters
        ----------
        f : float or np.array, optional
            effective frequency in kHz, by default 1
        x : float, optional
            x-position in um where to get the conductivity, by default 0
        t : float, optional
            simulation time in ms when to get the conductivity, by default 0
        unit : str, optional
            unit of the returned conductivity see `units`, by default "S/cm**2"
        mem_th : float, optional
            membrane thickness in um, by default 7*nm
        """
        Y = []
        f_keys = self.get_fascicle_key()
        for key in f_keys:
            Y_ = self[key].get_membrane_complexe_admitance(f=f, x=x, t=t, unit=unit, mem_th=mem_th)
            if Y_ is not None:
                Y = np.concatenate((Y,Y_))
            else:
                return None
        return Y
    
    ## Representation methods
    def plot_recruited_fibers(
        self, axes:plt.axes, contour_color:str="k", myel_color:str="r", unmyel_color:str="b", elec_color:str="gold",num:bool=False, **kwgs
    ):
        """

        """
        if MCH.do_master_only_work():
            ## plot contour
            axes.add_patch(plt.Circle((self.y_grav_center, self.z_grav_center),self.D/2,color=contour_color,
                                fill=False,linewidth=4,))
            fasc_keys = self.get_fascicle_key()
            for key in fasc_keys:
                fasc_res = self[key]
                fasc_res.plot_recruited_fibers(axes=axes,contour_color="grey",
                    myel_color=myel_color,unmyel_color=unmyel_color,num=num,)
            if self.extra_stim is not None:
                self.extra_stim.plot(axes=axes, color=elec_color, nerve_d=self.D)
            axes.set_xlim((-1.1*self.D/2, 1.1*self.D/2))
            axes.set_ylim((-1.1*self.D/2, 1.1*self.D/2))
