from neuron import h, _has_rxd
try:
    if _has_rxd['crxd']:
        from neuron.crxd import species
    else:
        from neuron.rxd import species
    have_rxd = True
except:
  have_rxd = False


def get_mech_names():
    _mech_names = []
    mt = h.MechanismType(0)
    mname = h.ref('')
    for i in range(int(mt.count())):
        mt.select(i)
        mt.selected(mname)
        _mech_names.append(mname[0])
    return _mech_names


def get_rangevars(mech_name):
    ms = h.MechanismStandard(mech_name)
    suffix = '_' + mech_name
    lensuffix = len(suffix)
    mname = h.ref('')
    result = []
    for i in range(int(ms.count())):
        ms.name(mname, i)
        my_mname = mname[0]
        if my_mname.endswith(suffix):
            my_mname = my_mname[:-lensuffix]
        result.append({'type': 'NEURONMECH', 'mech': mech_name, 'var': my_mname, 'name': mech_name + '.' + my_mname})
    return result


def _get_ptrs(plot_what):
    ptrs = []
    if plot_what['type'] == 'NEURON':
        plot_what_var = '_ref_' + plot_what['var']
        for sec in h.allsec():
            try:
                getattr(sec(0.5), plot_what_var)
                bad_section = False
            except NameError:
                bad_section = True
            except AttributeError:
                bad_section = True
            if not bad_section:
                ptrs += [getattr(seg, plot_what_var) for seg in sec]
            else:
                ptrs += [None] * sec.nseg
    elif plot_what['type'] == 'SECVAR':
        plot_what_var = '_ref_' + plot_what['var']
        for sec in h.allsec():
            ptrs += [getattr(sec, plot_what_var)] * sec.nseg
    elif plot_what['type'] == 'NEURONMECH':
        plot_what_mech = plot_what['mech']
        plot_what_var = '_ref_' + plot_what['var']
        for sec in h.allsec():
            for seg in sec:
                try:
                    mech = getattr(seg, plot_what_mech)
                except:
                    ptrs.append(None)
                    continue
                ptrs.append(getattr(mech, plot_what_var))
    elif plot_what['type'] == 'NEURONRXD':
        # TODO: there has to be something more efficient than this
        # TODO: this probably won't work well with 3D nodes
        sp = plot_what['species']
        region = plot_what['region']
        for s in species._all_species:
            s = s()
            if s is None: continue
            if s.name == sp:
                rs = s._regions
                if hasattr(rs, '__len__'):
                    rs = list(rs)
                else:
                    rs = list([rs])
                for r in rs:
                    if r._name == region:
                        base_obj = s[r].nodes
                        for sec in h.allsec():
                            sec_obj = base_obj(sec)
                            if sec_obj:
                                for seg in sec:
                                    nodes = sec_obj(seg.x)
                                    if len(nodes) != 1:
                                        print('gui: node length mismatch; 3D?')
                                    if not nodes:
                                        ptrs.append(None)
                                    else:
                                        # TODO: change this to use molecules when appropriate (stochastic sims)
                                        ptrs.append(nodes[0]._ref_concentration)
                        break
                else:
                    continue
                break
        else:
            # no such species found
            for sec in h.allsec():
                ptrs += [None] * sec.nseg
    
    elif plot_what['type'] is None:
        # no such species found
        for sec in h.allsec():
            ptrs += [None] * sec.nseg

    else:
        print('gui: unknown plot_what type', plot_what)

    return ptrs


def _neuron_var(name):
    return {'type': 'NEURON', 'var': name, 'name': name}


def _section_var(name):
    return {'type': 'SECVAR', 'var': name, 'name': name}


def get_rxd_vars(secs):
    # TODO: filter by sections present
    result = []
    if have_rxd:
        for sp in species._all_species:
            sp = sp()
            if sp is not None and sp.name is not None:
                regions = sp._regions
                if hasattr(regions, '__len__'):
                    regions = list(regions)
                else:
                    regions = list([regions])
                result += [{'type': 'NEURONRXD',
                            'name': '%s[%s]' % (sp.name, r._name),
                            'species': sp.name,
                            'region': r._name}
                             for r in regions if r._name is not None]
    return result


def rangevars_present(secs):
    result = []
    for name in get_mech_names():
        for sec in secs:
            try:
                if hasattr(sec(0.5), name):
                    if name[-4:] != '_ion':
                        result += get_rangevars(name)
                    else:
                        ionname = name[:-4]
                        result += [_neuron_var(name) for name in ['i' + ionname, ionname + 'i', ionname + 'o', 'e' + ionname]]
                    break
            except NameError:
                pass
    result += get_rxd_vars(secs)
    return [{'type': None, 'name': 'No coloration'}] + sorted([_neuron_var('cm'), _neuron_var('diam'), _neuron_var('i_cap'), _neuron_var('v')] + result, key=lambda n: n['name'].lower()) # _section_var('Ra')
