#define PY_SSIZE_T_CLEAN
#include "Python.h"
#include "numpy/arrayobject.h"
#include <math.h>
#include "src/lbfgsb.h"

#define PYERR(errobj,message) {PyErr_SetString(errobj,message); return NULL;}

static PyObject* lbfgsb_error;

static char doc_setulb[] = (
    "setulb(m, x, l, u, nbd, f, g, factr, pgtol, wa, iwa, task, lsave, isave, dsave, maxls, ln_task)\n\n"
    "Limited-memory Broyden-Fletcher-Goldfarb-Shanno with bounds.\n\n"
    "Low-level interface to the L-BFGS-B optimization algorithm. "
    "Performs a single iteration step.\n\n"
    "Parameters\n"
    "----------\n"
    "m : int\n"
    "    Number of limited-memory variable metric corrections.\n"
    "x : ndarray, shape (n,)\n"
    "    Current solution vector, float64. Modified in-place.\n"
    "l : ndarray, shape (n,)\n"
    "    Lower bounds on variables, float64.\n"
    "u : ndarray, shape (n,)\n"
    "    Upper bounds on variables, float64.\n"
    "nbd : ndarray, shape (n,)\n"
    "    Bound type for each variable, int32. 0: unbounded, 1: only lower, "
    "2: both, 3: only upper.\n"
    "f : float\n"
    "    Current objective function value.\n"
    "g : ndarray, shape (n,)\n"
    "    Gradient of objective function at current point, float64.\n"
    "factr : float\n"
    "    Factor for convergence test on objective function.\n"
    "pgtol : float\n"
    "    Projected gradient tolerance for convergence.\n"
    "wa : ndarray\n"
    "    Workspace array for internal computations, float64.\n"
    "iwa : ndarray\n"
    "    Integer workspace array for internal computations, int32.\n"
    "task : ndarray\n"
    "    Task control array encoding algorithm state, int32. Modified in-place.\n"
    "lsave : ndarray\n"
    "    Logical state array, int32. Modified in-place.\n"
    "isave : ndarray\n"
    "    Integer state array, int32. Modified in-place.\n"
    "dsave : ndarray\n"
    "    Real state array, float64. Modified in-place.\n"
    "maxls : int\n"
    "    Maximum number of line search steps per iteration.\n"
    "ln_task : ndarray\n"
    "    Task string length array, int32.\n\n"
    "Returns\n"
    "-------\n"
    "None\n"
    "    Modifies x, task, lsave, isave, and dsave in-place.\n\n"
);

static PyObject*
lbfgsb_setulb(PyObject *self, PyObject *args)
{

    Py_ssize_t maxls_input, m_input;
    CBLAS_INT m, n, maxls;
    double f, factr, pgtol;
    CBLAS_INT *nbd, *iwa, *isave, *lsave, *taskptr, *ln_taskptr;
    double *x, *l, *u, *g, *wa, *dsave;

    PyArrayObject *ap_x=NULL, *ap_l=NULL, *ap_u=NULL, *ap_g=NULL, *ap_nbd=NULL;
    PyArrayObject *ap_wa=NULL, *ap_iwa=NULL, *ap_task=NULL, *ap_ln_task=NULL;
    PyArrayObject *ap_lsave=NULL, *ap_isave=NULL, *ap_dsave=NULL;

    // Check argument types
    // m,x,l,u,nbd,f,g,factr,pgtol,wa,iwa,task,lsave,isave,dsave,maxls,ln_task
    if (!(PyArg_ParseTuple(args,
                           "nO!O!O!O!dO!ddO!O!O!O!O!O!nO!",
                           &m_input,                                // n
                           &PyArray_Type, (PyObject **)&ap_x,       // O!
                           &PyArray_Type, (PyObject **)&ap_l,       // O!
                           &PyArray_Type, (PyObject **)&ap_u,       // O!
                           &PyArray_Type, (PyObject **)&ap_nbd,     // O!
                           &f,                                      // d
                           &PyArray_Type, (PyObject **)&ap_g,       // O!
                           &factr,                                  // d
                           &pgtol,                                  // d
                           &PyArray_Type, (PyObject **)&ap_wa,      // O!
                           &PyArray_Type, (PyObject **)&ap_iwa,     // O!
                           &PyArray_Type, (PyObject **)&ap_task,    // O!
                           &PyArray_Type, (PyObject **)&ap_lsave,   // O!
                           &PyArray_Type, (PyObject **)&ap_isave,   // O!
                           &PyArray_Type, (PyObject **)&ap_dsave,   // O!
                           &maxls_input,                            // n
                           &PyArray_Type, (PyObject **)&ap_ln_task  // O!
                           ))) { return NULL; }

    maxls = (CBLAS_INT)maxls_input;
    m = (CBLAS_INT)m_input;
    int tcode = sizeof(CBLAS_INT) == sizeof(npy_int32) ? NPY_INT32 : NPY_INT64;

    // Check if arrays are contiguous and with the right dtype.
    // All arrays are 1D hence both F and C contiguous flags are set to True.
    if (!(PyArray_IS_C_CONTIGUOUS(ap_x)) || (PyArray_TYPE(ap_x) != NPY_FLOAT64))
    {
        PYERR(lbfgsb_error, " Argument (x) must be a contiguous array of type float64.");
    }
    if (!(PyArray_IS_C_CONTIGUOUS(ap_l)) || (PyArray_TYPE(ap_l) != NPY_FLOAT64))
    {
        PYERR(lbfgsb_error, " Argument (l) must be a contiguous array of type float64.");
    }
    if (!(PyArray_IS_C_CONTIGUOUS(ap_u)) || (PyArray_TYPE(ap_u) != NPY_FLOAT64))
    {
        PYERR(lbfgsb_error, " Argument (u) must be a contiguous array of type float64.");
    }
    if (!(PyArray_IS_C_CONTIGUOUS(ap_wa)) || (PyArray_TYPE(ap_wa) != NPY_FLOAT64))
    {
        PYERR(lbfgsb_error, " Argument (wa) must be a contiguous array of type float64.");
    }
    if (!(PyArray_IS_C_CONTIGUOUS(ap_iwa)) || (PyArray_TYPE(ap_iwa) != tcode))
    {
        PYERR(lbfgsb_error, " Argument (iwa) must be a contiguous array of type int32.");
    }
    if (!(PyArray_IS_C_CONTIGUOUS(ap_dsave)) || (PyArray_TYPE(ap_dsave) != NPY_FLOAT64))
    {
        PYERR(lbfgsb_error, " Argument (dsave) must be a contiguous array of type float64.");
    }
    if (!(PyArray_IS_C_CONTIGUOUS(ap_nbd)) || (PyArray_TYPE(ap_nbd) != tcode))
    {
        PYERR(lbfgsb_error, " Argument (nbd) must be a contiguous array of type int32.");
    }
    if (!(PyArray_IS_C_CONTIGUOUS(ap_task)) || (PyArray_TYPE(ap_task) != tcode))
    {
        PYERR(lbfgsb_error, " Argument (task) must be a contiguous array of type int32.");
    }
    if (!(PyArray_IS_C_CONTIGUOUS(ap_lsave)) || (PyArray_TYPE(ap_lsave) != tcode))
    {
        PYERR(lbfgsb_error, " Argument (lsave) must be a contiguous array of type int32.");
    }
    if (!(PyArray_IS_C_CONTIGUOUS(ap_isave)) || (PyArray_TYPE(ap_isave) != tcode))
    {
        PYERR(lbfgsb_error, " Argument (isave) must be a contiguous array of type int32.");
    }
    if (!(PyArray_IS_C_CONTIGUOUS(ap_ln_task)) || (PyArray_TYPE(ap_ln_task) != tcode))
    {
        PYERR(lbfgsb_error, " Argument (ln_task) must be a contiguous array of type int32.");
    }
    n = PyArray_DIMS(ap_x)[0];

    x = (double *)PyArray_DATA(ap_x);
    l = (double *)PyArray_DATA(ap_l);
    u = (double *)PyArray_DATA(ap_u);
    nbd = (CBLAS_INT *)PyArray_DATA(ap_nbd);
    g = (double *)PyArray_DATA(ap_g);
    wa = (double *)PyArray_DATA(ap_wa);
    iwa = (CBLAS_INT *)PyArray_DATA(ap_iwa);
    lsave = (CBLAS_INT *)PyArray_DATA(ap_lsave);
    isave = (CBLAS_INT *)PyArray_DATA(ap_isave);
    dsave = (double *)PyArray_DATA(ap_dsave);
    taskptr = (CBLAS_INT *)PyArray_DATA(ap_task);
    ln_taskptr = (CBLAS_INT *)PyArray_DATA(ap_ln_task);

    // Make the function call
    setulb(n, m, x, l, u, nbd, &f, g, factr, pgtol, wa, iwa, taskptr, lsave,
           isave, dsave, maxls, ln_taskptr);

    Py_RETURN_NONE;
}


static struct PyMethodDef lbfgsb_module_methods[] = {
  {"setulb", lbfgsb_setulb, METH_VARARGS, doc_setulb},
  {NULL,     NULL,          0,            NULL}
};


static int module_exec(PyObject *module) {

    if (_import_array() < 0) { return -1; }

    lbfgsb_error = PyErr_NewException("_lbfgsb.error", NULL, NULL);
    if (lbfgsb_error == NULL) { return -1; }

    if (PyModule_AddObject(module, "error", lbfgsb_error) < 0) {
        Py_DECREF(lbfgsb_error);
        return -1;
    }
#if Py_GIL_DISABLED
    PyUnstable_Module_SetGIL(module, Py_MOD_GIL_NOT_USED);
#endif
    return 0;
}


static struct PyModuleDef_Slot _lbfgsb_slots[] = {
    {Py_mod_exec, module_exec},
    {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
#if PY_VERSION_HEX >= 0x030d00f0  // Python 3.13+
    // signal that this module supports running without an active GIL
    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
#endif
    {0, NULL},
};


static struct PyModuleDef moduledef = {
    .m_base = PyModuleDef_HEAD_INIT,
    .m_name = "_lbfgsb",
    .m_size = 0,
    .m_methods = lbfgsb_module_methods,
    .m_slots = _lbfgsb_slots
};


PyMODINIT_FUNC
PyInit__lbfgsb(void)
{
    return PyModuleDef_Init(&moduledef);
}
