#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

// We don't actually need the NumPy C API for our implementation
// since we're working directly with Python lists

// Struct to represent a run of consecutive values
typedef struct {
    int start;
    int end;
    double sign;
    int length;
} Run;

// Compare function for sorting runs by length, then by start index
static int compare_runs(const void *a, const void *b) {
    const Run *run_a = (const Run *)a;
    const Run *run_b = (const Run *)b;

    if (run_a->length != run_b->length) {
        return run_a->length - run_b->length;
    }
    return run_a->start - run_b->start;
}

// Function to compute the sign of a double value (-1, 0, 1)
static double sign(double x) {
    if (x > 0) return 1.0;
    if (x < 0) return -1.0;
    return 0.0;
}

// Core implementation of filter_short_groups
static PyObject* filter_short_groups_c(PyObject *self, PyObject *args) {
    PyObject *presence_list_obj;
    int filter_size;
    char *device_id;
    char *dates_str;

    // Parse Python arguments
    if (!PyArg_ParseTuple(args, "Oiss", &presence_list_obj, &filter_size, &device_id, &dates_str)) {
        return NULL;
    }

    // Start timing
    clock_t start_time = clock();

    // Convert Python list to C array
    Py_ssize_t n = PyObject_Length(presence_list_obj);
    if (n < 0) {
        PyErr_SetString(PyExc_ValueError, "Invalid input list");
        return NULL;
    }

    // Early exit condition
    if (n == 0 || filter_size <= 1) {
        Py_INCREF(presence_list_obj);
        printf("C implementation: Early exit/no processing time: %.6fs\n",
               ((double)(clock() - start_time) / CLOCKS_PER_SEC));
        return presence_list_obj;
    }

    // Allocate memory for our working array
    double *result = (double *)malloc(n * sizeof(double));
    if (!result) {
        PyErr_SetString(PyExc_MemoryError, "Failed to allocate memory");
        return NULL;
    }

    // Copy data from Python list to C array
    for (Py_ssize_t i = 0; i < n; i++) {
        PyObject *item = PySequence_GetItem(presence_list_obj, i);
        if (!item) {
            free(result);
            return NULL;
        }
        result[i] = PyFloat_AsDouble(item);
        Py_DECREF(item);
        if (PyErr_Occurred()) {
            free(result);
            return NULL;
        }
    }

    // Allocate memory for signs array
    double *signs = (double *)malloc(n * sizeof(double));
    if (!signs) {
        free(result);
        PyErr_SetString(PyExc_MemoryError, "Failed to allocate memory for signs");
        return NULL;
    }

    // Allocate memory for change indices
    int *change_indices = (int *)malloc(n * sizeof(int));
    if (!change_indices) {
        free(result);
        free(signs);
        PyErr_SetString(PyExc_MemoryError, "Failed to allocate memory for change indices");
        return NULL;
    }

    // Allocate memory for boundaries
    int *boundaries = (int *)malloc((n + 1) * sizeof(int));
    if (!boundaries) {
        free(result);
        free(signs);
        free(change_indices);
        PyErr_SetString(PyExc_MemoryError, "Failed to allocate memory for boundaries");
        return NULL;
    }

    // Allocate memory for runs
    Run *runs = (Run *)malloc(n * sizeof(Run));
    if (!runs) {
        free(result);
        free(signs);
        free(change_indices);
        free(boundaries);
        PyErr_SetString(PyExc_MemoryError, "Failed to allocate memory for runs");
        return NULL;
    }

    // Simple cycle detection with hashset-like approach
    // We'll use a dynamically-sized array of states for simplicity
    int max_states = 100; // Initial capacity, will grow as needed
    int states_count = 0;
    double **previous_states = (double **)malloc(max_states * sizeof(double *));
    if (!previous_states) {
        free(result);
        free(signs);
        free(change_indices);
        free(boundaries);
        free(runs);
        PyErr_SetString(PyExc_MemoryError, "Failed to allocate memory for previous states");
        return NULL;
    }

    bool cycle_detected = false;

    // Main processing loop
    while (!cycle_detected) {
        // Check for cycle by comparing with previous states
        bool found_match = false;
        for (int s = 0; s < states_count; s++) {
            if (memcmp(previous_states[s], result, n * sizeof(double)) == 0) {
                found_match = true;
                cycle_detected = true;
                break;
            }
        }

        if (!found_match) {
            // Store current state
            if (states_count >= max_states) {
                // Grow the states array
                max_states *= 2;
                double **new_states = (double **)realloc(previous_states, max_states * sizeof(double *));
                if (!new_states) {
                    // Handle memory allocation failure
                    for (int i = 0; i < states_count; i++) {
                        free(previous_states[i]);
                    }
                    free(previous_states);
                    free(result);
                    free(signs);
                    free(change_indices);
                    free(boundaries);
                    free(runs);
                    PyErr_SetString(PyExc_MemoryError, "Failed to reallocate memory for previous states");
                    return NULL;
                }
                previous_states = new_states;
            }

            // Allocate memory for the new state
            previous_states[states_count] = (double *)malloc(n * sizeof(double));
            if (!previous_states[states_count]) {
                // Handle memory allocation failure
                for (int i = 0; i < states_count; i++) {
                    free(previous_states[i]);
                }
                free(previous_states);
                free(result);
                free(signs);
                free(change_indices);
                free(boundaries);
                free(runs);
                PyErr_SetString(PyExc_MemoryError, "Failed to allocate memory for new state");
                return NULL;
            }

            // Copy the current result to the new state
            memcpy(previous_states[states_count], result, n * sizeof(double));
            states_count++;

            // 1. Calculate signs
            for (Py_ssize_t i = 0; i < n; i++) {
                signs[i] = sign(result[i]);
            }

            // 2. Find indices where sign changes
            int change_count = 0;
            for (Py_ssize_t i = 1; i < n; i++) {
                if (signs[i] != signs[i-1]) {
                    change_indices[change_count++] = i;
                }
            }

            // 3. Define boundaries of consecutive runs
            boundaries[0] = 0;
            for (int i = 0; i < change_count; i++) {
                boundaries[i+1] = change_indices[i];
            }
            boundaries[change_count+1] = n;
            int boundaries_count = change_count + 2;

            // If there's only one segment, no further processing needed
            if (boundaries_count <= 2) {
                break;
            }

            // 4. Extract run properties
            int run_count = 0;
            for (int i = 0; i < boundaries_count - 1; i++) {
                int start = boundaries[i];
                int end = boundaries[i+1];
                int length = end - start;
                double run_sign = signs[start];

                // 5. Identify short runs
                if (length > 0 && length < filter_size) {
                    runs[run_count].start = start;
                    runs[run_count].end = end;
                    runs[run_count].sign = run_sign;
                    runs[run_count].length = length;
                    run_count++;
                }
            }

            // 6. Check if any modifiable short runs were found
            if (run_count == 0) {
                // No short runs found, the list is stable
                break;
            }

            // 7. Sort the short runs: shortest first, then by start index
            qsort(runs, run_count, sizeof(Run), compare_runs);

            // 8. Process the first (shortest) identified run
            Run run_to_process = runs[0];
            int start = run_to_process.start;
            int end = run_to_process.end;
            double run_sign = run_to_process.sign;

            // Determine replacement value
            double replacement_value = (run_sign == 0) ? 1.0 : 0.0;

            // 9. Apply the replacement
            for (int i = start; i < end; i++) {
                result[i] = replacement_value;
            }
            // Loop continues for the next pass
        }
    }

    // Create Python list for the result
    PyObject *result_list = PyList_New(n);
    if (!result_list) {
        // Cleanup and return
        for (int i = 0; i < states_count; i++) {
            free(previous_states[i]);
        }
        free(previous_states);
        free(result);
        free(signs);
        free(change_indices);
        free(boundaries);
        free(runs);
        return NULL;
    }

    // Fill result list
    for (Py_ssize_t i = 0; i < n; i++) {
        PyObject *item = PyFloat_FromDouble(result[i]);
        if (!item) {
            Py_DECREF(result_list);
            // Cleanup and return
            for (int i = 0; i < states_count; i++) {
                free(previous_states[i]);
            }
            free(previous_states);
            free(result);
            free(signs);
            free(change_indices);
            free(boundaries);
            free(runs);
            return NULL;
        }
        PyList_SET_ITEM(result_list, i, item);
    }

    // Print timing information
    printf("filter_short_groups_c time: %.6fs\n", ((double)(clock() - start_time) / CLOCKS_PER_SEC));

    // Cleanup
    for (int i = 0; i < states_count; i++) {
        free(previous_states[i]);
    }
    free(previous_states);
    free(result);
    free(signs);
    free(change_indices);
    free(boundaries);
    free(runs);

    return result_list;
}

// Method definitions
static PyMethodDef FilterShortGroupsMethods[] = {
    {"filter_short_groups_c", filter_short_groups_c, METH_VARARGS,
     "Optimized C version to remove groups of consecutive zeros or non-zeros shorter than filter_size."},
    {NULL, NULL, 0, NULL}
};

// Module definition
static struct PyModuleDef filter_short_groups_module = {
    PyModuleDef_HEAD_INIT,
    "filter_short_groups",
    "A module that provides optimized filter_short_groups function",
    -1,
    FilterShortGroupsMethods
};

// Module initialization
PyMODINIT_FUNC PyInit_filter_short_groups(void) {
    return PyModule_Create(&filter_short_groups_module);
}
