/**
 * Python C extension for SIMD-optimized polyphase resampling.
 *
 * Provides:
 *   - upsample(input, factor, half_length) - standard (time, channels) layout
 *   - upsample_transposed(input, factor, half_length) - transposed (channels, time) layout
 *   - get_impl_name() - return current SIMD implementation name
 *   - get_available_impls() - return list of available implementations
 *   - select_impl(name) - manually select implementation
 *   - get_cpu_features() - return detected CPU feature string
 */

#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <numpy/arrayobject.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

/* NEON intrinsics for ARM64 interleaved stores */
#if defined(__aarch64__) || defined(_M_ARM64)
#include <arm_neon.h>
#endif

#include "dispatch.h"
#include "cpu_detect.h"

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

/* ============================================================================
 * Memory allocation (replaces FFTW)
 * ============================================================================ */

static inline void* aligned_malloc(size_t size, size_t alignment) {
#if defined(_MSC_VER)
    return _aligned_malloc(size, alignment);
#elif defined(__APPLE__) || (defined(__STDC_VERSION__) && __STDC_VERSION__ >= 201112L)
    /* aligned_alloc requires size to be multiple of alignment */
    size_t aligned_size = ((size + alignment - 1) / alignment) * alignment;
    return aligned_alloc(alignment, aligned_size);
#else
    void* ptr = NULL;
    if (posix_memalign(&ptr, alignment, size) != 0) {
        return NULL;
    }
    return ptr;
#endif
}

static inline void aligned_free(void* ptr) {
#if defined(_MSC_VER)
    _aligned_free(ptr);
#else
    free(ptr);
#endif
}

/* ============================================================================
 * Kernel generation (Lanczos-windowed sinc)
 * ============================================================================ */

typedef struct {
    float* data;
    unsigned int length;
} Kernel;

/**
 * Generate polyphase decomposition of Lanczos-windowed sinc kernel.
 *
 * @param half_length Half-length of the kernel (typically 8-16)
 * @param factor Upsampling factor
 * @return Array of `factor` kernels, or NULL on error. Caller must free with free_kernels().
 */
static Kernel* generate_kernels(int half_length, int factor) {
    int kernel_length = 2 * half_length * factor + 1;
    int sub_len = 2 * half_length + 1;
    int center = kernel_length / 2;

    /* Allocate kernel array */
    Kernel* kernels = (Kernel*)malloc(sizeof(Kernel) * factor);
    if (!kernels) return NULL;

    /* Allocate master kernel with alignment for SIMD */
    float* master = (float*)aligned_malloc(kernel_length * sizeof(float), 64);
    if (!master) {
        free(kernels);
        return NULL;
    }

    /* Generate master Lanczos-windowed sinc kernel */
    for (int i = 0; i < kernel_length; i++) {
        int x = i - center;
        if (x == 0) {
            master[i] = 1.0f;
        } else {
            float xf = (float)x;
            float sinc_arg = (float)M_PI * xf / factor;
            float lanczos_arg = (float)M_PI * xf / center;
            master[i] = (sinf(sinc_arg) / sinc_arg) * (sinf(lanczos_arg) / lanczos_arg);
        }
    }

    /* Extract polyphase sub-kernels */
    for (int j = 0; j < factor; j++) {
        kernels[j].length = sub_len;
        kernels[j].data = (float*)aligned_malloc(sub_len * sizeof(float), 64);

        if (!kernels[j].data) {
            /* Cleanup on error */
            for (int k = 0; k < j; k++) {
                aligned_free(kernels[k].data);
            }
            aligned_free(master);
            free(kernels);
            return NULL;
        }

        /* Subsample master kernel with stride `factor`, reversed for convolution */
        for (int i = 0; i < sub_len; i++) {
            int idx = i * factor + j;
            if (idx < kernel_length) {
                kernels[j].data[sub_len - 1 - i] = master[idx];
            } else {
                kernels[j].data[sub_len - 1 - i] = 0.0f;
            }
        }
    }

    aligned_free(master);
    return kernels;
}

static void free_kernels(Kernel* kernels, int count) {
    if (!kernels) return;
    for (int i = 0; i < count; i++) {
        if (kernels[i].data) {
            aligned_free(kernels[i].data);
        }
    }
    free(kernels);
}

/* ============================================================================
 * Upsampling implementation - standard (time, channels) layout
 * ============================================================================ */

static PyObject* py_upsample(PyObject* self, PyObject* args, PyObject* kwargs) {
    PyArrayObject* input_array = NULL;
    int factor = 2;
    int half_length = 8;

    static char* kwlist[] = {"input", "factor", "half_length", NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O!|ii", kwlist,
                                      &PyArray_Type, &input_array,
                                      &factor, &half_length)) {
        return NULL;
    }

    /* Validate input */
    if (PyArray_NDIM(input_array) != 2) {
        PyErr_SetString(PyExc_ValueError, "Input must be a 2D array (time, channels)");
        return NULL;
    }

    if (PyArray_TYPE(input_array) != NPY_FLOAT32) {
        PyErr_SetString(PyExc_TypeError, "Input array must be float32 dtype");
        return NULL;
    }

    if (!PyArray_IS_C_CONTIGUOUS(input_array)) {
        PyErr_SetString(PyExc_ValueError, "Input array must be C-contiguous");
        return NULL;
    }

    if (factor < 2) {
        PyErr_SetString(PyExc_ValueError, "factor must be >= 2");
        return NULL;
    }

    if (half_length < 1) {
        PyErr_SetString(PyExc_ValueError, "half_length must be >= 1");
        return NULL;
    }

    npy_intp* dims = PyArray_DIMS(input_array);
    unsigned int n_samples = (unsigned int)dims[0];
    unsigned int channels = (unsigned int)dims[1];
    unsigned int kernel_len = 2 * half_length + 1;

    if (n_samples < kernel_len) {
        PyErr_SetString(PyExc_ValueError,
            "Input must have at least kernel_len samples (2*half_length+1)");
        return NULL;
    }

    /* Generate kernels */
    Kernel* kernels = generate_kernels(half_length, factor);
    if (!kernels) {
        PyErr_SetString(PyExc_MemoryError, "Failed to generate kernels");
        return NULL;
    }

    /* Calculate output size */
    unsigned int n_valid = n_samples - kernel_len + 1;
    unsigned int n_out = n_valid * factor;

    /* Create output array - use EMPTY instead of ZEROS since we write all values */
    npy_intp out_dims[2] = {n_out, channels};
    PyArrayObject* output_array = (PyArrayObject*)PyArray_EMPTY(2, out_dims, NPY_FLOAT32, 0);

    if (!output_array) {
        free_kernels(kernels, factor);
        return NULL;
    }

    float* input = (float*)PyArray_DATA(input_array);
    float* output = (float*)PyArray_DATA(output_array);

    /* Ensure dispatch is initialized */
    dispatch_init();

    /* Perform upsampling */
    Py_BEGIN_ALLOW_THREADS

    /* For standard layout, process all samples for each phase */
    for (unsigned int phase = 0; phase < (unsigned int)factor; phase++) {
        float* kernel = kernels[phase].data;

        /* Calculate number of outputs for this phase */
        unsigned int n_phase_out = (n_out - phase + factor - 1) / factor;

        /* Allocate temp buffer for contiguous output */
        float* temp = (float*)aligned_malloc(n_phase_out * channels * sizeof(float), 64);
        if (temp) {
            /* Process all samples for this phase in one call */
            g_convolve(
                temp,
                kernel,
                input,
                kernel_len,
                channels,
                n_phase_out
            );

            /* Scatter to interleaved output positions */
            for (unsigned int i = 0; i < n_phase_out; i++) {
                unsigned int out_idx = i * factor + phase;
                if (out_idx < n_out) {
                    memcpy(output + out_idx * channels,
                           temp + i * channels,
                           channels * sizeof(float));
                }
            }

            aligned_free(temp);
        }
    }

    Py_END_ALLOW_THREADS

    free_kernels(kernels, factor);
    return (PyObject*)output_array;
}

/* ============================================================================
 * Upsampling implementation - transposed (channels, time) layout
 * ============================================================================ */

static PyObject* py_upsample_transposed(PyObject* self, PyObject* args, PyObject* kwargs) {
    PyArrayObject* input_array = NULL;
    int factor = 2;
    int half_length = 8;

    static char* kwlist[] = {"input", "factor", "half_length", NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O!|ii", kwlist,
                                      &PyArray_Type, &input_array,
                                      &factor, &half_length)) {
        return NULL;
    }

    /* Validate input */
    if (PyArray_NDIM(input_array) != 2) {
        PyErr_SetString(PyExc_ValueError, "Input must be a 2D array (channels, time)");
        return NULL;
    }

    if (PyArray_TYPE(input_array) != NPY_FLOAT32) {
        PyErr_SetString(PyExc_TypeError, "Input array must be float32 dtype");
        return NULL;
    }

    if (!PyArray_IS_C_CONTIGUOUS(input_array)) {
        PyErr_SetString(PyExc_ValueError, "Input array must be C-contiguous");
        return NULL;
    }

    if (factor < 2) {
        PyErr_SetString(PyExc_ValueError, "factor must be >= 2");
        return NULL;
    }

    if (half_length < 1) {
        PyErr_SetString(PyExc_ValueError, "half_length must be >= 1");
        return NULL;
    }

    npy_intp* dims = PyArray_DIMS(input_array);
    unsigned int channels = (unsigned int)dims[0];
    unsigned int n_samples = (unsigned int)dims[1];
    unsigned int kernel_len = 2 * half_length + 1;

    if (n_samples < kernel_len) {
        PyErr_SetString(PyExc_ValueError,
            "Input must have at least kernel_len samples (2*half_length+1)");
        return NULL;
    }

    /* Generate kernels */
    Kernel* kernels = generate_kernels(half_length, factor);
    if (!kernels) {
        PyErr_SetString(PyExc_MemoryError, "Failed to generate kernels");
        return NULL;
    }

    /* Calculate output size */
    unsigned int n_valid = n_samples - kernel_len + 1;
    unsigned int n_out = n_valid * factor;

    /* Create output array - use EMPTY instead of ZEROS since we write all values */
    npy_intp out_dims[2] = {channels, n_out};
    PyArrayObject* output_array = (PyArrayObject*)PyArray_EMPTY(2, out_dims, NPY_FLOAT32, 0);

    if (!output_array) {
        free_kernels(kernels, factor);
        return NULL;
    }

    float* input = (float*)PyArray_DATA(input_array);
    float* output = (float*)PyArray_DATA(output_array);

    /* Ensure dispatch is initialized */
    dispatch_init();

    /* Build transposed kernel layout: kernel_T[k][phase] for polyphase function */
    float* kernel_T = (float*)aligned_malloc(kernel_len * factor * sizeof(float), 64);
    if (!kernel_T) {
        Py_DECREF(output_array);
        free_kernels(kernels, factor);
        PyErr_SetString(PyExc_MemoryError, "Failed to allocate kernel_T buffer");
        return NULL;
    }

    for (unsigned int k = 0; k < kernel_len; k++) {
        for (int p = 0; p < factor; p++) {
            kernel_T[k * factor + p] = kernels[p].data[k];
        }
    }

    /* Pre-allocate temp buffer for per-phase fallback */
    float* temp = (float*)aligned_malloc(n_valid * sizeof(float), 64);
    if (!temp) {
        aligned_free(kernel_T);
        Py_DECREF(output_array);
        free_kernels(kernels, factor);
        PyErr_SetString(PyExc_MemoryError, "Failed to allocate temp buffer");
        return NULL;
    }

    /* Check environment variable for method selection */
    const char* method_env = getenv("SGNL_UPSAMPLE_METHOD");
    int use_polyphase = (method_env == NULL || strcmp(method_env, "old") != 0);

    /* Perform upsampling - process each channel independently */
    Py_BEGIN_ALLOW_THREADS

    if (use_polyphase) {
        for (unsigned int ch = 0; ch < channels; ch++) {
            float* input_channel = input + ch * n_samples;
            float* output_channel = output + ch * n_out;

            g_upsample_polyphase(output_channel, kernel_T, input_channel,
                                 kernel_len, factor, n_valid);
        }
    } else {
        /* Old per-phase approach for comparison */
        for (unsigned int ch = 0; ch < channels; ch++) {
            float* input_channel = input + ch * n_samples;
            float* output_channel = output + ch * n_out;

            for (unsigned int phase = 0; phase < (unsigned int)factor; phase++) {
                float* kern = kernels[phase].data;

                g_convolve_transposed(temp, kern, input_channel, kernel_len, n_valid);

                for (unsigned int i = 0; i < n_valid; i++) {
                    output_channel[i * factor + phase] = temp[i];
                }
            }
        }
    }

    Py_END_ALLOW_THREADS

    aligned_free(temp);
    aligned_free(kernel_T);

    free_kernels(kernels, factor);
    return (PyObject*)output_array;
}

/* ============================================================================
 * Grouped correlation - correlate each channel with its own filter
 * ============================================================================ */

static PyObject* py_correlate(PyObject* self, PyObject* args, PyObject* kwargs) {
    PyArrayObject* data_array = NULL;
    PyArrayObject* filters_array = NULL;
    int groups = -1;  /* -1 means auto: use data channels */

    static char* kwlist[] = {"data", "filters", "groups", NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O!O!|i", kwlist,
                                      &PyArray_Type, &data_array,
                                      &PyArray_Type, &filters_array,
                                      &groups)) {
        return NULL;
    }

    /* Validate data input */
    if (PyArray_NDIM(data_array) != 2) {
        PyErr_SetString(PyExc_ValueError, "data must be a 2D array (channels, samples)");
        return NULL;
    }

    if (PyArray_TYPE(data_array) != NPY_FLOAT32) {
        PyErr_SetString(PyExc_TypeError, "data array must be float32 dtype");
        return NULL;
    }

    if (!PyArray_IS_C_CONTIGUOUS(data_array)) {
        PyErr_SetString(PyExc_ValueError, "data array must be C-contiguous");
        return NULL;
    }

    /* Validate filters input */
    if (PyArray_NDIM(filters_array) != 2) {
        PyErr_SetString(PyExc_ValueError, "filters must be a 2D array (n_filters, filter_length)");
        return NULL;
    }

    if (PyArray_TYPE(filters_array) != NPY_FLOAT32) {
        PyErr_SetString(PyExc_TypeError, "filters array must be float32 dtype");
        return NULL;
    }

    if (!PyArray_IS_C_CONTIGUOUS(filters_array)) {
        PyErr_SetString(PyExc_ValueError, "filters array must be C-contiguous");
        return NULL;
    }

    npy_intp* data_dims = PyArray_DIMS(data_array);
    npy_intp* filter_dims = PyArray_DIMS(filters_array);

    unsigned int n_data_channels = (unsigned int)data_dims[0];
    unsigned int n_samples = (unsigned int)data_dims[1];
    unsigned int n_filters = (unsigned int)filter_dims[0];
    unsigned int filter_len = (unsigned int)filter_dims[1];

    /* Auto-detect groups if not specified */
    if (groups < 0) {
        groups = (int)n_data_channels;
    }

    /* Validate groups parameter */
    if (groups < 1) {
        PyErr_SetString(PyExc_ValueError, "groups must be >= 1");
        return NULL;
    }

    if (n_data_channels % groups != 0) {
        PyErr_Format(PyExc_ValueError,
            "data channels (%u) must be divisible by groups (%d)",
            n_data_channels, groups);
        return NULL;
    }

    if (n_filters % groups != 0) {
        PyErr_Format(PyExc_ValueError,
            "n_filters (%u) must be divisible by groups (%d)",
            n_filters, groups);
        return NULL;
    }

    unsigned int channels_per_group = n_data_channels / groups;
    unsigned int filters_per_group = n_filters / groups;

    /* For this implementation, we require channels_per_group == 1
     * (matches PyTorch conv1d with filters shape (n_filters, 1, K)) */
    if (channels_per_group != 1) {
        PyErr_Format(PyExc_ValueError,
            "channels_per_group must be 1 (got %u). "
            "This matches PyTorch conv1d with weight shape (n_filters, 1, K)",
            channels_per_group);
        return NULL;
    }

    /* Check minimum size */
    if (n_samples < filter_len) {
        PyErr_Format(PyExc_ValueError,
            "data must have at least filter_len (%u) samples, got %u",
            filter_len, n_samples);
        return NULL;
    }

    /* Calculate output size */
    unsigned int n_out = n_samples - filter_len + 1;

    /* Output has n_filters channels */
    npy_intp out_dims[2] = {n_filters, n_out};
    PyArrayObject* output_array = (PyArrayObject*)PyArray_EMPTY(2, out_dims, NPY_FLOAT32, 0);

    if (!output_array) {
        return NULL;
    }

    float* data = (float*)PyArray_DATA(data_array);
    float* filters = (float*)PyArray_DATA(filters_array);
    float* output = (float*)PyArray_DATA(output_array);

    /* Ensure dispatch is initialized */
    dispatch_init();

    Py_BEGIN_ALLOW_THREADS

    /*
     * Grouped convolution matching PyTorch conv1d with groups parameter.
     *
     * With groups=g, n_data_channels=C_in, n_filters=C_out:
     * - Input is divided into g groups, each with C_in/g channels
     * - Filters are divided into g groups, each with C_out/g filters
     * - Each filter group operates only on its corresponding input group
     *
     * Since channels_per_group=1, each group has 1 input channel and
     * filters_per_group output channels.
     */
    if (groups == 1 && g_convolve_multi_filter) {
        /*
         * Optimized path for groups=1: apply all filters to single input channel.
         * Uses specialized kernel that keeps input data in L1 cache by processing
         * all filters for each output position before moving to the next.
         */
        g_convolve_multi_filter(
            output,
            filters,
            data,
            n_filters,
            filter_len,
            n_out
        );
    } else {
        /* General grouped convolution path */
        for (int g = 0; g < groups; g++) {
            /* Input channel for this group */
            float* data_channel = data + g * n_samples;

            /* Filters for this group */
            unsigned int filter_start = g * filters_per_group;

            /* Apply each filter in this group to the input channel */
            for (unsigned int f = 0; f < filters_per_group; f++) {
                unsigned int filter_idx = filter_start + f;
                float* filter = filters + filter_idx * filter_len;
                float* output_channel = output + filter_idx * n_out;

                g_convolve_transposed(
                    output_channel,
                    filter,
                    data_channel,
                    filter_len,
                    n_out
                );
            }
        }
    }

    Py_END_ALLOW_THREADS

    return (PyObject*)output_array;
}

/* ============================================================================
 * Introspection functions
 * ============================================================================ */

static PyObject* py_get_impl_name(PyObject* self, PyObject* args) {
    dispatch_init();
    return PyUnicode_FromString(g_impl_name ? g_impl_name : "Unknown");
}

static PyObject* py_get_available_impls(PyObject* self, PyObject* args) {
    const char** impls = dispatch_get_available();

    /* Count implementations */
    int count = 0;
    while (impls[count]) count++;

    PyObject* list = PyList_New(count);
    if (!list) return NULL;

    for (int i = 0; i < count; i++) {
        PyObject* name = PyUnicode_FromString(impls[i]);
        if (!name) {
            Py_DECREF(list);
            return NULL;
        }
        PyList_SET_ITEM(list, i, name);
    }

    return list;
}

static PyObject* py_select_impl(PyObject* self, PyObject* args) {
    const char* name;

    if (!PyArg_ParseTuple(args, "s", &name)) {
        return NULL;
    }

    int result = dispatch_select(name);

    if (result == -1) {
        PyErr_Format(PyExc_ValueError, "Unknown implementation: %s", name);
        return NULL;
    } else if (result == -2) {
        PyErr_Format(PyExc_RuntimeError,
            "CPU features required for '%s' are not available", name);
        return NULL;
    }

    Py_RETURN_NONE;
}

static PyObject* py_get_cpu_features(PyObject* self, PyObject* args) {
    unsigned int features = dispatch_get_cpu_features();
    return PyUnicode_FromString(cpu_features_string(features));
}

/* ============================================================================
 * Module definition
 * ============================================================================ */

static PyMethodDef module_methods[] = {
    {"upsample", (PyCFunction)py_upsample, METH_VARARGS | METH_KEYWORDS,
     "Upsample multichannel data using polyphase filtering.\n\n"
     "Parameters\n"
     "----------\n"
     "input : ndarray (float32)\n"
     "    Input array of shape (n_samples, n_channels). Must be C-contiguous.\n"
     "factor : int, optional\n"
     "    Upsampling factor (default: 2)\n"
     "half_length : int, optional\n"
     "    Half-length of the sinc kernel (default: 8)\n\n"
     "Returns\n"
     "-------\n"
     "output : ndarray (float32)\n"
     "    Upsampled array of shape ((n_samples - kernel_len + 1) * factor, n_channels)\n"
    },
    {"upsample_transposed", (PyCFunction)py_upsample_transposed, METH_VARARGS | METH_KEYWORDS,
     "Upsample multichannel data using polyphase filtering (transposed layout).\n\n"
     "Parameters\n"
     "----------\n"
     "input : ndarray (float32)\n"
     "    Input array of shape (n_channels, n_samples). Must be C-contiguous.\n"
     "factor : int, optional\n"
     "    Upsampling factor (default: 2)\n"
     "half_length : int, optional\n"
     "    Half-length of the sinc kernel (default: 8)\n\n"
     "Returns\n"
     "-------\n"
     "output : ndarray (float32)\n"
     "    Upsampled array of shape (n_channels, (n_samples - kernel_len + 1) * factor)\n"
    },
    {"correlate", (PyCFunction)py_correlate, METH_VARARGS | METH_KEYWORDS,
     "Grouped 1D correlation - correlate each channel with its own filter.\n\n"
     "This is equivalent to torch.nn.functional.conv1d with groups=channels,\n"
     "but optimized for CPU with SIMD instructions.\n\n"
     "Parameters\n"
     "----------\n"
     "data : ndarray (float32)\n"
     "    Input array of shape (channels, samples). Must be C-contiguous.\n"
     "filters : ndarray (float32)\n"
     "    Filter array of shape (channels, filter_length). Must be C-contiguous.\n"
     "    Each channel in data is correlated with the corresponding filter.\n\n"
     "Returns\n"
     "-------\n"
     "output : ndarray (float32)\n"
     "    Output array of shape (channels, samples - filter_length + 1)\n"
    },
    {"get_impl_name", py_get_impl_name, METH_NOARGS,
     "Get the name of the current SIMD implementation.\n\n"
     "Returns\n"
     "-------\n"
     "name : str\n"
     "    Implementation name (e.g., 'AVX2+FMA', 'NEON', 'Scalar')\n"
    },
    {"get_available_impls", py_get_available_impls, METH_NOARGS,
     "Get list of available SIMD implementations for this CPU.\n\n"
     "Returns\n"
     "-------\n"
     "impls : list of str\n"
     "    List of implementation names that can be used with select_impl()\n"
    },
    {"select_impl", py_select_impl, METH_VARARGS,
     "Manually select a SIMD implementation.\n\n"
     "Parameters\n"
     "----------\n"
     "name : str\n"
     "    Implementation name (from get_available_impls())\n\n"
     "Raises\n"
     "------\n"
     "ValueError\n"
     "    If implementation name is not recognized\n"
     "RuntimeError\n"
     "    If CPU features required for the implementation are not available\n"
    },
    {"get_cpu_features", py_get_cpu_features, METH_NOARGS,
     "Get detected CPU SIMD features.\n\n"
     "Returns\n"
     "-------\n"
     "features : str\n"
     "    Space-separated list of detected features\n"
    },
    {NULL, NULL, 0, NULL}
};

static struct PyModuleDef module_def = {
    PyModuleDef_HEAD_INIT,
    "_core",
    "SIMD-optimized polyphase resampling with runtime CPU detection",
    -1,
    module_methods
};

PyMODINIT_FUNC PyInit__core(void) {
    /* Initialize NumPy */
    import_array();

    /* Initialize dispatch on module load */
    dispatch_init();

    return PyModule_Create(&module_def);
}
