import numpy as np
import time
import ctypes
import subprocess, tempfile, os, sys

c_code = """
#include <math.h>
#include <stdlib.h>
#include <string.h>

#ifdef _OPENMP
#include <omp.h>
#endif

static inline double point_to_segment_dist_sq(
    const double* p, const double* a, const double* b, int dim
) {
    double ab_sq = 0.0;
    double ap_dot_ab = 0.0;
    for (int i = 0; i < dim; ++i) {
        double d_ab = b[i] - a[i];
        double d_ap = p[i] - a[i];
        ab_sq += d_ab * d_ab;
        ap_dot_ab += d_ap * d_ab;
    }
    if (ab_sq <= 1e-14) {
        double d_sq = 0.0;
        for (int i = 0; i < dim; ++i) {
            double d = p[i] - a[i];
            d_sq += d * d;
        }
        return d_sq;
    }
    double t = ap_dot_ab / ab_sq;
    if (t < 0.0) t = 0.0;
    else if (t > 1.0) t = 1.0;
    
    double dist_sq = 0.0;
    for (int i = 0; i < dim; ++i) {
        double proj = a[i] + t * (b[i] - a[i]);
        double d = p[i] - proj;
        dist_sq += d * d;
    }
    return dist_sq;
}

int rdp_simplify_polyline_c(
    const double* in_pts,
    int num_pts,
    int stride,
    double tol,
    int preserve_topology,
    double* out_pts
) {
    if (num_pts <= 2 || tol <= 0.0) {
        memcpy(out_pts, in_pts, num_pts * stride * sizeof(double));
        return num_pts;
    }

    double tol_sq = tol * tol;
    char* keep = (char*)calloc(num_pts, sizeof(char));
    if (!keep) {
        memcpy(out_pts, in_pts, num_pts * stride * sizeof(double));
        return num_pts;
    }

    keep[0] = 1;
    keep[num_pts - 1] = 1;

    // Check if closed loop
    int is_closed = 0;
    double d_close_sq = 0.0;
    for (int s = 0; s < 3 && s < stride; ++s) {
        double d = in_pts[s] - in_pts[(num_pts - 1) * stride + s];
        d_close_sq += d * d;
    }
    if (d_close_sq < 1e-12) {
        is_closed = 1;
    }

    // Stack for iterative RDP: pairs of (start_idx, end_idx)
    int stack_cap = 256;
    int stack_size = 0;
    int stack_buf[512];
    int* stack = stack_buf;
    int heap_alloc = 0;

    if (is_closed) {
        double max_d_sq = -1.0;
        int max_idx = -1;
        const double* p0 = in_pts;
        for (int i = 1; i < num_pts - 1; ++i) {
            double d_sq = 0.0;
            for (int s = 0; s < 3 && s < stride; ++s) {
                double diff = in_pts[i * stride + s] - p0[s];
                d_sq += diff * diff;
            }
            if (d_sq > max_d_sq) {
                max_d_sq = d_sq;
                max_idx = i;
            }
        }
        if (max_idx > 0 && max_d_sq > tol_sq) {
            keep[max_idx] = 1;
            stack[stack_size * 2 + 0] = 0;
            stack[stack_size * 2 + 1] = max_idx;
            stack_size++;
            stack[stack_size * 2 + 0] = max_idx;
            stack[stack_size * 2 + 1] = num_pts - 1;
            stack_size++;
        }
    } else {
        stack[stack_size * 2 + 0] = 0;
        stack[stack_size * 2 + 1] = num_pts - 1;
        stack_size++;
    }

    while (stack_size > 0) {
        stack_size--;
        int start = stack[stack_size * 2 + 0];
        int end = stack[stack_size * 2 + 1];

        if (end - start <= 1) continue;

        const double* a = in_pts + (start * stride);
        const double* b = in_pts + (end * stride);

        double max_d_sq = 0.0;
        int max_idx = -1;

        for (int i = start + 1; i < end; ++i) {
            const double* p = in_pts + (i * stride);
            double d_sq = point_to_segment_dist_sq(p, a, b, (stride >= 3 ? 3 : stride));
            if (d_sq > max_d_sq) {
                max_d_sq = d_sq;
                max_idx = i;
            }
        }

        if (max_d_sq > tol_sq && max_idx > 0) {
            keep[max_idx] = 1;
            if (stack_size + 2 >= stack_cap) {
                stack_cap *= 2;
                if (!heap_alloc) {
                    stack = (int*)malloc(stack_cap * 2 * sizeof(int));
                    memcpy(stack, stack_buf, stack_size * 2 * sizeof(int));
                    heap_alloc = 1;
                } else {
                    stack = (int*)realloc(stack, stack_cap * 2 * sizeof(int));
                }
            }
            stack[stack_size * 2 + 0] = start;
            stack[stack_size * 2 + 1] = max_idx;
            stack_size++;
            stack[stack_size * 2 + 0] = max_idx;
            stack[stack_size * 2 + 1] = end;
            stack_size++;
        }
    }

    if (heap_alloc) {
        free(stack);
    }

    if (preserve_topology && is_closed) {
        int kept_count = 0;
        for (int i = 0; i < num_pts; ++i) if (keep[i]) kept_count++;
        if (kept_count < 4 && num_pts >= 4) {
            keep[num_pts / 3] = 1;
            keep[(2 * num_pts) / 3] = 1;
        }
    }

    int out_count = 0;
    for (int i = 0; i < num_pts; ++i) {
        if (keep[i]) {
            memcpy(out_pts + (out_count * stride), in_pts + (i * stride), stride * sizeof(double));
            out_count++;
        }
    }

    free(keep);
    return out_count;
}

// Batch simplify with OpenMP parallelization
void rdp_simplify_batch_c(
    const double** in_pts_arr,
    const int* num_pts_arr,
    const int* stride_arr,
    int num_paths,
    double tol,
    int preserve_topology,
    double** out_pts_arr,
    int* out_counts_arr
) {
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < num_paths; ++i) {
        out_counts_arr[i] = rdp_simplify_polyline_c(
            in_pts_arr[i],
            num_pts_arr[i],
            stride_arr[i],
            tol,
            preserve_topology,
            out_pts_arr[i]
        );
    }
}
"""

with tempfile.NamedTemporaryFile("w", suffix=".c", delete=False) as f:
    f.write(c_code)
    c_path = f.name

so_path = c_path.replace(".c", ".so")
subprocess.run(["gcc", "-O3", "-fopenmp", "-fPIC", "-shared", c_path, "-o", so_path], check=True)
lib = ctypes.CDLL(so_path)
lib.rdp_simplify_batch_c.argtypes = [
    ctypes.POINTER(ctypes.POINTER(ctypes.c_double)),
    ctypes.POINTER(ctypes.c_int),
    ctypes.POINTER(ctypes.c_int),
    ctypes.c_int,
    ctypes.c_double,
    ctypes.c_int,
    ctypes.POINTER(ctypes.POINTER(ctypes.c_double)),
    ctypes.POINTER(ctypes.c_int)
]

sys.path.append(os.path.dirname(__file__))
from benchmark_simplify import sp

paths_c = [np.ascontiguousarray(p, dtype=np.float64) for p in sp]
num_paths = len(paths_c)

# Setup pointers
in_ptrs = (ctypes.POINTER(ctypes.c_double) * num_paths)()
out_ptrs = (ctypes.POINTER(ctypes.c_double) * num_paths)()
num_pts = (ctypes.c_int * num_paths)()
strides = (ctypes.c_int * num_paths)()
out_counts = (ctypes.c_int * num_paths)()

out_bufs = [np.empty_like(p) for p in paths_c]
for i in range(num_paths):
    in_ptrs[i] = paths_c[i].ctypes.data_as(ctypes.POINTER(ctypes.c_double))
    out_ptrs[i] = out_bufs[i].ctypes.data_as(ctypes.POINTER(ctypes.c_double))
    num_pts[i] = len(paths_c[i])
    strides[i] = paths_c[i].shape[1]

t0 = time.perf_counter()
lib.rdp_simplify_batch_c(
    in_ptrs, num_pts, strides, num_paths, 0.05, 1, out_ptrs, out_counts
)
res = [out_bufs[i][:out_counts[i]].copy() for i in range(num_paths)]
t1 = time.perf_counter()

print(f"OpenMP Batch C RDP simplify: {t1 - t0:.5f} s, output points: {sum(len(p) for p in res)}")