#include <../../nrnconf.h>

#include <stdio.h>
#include <InterViews/resource.h>
#include <nrnoc2iv.h>
#include <classreg.h>
#include <nrnpython.h>
#include <hoccontext.h>
#include "nrnpy_utils.h"

#include "parse.hpp"
extern void hoc_nopop();
extern void hoc_pop_defer();
extern Object* hoc_new_object(Symbol*, void*);
extern int hoc_stack_type();
extern char** hoc_strpop();
extern Object** hoc_objpop();
extern Object* hoc_pop_object();
extern void hoc_tobj_unref(Object**);
extern int hoc_ipop();
PyObject* nrnpy_hoc2pyobject(Object*);
PyObject* hocobj_call_arg(int);
Object* nrnpy_pyobject_in_obj(PyObject*);
int nrnpy_ho_eq_po(Object*, PyObject*);
char* nrnpyerr_str();
extern void* (*nrnpy_save_thread)();
extern void (*nrnpy_restore_thread)(void*);
static void* save_thread() {
    return PyEval_SaveThread();
}
static void restore_thread(void* g) {
    PyEval_RestoreThread((PyThreadState*) g);
}
extern Symbol* nrnpy_pyobj_sym_;
extern void (*nrnpy_py2n_component)(Object*, Symbol*, int, int);
extern void (*nrnpy_hpoasgn)(Object*, int);
extern double (*nrnpy_praxis_efun)(Object*, Object*);
extern int (*nrnpy_hoccommand_exec)(Object*);
extern int (*nrnpy_hoccommand_exec_strret)(Object*, char*, int);
extern void (*nrnpy_cmdtool)(Object*, int, double, double, int);
extern double (*nrnpy_func_call)(Object*, int, int*);
extern Object* (*nrnpy_callable_with_args)(Object*, int);
extern double (*nrnpy_guigetval)(Object*);
extern void (*nrnpy_guisetval)(Object*, double);
extern int (*nrnpy_guigetstr)(Object*, char**);
extern char* (*nrnpy_po2pickle)(Object*, size_t* size);
extern Object* (*nrnpy_pickle2po)(char*, size_t size);
extern char* (*nrnpy_callpicklef)(char*, size_t size, int narg, size_t* retsize);
extern int (*nrnpy_pysame)(Object*, Object*);  // contain same Python object
extern Object* (*nrnpympi_alltoall_type)(int, int);
typedef struct {
    PyObject_HEAD Section* sec_;
    char* name_;
    PyObject* cell_;
} NPySecObj;
extern NPySecObj* newpysechelp(Section* sec);
extern void (*nrnpy_call_python_with_section)(Object*, Section*);
extern "C" void nrnpython_reg_real();
PyObject* nrnpy_ho2po(Object*);
void nrnpy_decref_defer(PyObject*);
PyObject* nrnpy_pyCallObject(PyObject*, PyObject*);

Object* nrnpy_po2ho(PyObject*);
static void py2n_component(Object*, Symbol*, int, int);
static void hpoasgn(Object*, int);
static double praxis_efun(Object*, Object*);
static int hoccommand_exec(Object*);
static int hoccommand_exec_strret(Object*, char*, int);
static void grphcmdtool(Object*, int, double, double, int);
static double func_call(Object*, int, int*);
static Object* callable_with_args(Object*, int);
static double guigetval(Object*);
static void guisetval(Object*, double);
static int guigetstr(Object*, char**);
static char* po2pickle(Object*, size_t* size);
static Object* pickle2po(char*, size_t size);
static char* call_picklef(char*, size_t size, int narg, size_t* retsize);
static Object* py_alltoall_type(int, int);
static int pysame(Object*, Object*);
static PyObject* main_module;
static PyObject* main_namespace;
static hoc_List* dlist;
#if NRNPYTHON_DYNAMICLOAD
extern int nrnpy_site_problem;
extern int* nrnpy_site_problem_p;
#endif

class Py2Nrn {
  public:
    Py2Nrn();
    virtual ~Py2Nrn();
    int type_;  // 0 toplevel
    PyObject* po_;
};

static void* p_cons(Object*) {
    Py2Nrn* p = new Py2Nrn();
    return p;
}

static void p_destruct(void* v) {
    delete (Py2Nrn*) v;
}

Member_func p_members[] = {0, 0};

static void call_python_with_section(Object* pyact, Section* sec) {
    PyObject* po = ((Py2Nrn*) pyact->u.this_pointer)->po_;
    PyObject* r;
    PyLockGIL lock;

    PyObject* args = PyTuple_Pack(1, (PyObject*) newpysechelp(sec));
    r = nrnpy_pyCallObject(po, args);
    Py_XDECREF(args);
    Py_XDECREF(r);
    if (!r) {
        char* mes = nrnpyerr_str();
        if (mes) {
            Fprintf(stderr, "%s\n", mes);
            free(mes);
            lock.release();
            hoc_execerror("Call of Python Callable failed", NULL);
        }
        if (PyErr_Occurred()) {
            PyErr_Print();
        }
    }
}

extern void* (*nrnpy_opaque_obj2pyobj_p_)(Object*);
static void* opaque_obj2pyobj(Object* ho) {
    assert(ho && ho->ctemplate->sym == nrnpy_pyobj_sym_);
    PyObject* po = ((Py2Nrn*) ho->u.this_pointer)->po_;
    assert(po);
    return po;
}

extern "C" void nrnpython_reg_real() {
    // printf("nrnpython_reg_real()\n");
    class2oc("PythonObject", p_cons, p_destruct, p_members, NULL, NULL, NULL);
    Symbol* s = hoc_lookup("PythonObject");
    assert(s);
    nrnpy_pyobj_sym_ = s;
    nrnpy_py2n_component = py2n_component;
    nrnpy_call_python_with_section = call_python_with_section;
    nrnpy_hpoasgn = hpoasgn;
    nrnpy_praxis_efun = praxis_efun;
    nrnpy_hoccommand_exec = hoccommand_exec;
    nrnpy_hoccommand_exec_strret = hoccommand_exec_strret;
    nrnpy_cmdtool = grphcmdtool;
    nrnpy_func_call = func_call;
    nrnpy_callable_with_args = callable_with_args;
    nrnpy_guigetval = guigetval;
    nrnpy_guisetval = guisetval;
    nrnpy_guigetstr = guigetstr;
    nrnpy_po2pickle = po2pickle;
    nrnpy_pickle2po = pickle2po;
    nrnpy_callpicklef = call_picklef;
    nrnpympi_alltoall_type = py_alltoall_type;
    nrnpy_pysame = pysame;
    nrnpy_save_thread = save_thread;
    nrnpy_restore_thread = restore_thread;
    nrnpy_opaque_obj2pyobj_p_ = opaque_obj2pyobj;
    dlist = hoc_l_newlist();
#if NRNPYTHON_DYNAMICLOAD
    nrnpy_site_problem_p = &nrnpy_site_problem;
#endif
}


Py2Nrn::Py2Nrn() {
    po_ = NULL;
    type_ = 0;
    //	printf("Py2Nrn() %p\n", this);
}
Py2Nrn::~Py2Nrn() {
    PyLockGIL lock;
    Py_XDECREF(po_);
    //	printf("~Py2Nrn() %p\n", this);
}

int nrnpy_ho_eq_po(Object* ho, PyObject* po) {
    if (ho->ctemplate->sym == nrnpy_pyobj_sym_) {
        return ((Py2Nrn*) ho->u.this_pointer)->po_ == po;
    }
    return 0;
}

int pysame(Object* o1, Object* o2) {
    if (o2->ctemplate->sym == nrnpy_pyobj_sym_) {
        return nrnpy_ho_eq_po(o1, ((Py2Nrn*) o2->u.this_pointer)->po_);
    }
    return 0;
}

PyObject* nrnpy_hoc2pyobject(Object* ho) {
    PyObject* po = ((Py2Nrn*) ho->u.this_pointer)->po_;
    if (!po) {
        if (!main_module) {
            main_module = PyImport_AddModule("__main__");
            main_namespace = PyModule_GetDict(main_module);
            Py_INCREF(main_module);
            Py_INCREF(main_namespace);
        }
        po = main_module;
    }
    return po;
}

Object* nrnpy_pyobject_in_obj(PyObject* po) {
    Py2Nrn* pn = new Py2Nrn();
    pn->po_ = po;
    Py_INCREF(po);
    pn->type_ = 1;
    Object* on = hoc_new_object(nrnpy_pyobj_sym_, (void*) pn);
    hoc_obj_ref(on);
    return on;
}

PyObject* nrnpy_pyCallObject(PyObject* callable, PyObject* args) {
    // When hoc calls a PythonObject method, then in case python
    // calls something back in hoc, the hoc interpreter must be
    // at the top level
    HocTopContextSet PyObject* p = PyObject_CallObject(callable, args);
#if 0
printf("PyObject_CallObject callable\n");
PyObject_Print(callable, stdout, 0);
printf("\nargs\n");
PyObject_Print(args, stdout, 0);
printf("\nreturn %p\n", p);
#endif
    HocContextRestore
        // It would be nice to handle the error here, ending with a hoc_execerror
        // for any Exception (note, that does not include SystemExit). However
        // since many, but not all, of the callers need to clean up and
        // release the GIL, errors get handled by the caller or higher up.
        // The almost generic idiom is:
        /**
        if (!p) {
          char* mes = nrnpyerr_str();
          if (mes) {
            Fprintf(stderr, "%s\n", mes);
            free(mes);
            hoc_execerror("Call of Python Callable failed", NULL);
          }
          if (PyErr_Occurred()) {
            PyErr_Print(); // Python process will exit with the error code specified by the
        SystemExit instance.
          }
        }
        **/
        return p;
}

void py2n_component(Object* ob, Symbol* sym, int nindex, int isfunc) {
#if 0
	if (isfunc) {
	printf("py2n_component %s.%s(%d)\n", hoc_object_name(ob), sym->name, nindex);
	}else{
	printf("py2n_component %s.%s[%d]\n", hoc_object_name(ob), sym->name, nindex);
	}
#endif
    int i;
    Py2Nrn* pn = (Py2Nrn*) ob->u.this_pointer;
    PyObject* head = pn->po_;
    PyObject* tail;
    PyLockGIL lock;
    if (pn->type_ == 0) {  // top level
        if (!main_module) {
            main_module = PyImport_AddModule("__main__");
            main_namespace = PyModule_GetDict(main_module);
            Py_INCREF(main_module);
            Py_INCREF(main_namespace);
        }
        tail = PyRun_String(sym->name, Py_eval_input, main_namespace, main_namespace);
    } else {
        Py_INCREF(head);
        if (strcmp(sym->name, "_") == 0) {
            tail = head;
            Py_INCREF(tail);
        } else {
            tail = PyObject_GetAttrString(head, sym->name);
        }
    }
    if (!tail) {
        PyErr_Print();
        lock.release();
        hoc_execerror("No attribute:", sym->name);
    }
    PyObject* args = 0;
    Object* on;
    PyObject* result = 0;
    if (isfunc) {
        args = PyTuple_New(nindex);
        for (i = 0; i < nindex; ++i) {
            PyObject* arg = nrnpy_hoc_pop();
            // PyObject_Print(arg, stdout, 0);
            // printf(" %d   arg %d\n", arg->ob_refcnt,  i);
            if (PyTuple_SetItem(args, nindex - 1 - i, arg)) {
                assert(0);
            }
        }
        // printf("PyObject_CallObject %s %p\n", sym->name, tail);
        result = nrnpy_pyCallObject(tail, args);
        Py_DECREF(args);
        // PyObject_Print(result, stdout, 0);
        // printf("  result of call\n");
        if (!result) {
            char* mes = nrnpyerr_str();
            Py_XDECREF(tail);
            Py_XDECREF(head);
            if (mes) {
                Fprintf(stderr, "%s\n", mes);
                free(mes);
                lock.release();
                hoc_execerror("PyObject method call failed:", sym->name);
            }
            if (PyErr_Occurred()) {
                PyErr_Print();
            }
            return;
        }
    } else if (nindex) {
        PyObject* arg;
        if (hoc_stack_type() == NUMBER) {
            arg = Py_BuildValue("l", (long) hoc_xpop());
        } else {
            arg = nrnpy_hoc_pop();
        }
        result = PyObject_GetItem(tail, arg);
        if (!result) {
            PyErr_Print();
            lock.release();
            hoc_execerror("Python get item failed:", hoc_object_name(ob));
        }
    } else {
        result = tail;
        Py_INCREF(result);
    }
    // printf("py2n_component %s %d %s result refcount=%d\n", hoc_object_name(ob),
    // ob->refcount, sym->name, result->ob_refcnt);
    // if numeric, string, or python HocObject return those
    if (nrnpy_numbercheck(result)) {
        hoc_pop_defer();
        PyObject* pn = PyNumber_Float(result);
        hoc_pushx(PyFloat_AsDouble(pn));
        Py_XDECREF(pn);
        Py_XDECREF(result);
    } else if (is_python_string(result)) {
        char** ts = hoc_temp_charptr();
        Py2NRNString str(result, true);
        *ts = str.c_str();
        hoc_pop_defer();
        hoc_pushstr(ts);
        // how can we defer the result unref til the string is popped
        nrnpy_decref_defer(result);
    } else {
        // PyObject_Print(result, stdout, 0);
        on = nrnpy_po2ho(result);
        hoc_pop_defer();
        hoc_push_object(on);
        if (on) {
            on->refcount--;
        }
        Py_XDECREF(result);
    }
    Py_XDECREF(head);
    Py_DECREF(tail);
}

static void hpoasgn(Object* o, int type) {
    int err = 0;
    int nindex;
    Symbol* sym;
    PyObject* poleft;
    PyObject* poright;
    if (type == NUMBER) {
        poright = PyFloat_FromDouble(hoc_xpop());
    } else if (type == STRING) {
        poright = Py_BuildValue("s", *hoc_strpop());
    } else if (type == OBJECTVAR || type == OBJECTTMP) {
        Object** po2 = hoc_objpop();
        poright = nrnpy_ho2po(*po2);
        hoc_tobj_unref(po2);
    } else {
        hoc_execerror("Cannot assign that type to PythonObject", (char*) 0);
    }
    Object* stack_value = hoc_pop_object();
    assert(o == stack_value);
    poleft = nrnpy_hoc2pyobject(o);
    sym = hoc_spop();
    nindex = hoc_ipop();
    // printf("hpoasgn %s %s %d\n", hoc_object_name(o), sym->name, nindex);
    if (nindex == 0) {
        err = PyObject_SetAttrString(poleft, sym->name, poright);
    } else if (nindex == 1) {
        PyObject* key = PyLong_FromDouble(hoc_xpop());
        PyObject* a = PyObject_GetAttrString(poleft, sym->name);
        if (a) {
            err = PyObject_SetItem(a, key, poright);
            Py_DECREF(a);
        } else {
            err = -1;
        }
        Py_DECREF(key);
    } else {
        char buf[512];
        sprintf(buf, "%s.%s[][]...=...:", hoc_object_name(o), sym->name);
        hoc_execerror(buf, "HOC cannot handle PythonObject assignment with more than one index.");
    }
    Py_DECREF(poright);
    if (err) {
        PyErr_Print();
        hoc_execerror("Assignment to PythonObject failed", NULL);
    }
}

void nrnpy_decref_defer(PyObject* po) {
    if (po) {
#if 0
		PyObject* ps = PyObject_Str(po);
		printf("defer %s\n", PyString_AsString(ps));
		Py_DECREF(ps);
#endif
        hoc_l_lappendvoid(dlist, (void*) po);
    }
}

static PyObject* hoccommand_exec_help1(PyObject* po) {
    PyObject* r;
    // PyObject_Print(po, stdout, 0);
    // printf("\n");
    if (PyTuple_Check(po)) {
        PyObject* args = PyTuple_GetItem(po, 1);
        if (!PyTuple_Check(args)) {
            args = PyTuple_Pack(1, args);
        } else {
            Py_INCREF(args);
        }
        // PyObject_Print(PyTuple_GetItem(po, 0), stdout, 0);
        // printf("\n");
        // PyObject_Print(args, stdout, 0);
        // printf("\n");
        // printf("threadstate %p\n", PyThreadState_GET());
        r = nrnpy_pyCallObject(PyTuple_GetItem(po, 0), args);
        Py_DECREF(args);
    } else {
        PyObject* args = PyTuple_New(0);
        r = nrnpy_pyCallObject(po, args);
        Py_DECREF(args);
    }
    return r;
}

static PyObject* hoccommand_exec_help(Object* ho) {
    PyObject* po = ((Py2Nrn*) ho->u.this_pointer)->po_;
    // printf("%s\n", hoc_object_name(ho));
    return hoccommand_exec_help1(po);
}

static double praxis_efun(Object* ho, Object* v) {
    PyLockGIL lock;

    PyObject* pc = nrnpy_ho2po(ho);
    PyObject* pv = nrnpy_ho2po(v);
    PyObject* po = Py_BuildValue("(OO)", pc, pv);
    Py_XDECREF(pc);
    Py_XDECREF(pv);
    PyObject* r = hoccommand_exec_help1(po);
    Py_XDECREF(po);
    if (!r) {
        char* mes = nrnpyerr_str();
        if (mes) {
            Fprintf(stderr, "%s\n", mes);
            free(mes);
            lock.release();
            hoc_execerror("Call of Python Callable failed in praxis_efun", NULL);
        }
        if (PyErr_Occurred()) {
            PyErr_Print();
        }
        return 1e9;  // SystemExit?
    }
    PyObject* pn = PyNumber_Float(r);
    double x = PyFloat_AsDouble(pn);
    Py_XDECREF(pn);
    Py_XDECREF(r);
    return x;
}

static int hoccommand_exec(Object* ho) {
    PyLockGIL lock;

    PyObject* r = hoccommand_exec_help(ho);
    if (r == NULL) {
        char* mes = nrnpyerr_str();
        if (mes) {
            Fprintf(stderr, "%s\n", mes);
            free(mes);
            lock.release();
            hoc_execerror("Python Callback failed", 0);
        }
        if (PyErr_Occurred()) {
            PyErr_Print();
        }
    }
    Py_XDECREF(r);
    return (r != NULL);
}

static int hoccommand_exec_strret(Object* ho, char* buf, int size) {
    PyLockGIL lock;

    PyObject* r = hoccommand_exec_help(ho);
    if (r) {
        PyObject* pn = PyObject_Str(r);
        Py2NRNString str(pn);
        Py_XDECREF(pn);
        strncpy(buf, str.c_str(), size);
        buf[size - 1] = '\0';
        Py_XDECREF(r);
    } else {
        char* mes = nrnpyerr_str();
        if (mes) {
            Fprintf(stderr, "%s\n", mes);
            free(mes);
            lock.release();
            hoc_execerror("Python Callback failed", 0);
        }
        if (PyErr_Occurred()) {
            PyErr_Print();
        }
    }
    return (r != NULL);
}

static void grphcmdtool(Object* ho, int type, double x, double y, int key) {
    PyObject* po = ((Py2Nrn*) ho->u.this_pointer)->po_;
    PyObject* r;
    PyLockGIL lock;

    PyObject* args = PyTuple_Pack(
        4, PyInt_FromLong(type), PyFloat_FromDouble(x), PyFloat_FromDouble(y), PyInt_FromLong(key));
    r = nrnpy_pyCallObject(po, args);
    Py_XDECREF(args);
    Py_XDECREF(r);
    if (!r) {
        char* mes = nrnpyerr_str();
        if (mes) {
            Fprintf(stderr, "%s\n", mes);
            free(mes);
            lock.release();
            hoc_execerror("Python Callback failed", 0);
        }
        if (PyErr_Occurred()) {
            PyErr_Print();
        }
    }
}

static Object* callable_with_args(Object* ho, int narg) {
    PyObject* po = ((Py2Nrn*) ho->u.this_pointer)->po_;
    PyLockGIL lock;

    PyObject* args = PyTuple_New((Py_ssize_t) narg);
    if (args == NULL) {
        lock.release();
        hoc_execerror("PyTuple_New failed", 0);
    }
    for (int i = 0; i < narg; ++i) {
        PyObject* item = nrnpy_hoc_pop();
        if (item == NULL) {
            Py_XDECREF(args);
            lock.release();
            hoc_execerror("nrnpy_hoc_pop failed", 0);
        }
        if (PyTuple_SetItem(args, (Py_ssize_t) (narg - i - 1), item) != 0) {
            Py_XDECREF(args);
            lock.release();
            hoc_execerror("PyTuple_SetItem failed", 0);
        }
    }

    PyObject* r = PyTuple_New(2);
    PyTuple_SetItem(r, 1, args);
    Py_INCREF(po);  // when r is destroyed, do not want po refcnt to go to 0
    PyTuple_SetItem(r, 0, po);

    Object* hr = nrnpy_po2ho(r);
    Py_XDECREF(r);

    return hr;
}

static double func_call(Object* ho, int narg, int* err) {
    PyObject* po = ((Py2Nrn*) ho->u.this_pointer)->po_;
    PyObject* r;
    PyLockGIL lock;

    PyObject* args = PyTuple_New((Py_ssize_t) narg);
    if (args == NULL) {
        lock.release();
        hoc_execerror("PyTuple_New failed", 0);
    }
    for (int i = 0; i < narg; ++i) {
        PyObject* item = nrnpy_hoc_pop();
        if (item == NULL) {
            Py_XDECREF(args);
            lock.release();
            hoc_execerror("nrnpy_hoc_pop failed", 0);
        }
        if (PyTuple_SetItem(args, (Py_ssize_t) (narg - i - 1), item) != 0) {
            Py_XDECREF(args);
            lock.release();
            hoc_execerror("PyTuple_SetItem failed", 0);
        }
    }

    r = nrnpy_pyCallObject(po, args);
    Py_XDECREF(args);
    double rval = 0.0;
    if (r == NULL) {
        if (!err || *err) {
            char* mes = nrnpyerr_str();
            if (mes) {
                Fprintf(stderr, "%s\n", mes);
                free(mes);
            }
            if (PyErr_Occurred()) {
                PyErr_Print();
            }
        } else {
            PyErr_Clear();
        }
        if (!err || *err) {
            lock.release();
            hoc_execerror("func_call failed", NULL);
        }
        if (err) {
            *err = 1;
        }
    } else {
        if (nrnpy_numbercheck(r)) {
            PyObject* pn = PyNumber_Float(r);
            rval = PyFloat_AsDouble(pn);
            Py_XDECREF(pn);
        }
        Py_XDECREF(r);
        if (err) {
            *err = 0;
        }  // success
    }
    return rval;
}

static double guigetval(Object* ho) {
    PyObject* po = ((Py2Nrn*) ho->u.this_pointer)->po_;
    PyLockGIL lock;
    PyObject* r = NULL;
    PyObject* p = PyTuple_GetItem(po, 0);
    if (PySequence_Check(p) || PyMapping_Check(p)) {
        r = PyObject_GetItem(p, PyTuple_GetItem(po, 1));
    } else {
        r = PyObject_GetAttr(p, PyTuple_GetItem(po, 1));
    }
    PyObject* pn = PyNumber_Float(r);
    double x = PyFloat_AsDouble(pn);
    Py_XDECREF(pn);
    return x;
}

static void guisetval(Object* ho, double x) {
    PyObject* po = ((Py2Nrn*) ho->u.this_pointer)->po_;
    PyLockGIL lock;
    PyObject* pn = PyFloat_FromDouble(x);
    PyObject* p = PyTuple_GetItem(po, 0);
    if (PySequence_Check(p) || PyMapping_Check(p)) {
        PyObject_SetItem(p, PyTuple_GetItem(po, 1), pn);
    } else {
        PyObject_SetAttr(p, PyTuple_GetItem(po, 1), pn);
    }
    Py_XDECREF(pn);
}

static int guigetstr(Object* ho, char** cpp) {
    PyObject* po = ((Py2Nrn*) ho->u.this_pointer)->po_;
    PyLockGIL lock;

    PyObject* r = PyObject_GetAttr(PyTuple_GetItem(po, 0), PyTuple_GetItem(po, 1));
    PyObject* pn = PyObject_Str(r);
    Py2NRNString name(pn);
    Py_DECREF(pn);
    char* cp = name.c_str();
    if (*cpp && strcmp(*cpp, cp) == 0) {
        return 0;
    }
    if (*cpp) {
        delete[] * cpp;
    }
    *cpp = new char[strlen(cp) + 1];
    strcpy(*cpp, cp);
    return 1;
}

static PyObject* loads;
static PyObject* dumps;

static void setpickle() {
    PyObject* pickle;
    if (!dumps) {
        pickle = PyImport_ImportModule("pickle");
        if (pickle) {
            Py_INCREF(pickle);
            dumps = PyObject_GetAttrString(pickle, "dumps");
            loads = PyObject_GetAttrString(pickle, "loads");
            if (dumps) {
                Py_INCREF(dumps);
                Py_INCREF(loads);
            }
        }
        if (!dumps || !loads) {
            hoc_execerror("Neither Python cPickle nor pickle are available", 0);
        }
    }
}

// note that *size includes the null terminating character if it exists
static char* pickle(PyObject* p, size_t* size) {
    PyObject* arg = PyTuple_Pack(1, p);
    PyObject* r = nrnpy_pyCallObject(dumps, arg);
    Py_XDECREF(arg);
    if (!r && PyErr_Occurred()) {
        PyErr_Print();
    }
    assert(r);
    assert(PyBytes_Check(r));
    *size = PyBytes_Size(r);
    char* buf1 = PyBytes_AsString(r);
    char* buf = new char[*size];
    for (size_t i = 0; i < *size; ++i) {
        buf[i] = buf1[i];
    }
    Py_XDECREF(r);
    return buf;
}

static char* po2pickle(Object* ho, size_t* size) {
    setpickle();
    if (ho && ho->ctemplate->sym == nrnpy_pyobj_sym_) {
        PyObject* po = nrnpy_hoc2pyobject(ho);
        char* buf = pickle(po, size);
        return buf;
    } else {
        return 0;
    }
}

static PyObject* unpickle(char* s, size_t size) {
    PyObject* ps = PyBytes_FromStringAndSize(s, size);
    PyObject* arg = PyTuple_Pack(1, ps);
    PyObject* po = nrnpy_pyCallObject(loads, arg);
    assert(po);
    Py_XDECREF(arg);
    Py_XDECREF(ps);
    return po;
}

static Object* pickle2po(char* s, size_t size) {
    setpickle();
    PyObject* po = unpickle(s, size);
    Object* ho = nrnpy_pyobject_in_obj(po);
    Py_XDECREF(po);
    return ho;
}

/** Full python traceback error message returned as string.
 *  Caller should free the return value if not NULL
 **/
char* nrnpyerr_str() {
    if (PyErr_Occurred() && PyErr_ExceptionMatches(PyExc_Exception)) {
        PyObject *ptype, *pvalue, *ptraceback;
        PyErr_Fetch(&ptype, &pvalue, &ptraceback);
        PyErr_NormalizeException(&ptype, &pvalue, &ptraceback);
        // try for full backtrace
        PyObject* module_name = NULL;
        PyObject* pyth_module = NULL;
        PyObject* pyth_func = NULL;
        PyObject* py_str = NULL;
        char* cmes = NULL;

        // Since traceback.format_exception returns list of strings, wrap
        // in neuron.format_exception that returns a string.
        if (!ptraceback) {
            ptraceback = Py_None;
            Py_INCREF(ptraceback);
        }
        module_name = PyString_FromString("neuron");
        if (module_name) {
            pyth_module = PyImport_Import(module_name);
        }
        if (pyth_module) {
            pyth_func = PyObject_GetAttrString(pyth_module, "format_exception");
            if (pyth_func) {
                py_str = PyObject_CallFunctionObjArgs(pyth_func, ptype, pvalue, ptraceback, NULL);
            }
        }
        if (py_str) {
            Py2NRNString mes(py_str);
            if (mes.err()) {
                Fprintf(stderr, "nrnperr_str: Py2NRNString failed\n");
            } else {
                cmes = strdup(mes.c_str());
                if (!cmes) {
                    Fprintf(stderr, "nrnpyerr_str: strdup failed\n");
                }
            }
        }

        if (!py_str) {
            PyErr_Print();
            Fprintf(stderr, "nrnpyerr_str failed\n");
        }

        Py_XDECREF(module_name);
        Py_XDECREF(pyth_func);
        Py_XDECREF(pyth_module);
        Py_XDECREF(ptype);
        Py_XDECREF(pvalue);
        Py_XDECREF(ptraceback);
        Py_XDECREF(py_str);

        return cmes;
    }
    return NULL;
}

char* call_picklef(char* fname, size_t size, int narg, size_t* retsize) {
    // fname is a pickled callable, narg is the number of args on the
    // hoc stack with types double, char*, hoc Vector, and PythonObject
    // callable return must be pickleable.
    PyObject* args = 0;
    PyObject* result = 0;
    PyObject* callable;

    setpickle();
    PyObject* ps = PyBytes_FromStringAndSize(fname, size);
    args = PyTuple_Pack(1, ps);
    callable = nrnpy_pyCallObject(loads, args);
    assert(callable);
    Py_XDECREF(args);
    Py_XDECREF(ps);

    args = PyTuple_New(narg);
    for (int i = 0; i < narg; ++i) {
        PyObject* arg = nrnpy_hoc_pop();
        if (PyTuple_SetItem(args, narg - 1 - i, arg)) {
            assert(0);
        }
        // Py_XDECREF(arg);
    }
    result = nrnpy_pyCallObject(callable, args);
    Py_DECREF(callable);
    Py_DECREF(args);
    if (!result) {
        char* mes = nrnpyerr_str();
        if (mes) {
            Fprintf(stderr, "%s\n", mes);
            free(mes);
            hoc_execerror("PyObject method call failed:", NULL);
        }
        if (PyErr_Occurred()) {
            PyErr_Print();
        }
    }
    char* rs = pickle(result, retsize);
    Py_XDECREF(result);
    return rs;
}

#include "nrnmpi.h"

int* mk_displ(int* cnts) {
    int* displ = new int[nrnmpi_numprocs + 1];
    displ[0] = 0;
    for (int i = 0; i < nrnmpi_numprocs; ++i) {
        displ[i + 1] = displ[i] + cnts[i];
    }
    return displ;
}

static PyObject* char2pylist(char* buf, int np, int* cnt, int* displ) {
    PyObject* plist = PyList_New(np);
    assert(plist != NULL);
    for (int i = 0; i < np; ++i) {
        if (cnt[i] == 0) {
            Py_INCREF(Py_None);  // 'Fatal Python error: deallocating None' eventually
            PyList_SetItem(plist, i, Py_None);
        } else {
            PyObject* p = unpickle(buf + displ[i], cnt[i]);
            PyList_SetItem(plist, i, p);
        }
    }
    return plist;
}

#if NRNMPI
static PyObject* py_allgather(PyObject* psrc) {
    int np = nrnmpi_numprocs;
    size_t sz;
    char* sbuf = pickle(psrc, &sz);
    // what are the counts from each rank
    int* rcnt = new int[np];
    rcnt[nrnmpi_myid] = int(sz);
    nrnmpi_int_allgather_inplace(rcnt, 1);
    int* rdispl = mk_displ(rcnt);
    char* rbuf = new char[rdispl[np]];

    nrnmpi_char_allgatherv(sbuf, rbuf, rcnt, rdispl);
    delete[] sbuf;

    PyObject* pdest = char2pylist(rbuf, np, rcnt, rdispl);
    delete[] rbuf;
    delete[] rcnt;
    delete[] rdispl;
    return pdest;
}

static PyObject* py_gather(PyObject* psrc, int root) {
    int np = nrnmpi_numprocs;
    size_t sz;
    char* sbuf = pickle(psrc, &sz);
    // what are the counts from each rank
    int scnt = int(sz);
    int* rcnt = NULL;
    if (root == nrnmpi_myid) {
        rcnt = new int[np];
    }
    nrnmpi_int_gather(&scnt, rcnt, 1, root);
    int* rdispl = NULL;
    char* rbuf = NULL;
    if (root == nrnmpi_myid) {
        rdispl = mk_displ(rcnt);
        rbuf = new char[rdispl[np]];
    }

    nrnmpi_char_gatherv(sbuf, scnt, rbuf, rcnt, rdispl, root);
    delete[] sbuf;

    PyObject* pdest = Py_None;
    if (root == nrnmpi_myid) {
        pdest = char2pylist(rbuf, np, rcnt, rdispl);
        delete[] rbuf;
        delete[] rcnt;
        delete[] rdispl;
    } else {
        Py_INCREF(pdest);
    }
    return pdest;
}

static PyObject* py_broadcast(PyObject* psrc, int root) {
    // Note: root returns reffed psrc.
    char* buf = NULL;
    int cnt = 0;
    if (root == nrnmpi_myid) {
        size_t sz;
        buf = pickle(psrc, &sz);
        cnt = int(sz);
    }
    nrnmpi_int_broadcast(&cnt, 1, root);
    if (root != nrnmpi_myid) {
        buf = new char[cnt];
    }
    nrnmpi_char_broadcast(buf, cnt, root);
    PyObject* pdest = psrc;
    if (root != nrnmpi_myid) {
        pdest = unpickle(buf, size_t(cnt));
    } else {
        Py_INCREF(pdest);
    }
    delete[] buf;
    return pdest;
}
#endif

// type 1-alltoall, 2-allgather, 3-gather, 4-broadcast, 5-scatter
// size for 3, 4, 5 refer to rootrank.
Object* py_alltoall_type(int size, int type) {
    int np = nrnmpi_numprocs;  // of subworld communicator
    PyObject* psrc = NULL;
    PyObject* pdest = NULL;

    if (type == 1 || type == 5) {  // alltoall, scatter
        Object* o = *hoc_objgetarg(1);
        if (type == 1 || nrnmpi_myid == size) {  // if scatter only root must be a list
            psrc = nrnpy_hoc2pyobject(o);
            if (!PyList_Check(psrc)) {
                hoc_execerror("Argument must be a Python list", 0);
            }
            if (PyList_Size(psrc) != np) {
                if (type == 1) {
                    hoc_execerror("py_alltoall list size must be nhost", 0);
                } else {
                    hoc_execerror("py_scatter list size must be nhost", 0);
                }
            }
        }
        if (np == 1) {
            if (type == 1) {
                return o;
            } else {  // return psrc[0]
                pdest = PyList_GetItem(psrc, 0);
                Py_INCREF(pdest);
                Object* ho = nrnpy_po2ho(pdest);
                if (ho) {
                    --ho->refcount;
                }
                Py_XDECREF(pdest);
                return ho;
            }
        }
    } else {
        // Get the raw PyObject* arg. So things like None, int, bool are preserved.
        psrc = hocobj_call_arg(0);
        Py_INCREF(psrc);

        if (np == 1) {
            if (type == 4) {  // broadcast is just the PyObject
                pdest = psrc;
            } else {  // allgather and gather must wrap psrc in list
                pdest = PyList_New(1);
                PyList_SetItem(pdest, 0, psrc);
            }
            Object* ho = nrnpy_po2ho(pdest);
            if (ho) {
                --ho->refcount;
            }
            Py_XDECREF(pdest);
            return ho;
        }
    }

#if NRNMPI
    setpickle();
    int root;

    if (type == 2) {
        pdest = py_allgather(psrc);
        Py_DECREF(psrc);
    } else if (type != 1 && type != 5) {
        root = size;
        if (root < 0 || root >= np) {
            hoc_execerror("root rank must be >= 0 and < nhost", 0);
        }
        if (type == 3) {
            pdest = py_gather(psrc, root);
        } else if (type == 4) {
            pdest = py_broadcast(psrc, root);
        }
        Py_DECREF(psrc);
    } else {
        if (type == 5) {  // scatter
            root = size;
            size = 0;  // calculate dest size (cannot be -1 so cannot return it)
        }

        char* s = NULL;
        int* scnt = NULL;
        int* sdispl = NULL;
        char* r = NULL;
        int* rcnt = NULL;
        int* rdispl = NULL;

        // setup source buffer for transfer s, scnt, sdispl
        // for alltoall, each rank handled identically
        // for scatter, root handled as list all, other ranks handled as None
        if (type == 1 || nrnmpi_myid == root) {  // psrc is list of nhost items

            scnt = new int[np];
            for (int i = 0; i < np; ++i) {
                scnt[i] = 0;
            }

            PyObject* iterator = PyObject_GetIter(psrc);
            PyObject* p;

            size_t bufsz = 100000;  // 100k buffer to start with
            if (size > 0) {         // or else the positive number specified
                bufsz = size;
            }
            if (size >= 0) {  // otherwise count only
                s = new char[bufsz];
            }
            size_t curpos = 0;
            for (size_t i = 0; (p = PyIter_Next(iterator)) != NULL; ++i) {
                if (p == Py_None) {
                    scnt[i] = 0;
                    Py_DECREF(p);
                    continue;
                }
                size_t sz;
                char* b = pickle(p, &sz);
                if (size >= 0) {
                    if (curpos + sz >= bufsz) {
                        bufsz = bufsz * 2 + sz;
                        char* s2 = new char[bufsz];
                        for (size_t i = 0; i < curpos; ++i) {
                            s2[i] = s[i];
                        }
                        delete[] s;
                        s = s2;
                    }
                    for (size_t j = 0; j < sz; ++j) {
                        s[curpos + j] = b[j];
                    }
                }
                curpos += sz;
                scnt[i] = sz;
                delete[] b;
                Py_DECREF(p);
            }
            Py_DECREF(iterator);

            // scatter equivalent to alltoall NONE list for not root ranks.
        } else if (type == 5 && nrnmpi_myid != root) {
            // nothing to setup, s, scnt, sdispl already NULL
        }

        // destinations need to know receive counts. Then transfer data.
        if (type == 1) {  // alltoall

            // what are destination counts
            int* ones = new int[np];
            for (int i = 0; i < np; ++i) {
                ones[i] = 1;
            }
            sdispl = mk_displ(ones);
            rcnt = new int[np];
            nrnmpi_int_alltoallv(scnt, ones, sdispl, rcnt, ones, sdispl);
            delete[] ones;
            delete[] sdispl;

            // exchange
            sdispl = mk_displ(scnt);
            rdispl = mk_displ(rcnt);
            char* r = 0;
            if (size < 0) {
                pdest = PyTuple_New(2);
                PyTuple_SetItem(pdest, 0, Py_BuildValue("l", (long) sdispl[np]));
                PyTuple_SetItem(pdest, 1, Py_BuildValue("l", (long) rdispl[np]));
                delete[] scnt;
                delete[] sdispl;
                delete[] rcnt;
                delete[] rdispl;
            } else {
                r = new char[rdispl[np] + 1];  // force > 0 for all None case
                nrnmpi_char_alltoallv(s, scnt, sdispl, r, rcnt, rdispl);
                delete[] s;
                delete[] scnt;
                delete[] sdispl;

                pdest = char2pylist(r, np, rcnt, rdispl);

                delete[] r;
                delete[] rcnt;
                delete[] rdispl;
            }

        } else {  // scatter

            // destination counts
            rcnt = new int[1];
            nrnmpi_int_scatter(scnt, rcnt, 1, root);
            r = new char[rcnt[0] + 1];  // rcnt[0] can be 0

            // exchange
            if (nrnmpi_myid == root) {
                sdispl = mk_displ(scnt);
            }
            nrnmpi_char_scatterv(s, scnt, sdispl, r, rcnt[0], root);
            if (s)
                delete[] s;
            if (scnt)
                delete[] scnt;
            if (sdispl)
                delete[] sdispl;

            if (rcnt[0]) {
                pdest = unpickle(r, size_t(rcnt[0]));
            } else {
                pdest = Py_None;
                Py_INCREF(pdest);
            }

            delete[] r;
            delete[] rcnt;
            assert(rdispl == NULL);
        }
    }
    Object* ho = nrnpy_po2ho(pdest);
    Py_XDECREF(pdest);
    if (ho) {
        --ho->refcount;
    }
    return ho;
#else
    assert(0);
    return NULL;
#endif
}
