import numpy as np
import time
import ctypes

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

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 = (int*)malloc(stack_cap * 2 * sizeof(int));

    if (is_closed) {
        // For closed loop, find point furthest from start point
        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;
                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++;
        }
    }

    free(stack);

    // If preserve_topology is requested on closed loop and fewer than 4 points kept:
    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 at least 3 distinct vertices (e.g. 0, num_pts/3, 2*num_pts/3, num_pts-1)
            keep[num_pts / 3] = 1;
            keep[(2 * num_pts) / 3] = 1;
        }
    }

    // Collect output points
    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;
}
"""

import subprocess, tempfile, os

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", "-fPIC", "-shared", c_path, "-o", so_path], check=True)
lib = ctypes.CDLL(so_path)
lib.rdp_simplify_polyline_c.argtypes = [
    ctypes.POINTER(ctypes.c_double),
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_double,
    ctypes.c_int,
    ctypes.POINTER(ctypes.c_double)
]
lib.rdp_simplify_polyline_c.restype = ctypes.c_int

def simplify_c(p: np.ndarray, tol: float, preserve_topology: bool = True) -> np.ndarray:
    if len(p) <= 2 or tol <= 0:
        return p
    p_c = np.ascontiguousarray(p, dtype=np.float64)
    out_buf = np.empty_like(p_c)
    n_out = lib.rdp_simplify_polyline_c(
        p_c.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
        len(p_c),
        p_c.shape[1],
        float(tol),
        1 if preserve_topology else 0,
        out_buf.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
    )
    return out_buf[:n_out].copy()

# Test benchmark on dataset
import sys
sys.path.append(os.path.dirname(__file__))
from benchmark_simplify import sp

print(f"Testing on {len(sp)} paths...")
t0 = time.perf_counter()
res_c = [simplify_c(p, 0.05, preserve_topology=True) for p in sp]
t1 = time.perf_counter()
print(f"C RDP simplify: {t1 - t0:.4f} s, output points: {sum(len(p) for p in res_c)}")

# Check accuracy/geometry similarity
from shapely.geometry import LineString
p_sample = sp[0]
p_shapely = np.array(LineString(p_sample).simplify(0.05, preserve_topology=True).coords)
p_c_sample = simplify_c(p_sample, 0.05, preserve_topology=True)

print(f"Sample path: original {len(p_sample)} pts -> Shapely: {len(p_shapely)} pts, C: {len(p_c_sample)} pts")