/*
 reinhard@finalmedia.de
 Do 9. Jul 22:24:46 CEST 2026
 Public Domain

 ply2fsplat.c
 liest ply files von stdin und schreibt fsplat files auf stdout

 musl-gcc -O3 -march=native -static-pie -fPIE ply2fsplat.c -o ply2fsplat -lm

*/

#include <math.h>
#include <stdint.h>
#include <unistd.h>

#define MAX_SPLATS 8500000
#define MAX_PATTERNS 256
#define KMEANS_ITERATIONS 15

typedef struct { float c[48]; } SHData;

typedef struct {
    float x, y, z, sx, sy, sz;
    unsigned char opacity, q0, q1, q2, q3;
    unsigned char sh_idx;
} __attribute__((packed)) FSplat;


float xyz_scale[MAX_SPLATS * 6];
float opacity_rot[MAX_SPLATS * 5];
SHData sh_data[MAX_SPLATS];
uint8_t labels[MAX_SPLATS];

float centroids[MAX_PATTERNS][48];
float new_centroids[MAX_PATTERNS * 48];
uint32_t counts[MAX_PATTERNS];

#define OUT_BUF_SPLAT_COUNT 8192
FSplat out_buf[OUT_BUF_SPLAT_COUNT];

static inline void byte_copy(void* dst, const void* src, size_t n) {
    char* d = (char*)dst;
    const char* s = (const char*)src;
    while (n--) *d++ = *s++;
}

static inline void byte_zero(void* dst, size_t n) {
    char* d = (char*)dst;
    while (n--) *d++ = 0;
}

static inline int str_start(const char* s, const char* prefix) {
    while (*prefix) {
        if (*s++ != *prefix++) return 0;
    }
    return 1;
}

void run_kmeans(uint32_t n_splats) {
    for (int i = 0; i < MAX_PATTERNS; i++) {
        uint32_t idx = (i * (n_splats / MAX_PATTERNS)) % n_splats;
        byte_copy(centroids[i], sh_data[idx].c, 48 * sizeof(float));
    }

    for (int iter = 0; iter < KMEANS_ITERATIONS; iter++) {
        byte_zero(new_centroids, MAX_PATTERNS * 48 * sizeof(float));
        byte_zero(counts, MAX_PATTERNS * sizeof(uint32_t));

        for (uint32_t i = 0; i < n_splats; i++) {
            float min_dist = 1e30f;
            uint8_t best_cluster = 0;

            for (int k = 0; k < MAX_PATTERNS; k++) {
                float dist = 0.0f;
                for (int d = 0; d < 48; d++) {
                    float diff = sh_data[i].c[d] - centroids[k][d];
                    dist += diff * diff;
                }
                if (dist < min_dist) {
                    min_dist = dist;
                    best_cluster = k;
                }
            }
            labels[i] = best_cluster;
            counts[best_cluster]++;
            for (int d = 0; d < 48; d++) {
                new_centroids[best_cluster * 48 + d] += sh_data[i].c[d];
            }
        }

        for (int k = 0; k < MAX_PATTERNS; k++) {
            if (counts[k] > 0) {
                for (int d = 0; d < 48; d++) {
                    centroids[k][d] = new_centroids[k * 48 + d] / counts[k];
                }
            }
        }
    }
}

uint32_t parse_ply_header_from_stdin(void) {
    char buf[256];
    int bp = 0;
    char c;
    uint32_t count = 0;

    while (read(0, &c, 1) == 1) {
        if (c == '\n' || bp >= 255) {
            buf[bp] = '\0';
            if (str_start(buf, "element vertex ")) {
                char *p = buf + 15;
                while (*p >= '0' && *p <= '9') {
                    count = count * 10 + (*p - '0');
                    p++;
                }
            }
            if (str_start(buf, "end_header")) break;
            bp = 0;
        } else {
            buf[bp++] = c;
        }
    }
    return count;
}

int main(void) {
    uint32_t n_splats = parse_ply_header_from_stdin();
    if (n_splats == 0 || n_splats > MAX_SPLATS) return 1;

    float dummy_normal[3];
    for (uint32_t i = 0; i < n_splats; i++) {
        read(0, &xyz_scale[i*6], 3 * sizeof(float));
        read(0, dummy_normal, 3 * sizeof(float));
        read(0, &sh_data[i].c[0], 3 * sizeof(float));
        read(0, &sh_data[i].c[3], 45 * sizeof(float));
        read(0, &opacity_rot[i*5], 1 * sizeof(float));
        read(0, &xyz_scale[i*6 + 3], 3 * sizeof(float));
        read(0, &opacity_rot[i*5 + 1], 4 * sizeof(float));
    }

    run_kmeans(n_splats);

    uint32_t num_patterns = MAX_PATTERNS;
    write(1, &num_patterns, sizeof(uint32_t));
    write(1, centroids, sizeof(float) * 48 * MAX_PATTERNS);
    write(1, &n_splats, sizeof(uint32_t));

    int out_p = 0;
    for (uint32_t i = 0; i < n_splats; i++) {
        out_buf[out_p].x = xyz_scale[i*6 + 0];
        out_buf[out_p].y = xyz_scale[i*6 + 1];
        out_buf[out_p].z = xyz_scale[i*6 + 2];
        out_buf[out_p].sx = xyz_scale[i*6 + 3];
        out_buf[out_p].sy = xyz_scale[i*6 + 4];
        out_buf[out_p].sz = xyz_scale[i*6 + 5];

        float op = opacity_rot[i*5 + 0];
        out_buf[out_p].opacity = (unsigned char)((op < 0.f ? 0.f : (op > 1.f ? 1.f : op)) * 255.0f);
        out_buf[out_p].q0 = (unsigned char)((opacity_rot[i*5 + 1] + 1.0f) * 127.5f);
        out_buf[out_p].q1 = (unsigned char)((opacity_rot[i*5 + 2] + 1.0f) * 127.5f);
        out_buf[out_p].q2 = (unsigned char)((opacity_rot[i*5 + 3] + 1.0f) * 127.5f);
        out_buf[out_p].q3 = (unsigned char)((opacity_rot[i*5 + 4] + 1.0f) * 127.5f);
        out_buf[out_p].sh_idx = labels[i];

        out_p++;
        if (out_p >= OUT_BUF_SPLAT_COUNT) {
            write(1, out_buf, OUT_BUF_SPLAT_COUNT * sizeof(FSplat));
            out_p = 0;
        }
    }
    if (out_p > 0) write(1, out_buf, out_p * sizeof(FSplat));

    return 0;
}