// #include <cheri_init_globals.h>
// #include <cheritypes.h>
#include <cheriintrin.h>
#include <stdint.h>
#include <stddef.h>
#include <stdint.h>
#include <float.h>
// #include <math.h>

/* --- Simple Bare-Metal CHERI Malloc --- */

// #define HEAP_SIZE 0x10000
// static uint8_t raw_heap[HEAP_SIZE] __attribute__((aligned(16)));
// static size_t heap_ptr = 0;

// void* cheri_malloc(size_t size) {
//     // Ensure 16-byte alignment for CHERI capability representability
//     size = (size + 15) & ~15;
    
//     if (heap_ptr + size > HEAP_SIZE) return NULL;

//     // Get a capability to the heap slice
//     void *ptr = cheri_get_base(&raw_heap[heap_ptr]);
    
//     // Set strict bounds on the returned capability
//     ptr = cheri_set_bounds(ptr, size);
    
//     heap_ptr += size;
//     return ptr;
// }

void free(void * __capability ptr);
void * __capability malloc(size_t size);

/* --- K-Means Hybrid Structures --- */

typedef struct {
    float x, y;
} Point;

typedef struct {
    float x, y;
    int count;
} Centroid;

float get_distance(Point p, Centroid c) {
    float dx = p.x - c.x;
    float dy = p.y - c.y;
    return (dx * dx) + (dy * dy); // Squared Euclidean
}

/* --- Main Logic --- */

void run_kmeans(Point * __capability points, int num_points, int k, int iterations) {
    // Allocate centroids using our CHERI-bounded malloc
    Centroid * __capability centroids = (Centroid * __capability)malloc(sizeof(Centroid) * k);
    int * __capability assignments = (int * __capability)malloc(sizeof(int) * num_points);

    if (!centroids || !assignments) return;

    // Initialize Centroids (Simple sequential pick)
    for (int i = 0; i < k; i++) {
        centroids[i].x = points[i].x;
        centroids[i].y = points[i].y;
    }

    for (int iter = 0; iter < iterations; iter++) {
        // 1. Assignment Step
        for (int i = 0; i < num_points; i++) {
            float min_dist = FLT_MAX;
            int best_cluster = 0;
            for (int j = 0; j < k; j++) {
                float d = get_distance(points[i], centroids[j]);
                if (d < min_dist) {
                    min_dist = d;
                    best_cluster = j;
                }
            }
            assignments[i] = best_cluster;
        }

        // 2. Update Step
        for (int i = 0; i < k; i++) {
            centroids[i].x = 0;
            centroids[i].y = 0;
            centroids[i].count = 0;
        }

        for (int i = 0; i < num_points; i++) {
            int cluster = assignments[i];
            centroids[cluster].x += points[i].x;
            centroids[cluster].y += points[i].y;
            centroids[cluster].count++;
        }

        for (int i = 0; i < k; i++) {
            if (centroids[i].count > 0) {
                centroids[i].x /= centroids[i].count;
                centroids[i].y /= centroids[i].count;
            }
        }
    }
}

int main() {
    int n = 10;
    int k = 3;

    // Allocate point data on our CHERI heap
    Point * __capability data = (Point * __capability)malloc(sizeof(Point) * n);
    
    if (data) {
        // Mock data initialization
        for(int i = 0; i < n; i++) {
            data[i].x = (float)(i % 10);
            data[i].y = (float)(i / 10);
        }

        run_kmeans(data, n, k, 10);
    }

    return 0;
}