/*
 reinhard@finalmedia.de
 Fr 10. Jul 17:07:34 CEST 2026

 3DGS (Gaussian Splatting)
 Multithreaded SH-Codebuch-Lookup Szenen-Streaming

 Liest fortlaufend konkatenierte .fsplat-Dateien über stdin.
 Liest interaktive Kamerapositionen zeilenweise über Deskriptor 4.
 Gibt fdraw-Pixeldaten auf stdout aus.

 ### Compilieren mittels:
 musl-gcc -Q -mcmodel=large -march=native -O3 -static-pie -fPIE -fstack-protector-strong -D_FORTIFY_SOURCE=2 -Wl,-z,relro -Wl,-z,now fsplat.c -o fsplat -lm -pthread

 ### Kamerapositionen
 Position im Raum, Schwenk, Neigung und Zoom werden als
 befehl zeilenweise übergeben, z.B. womit auch animationen möglich sind.

 "P 4.0 0.0 .20 0 0 100"

  Die angegebenen float Werte (x y und h)
  Und int werte für yaw pitch und zoom

    x   y   h   yaw pitch zoom

 ### Offset Marker
 Mit dem Befehl O kann für alle ab dann auf stdin folgenden Szenendaten
 ein globaler Offset definiert werden, der die Szene verschiebt.

 "O 1.2 3.3 0.3"

 Aus dies wird bei Bedarf in den Descriptor 4 eingestreut.

 ### Nutzungsbeispiel

 echo P 4.0 0.0 .05 0 0 100 > camlines
 echo P 4.0 0.0 .10 0 0 100 >> camlines
 echo P 4.0 0.0 .15 0 0 100 >> camlines

 export SCREEN_WIDTH=120 SCREEN_HEIGHT=80
 zcat szene.fsplat.gz | fsplat 4<camlines | fdrawterm

*/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <unistd.h>
#include <stdint.h>
#include <pthread.h>

#define MAX_SCENES 16
#define MAX_SPLATS_PER_SCENE 5500000
#define MAX_RES_P_IDX (1920 * 1080)
#define MAX_PATTERNS 256
#define COEFFS_COUNT 48

#define PI 3.1415926535f
#define L_SZ 1024
float lut[L_SZ];

typedef struct { float r, g, b; } RGB;

RGB img[MAX_RES_P_IDX];
float rem_alpha[MAX_RES_P_IDX];

void init_lut(void) {
    int i;
    for (i = 0; i < L_SZ; i++)
        lut[i] = expf(-((float)i / L_SZ) * 4.0f);
}

inline float f_exp(float v) {
    int idx;
    if (v < 0.0f || isnan(v) || isinf(v)) return (v < 0.0f) ? 1.0f : 0.0f;
    idx = (int)((v * 0.25f) * L_SZ);
    if (idx < 0 || idx >= L_SZ) return 0.0f;
    return lut[idx];
}

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

typedef struct { float x, y, z, c11, c12, c22; unsigned char r, g, b, a; } TSplat;

typedef struct { float coeffs[COEFFS_COUNT]; } SHPattern;

typedef struct {
    float x, y, z;
    float v0, v1, v2, v3, v4, v5;
    unsigned char opacity;
    unsigned char sh_idx;
} SCached;

float global_offset_x = 0.0f;
float global_offset_y = 0.0f;
float global_offset_z = 0.0f;

SCached s_data[MAX_SCENES][MAX_SPLATS_PER_SCENE];
TSplat v_data[MAX_SCENES][MAX_SPLATS_PER_SCENE];
TSplat tmp_v_data[MAX_SCENES][MAX_SPLATS_PER_SCENE];
SHPattern codebooks[MAX_SCENES][MAX_PATTERNS];
RGB frame_luts[MAX_SCENES][MAX_PATTERNS];

typedef struct {
    size_t scene_idx;
    size_t n_s;
    size_t v_cnt;
    uint32_t num_patterns;
    float px, py, ph, view_dx, view_dy, view_dz, fov, hw, hh, cy, sy, cp, sp;
} SceneContext;

int w = 800, h = 600;

SceneContext scenes[MAX_SCENES];
size_t total_scenes_loaded = 0;

#define O_SZ (4 * 1024 * 1024)
char ob[O_SZ];
size_t op = 0;

static inline void out(char c) {
    ob[op++] = c;
    if (op >= O_SZ) { write(1, ob, op); op = 0; }
}

static inline void outi(int n) {
    char b[12]; int i = 0;
    if (n == 0) { out('0'); return; }
    if (n < 0) { out('-'); n = -n; }
    while (n > 0) { b[i++] = (n % 10) + '0'; n /= 10; }
    while (i > 0) out(b[--i]);
}

static inline void err_out(const char *s_ptr, size_t len) {
    write(2, s_ptr, len);
}

static inline void err_outi(int n) {
    char b[12]; int i = 0;
    if (n == 0) { err_out("0", 1); return; }
    if (n < 0) { err_out("-", 1); n = -n; }
    while (n > 0) { b[i++] = (n % 10) + '0'; n /= 10; }
    while (i > 0) write(2, &b[--i], 1);
}

#define SH_C0 0.28209479177387814f
#define SH_C1 0.4886025119029199f
#define SH_C2_0 1.0925484305920792f
#define SH_C2_1 -1.0925484305920792f
#define SH_C2_2 0.31539156525252005f
#define SH_C2_3 -1.0925484305920792f
#define SH_C2_4 0.5462742152960396f
#define SH_C3_0 -0.5900435899266435f
#define SH_C3_1 2.890611442640554f
#define SH_C3_2 -0.4570457994644658f
#define SH_C3_3 0.3731763325901154f
#define SH_C3_4 -0.4570457994644658f
#define SH_C3_5 1.445305721320277f
#define SH_C3_6 -0.5900435899266435f

static inline void compute_sh_element_local(size_t s_idx, int lut_idx, float dx, float dy, float dz) {
    float *c = codebooks[s_idx][lut_idx].coeffs;

    float x = dx, y = dy, z = dz;
    float xx = x*x, yy = y*y, zz = z*z;
    float xy = x*y, yz = y*z, xz = x*z;

    // Ordnung 0 & 1
    float r = SH_C0 * *(c+0) - SH_C1 * y * *(c+1) + SH_C1 * z * *(c+2) - SH_C1 * x * *(c+3);
    float g = SH_C0 * *(c+16) - SH_C1 * y * *(c+17) + SH_C1 * z * *(c+18) - SH_C1 * x * *(c+19);
    float b = SH_C0 * *(c+32) - SH_C1 * y * *(c+33) + SH_C1 * z * *(c+34) - SH_C1 * x * *(c+35);

    // Ordnung 2
    r += SH_C2_0 * xy * *(c+4) + SH_C2_1 * yz * *(c+5) + SH_C2_2 * (2.0f*zz - xx - yy) * *(c+6) + SH_C2_3 * xz * *(c+7) + SH_C2_4 * (xx - yy) * *(c+8);
    g += SH_C2_0 * xy * *(c+20) + SH_C2_1 * yz * *(c+21) + SH_C2_2 * (2.0f*zz - xx - yy) * *(c+22) + SH_C2_3 * xz * *(c+23) + SH_C2_4 * (xx - yy) * *(c+24);
    b += SH_C2_0 * xy * *(c+36) + SH_C2_1 * yz * *(c+37) + SH_C2_2 * (2.0f*zz - xx - yy) * *(c+38) + SH_C2_3 * xz * *(c+39) + SH_C2_4 * (xx - yy) * *(c+40);

    // Ordnung 3
    r += SH_C3_0 * y * (3.0f*xx - yy) * *(c+9) + SH_C3_1 * xy * z * *(c+10) + SH_C3_2 * y * (4.0f*zz - xx - yy) * *(c+11) + SH_C3_3 * z * (2.0f*zz - 3.0f*xx - 3.0f*yy) * *(c+12) + SH_C3_4 * x * (4.0f*zz - xx - yy) * *(c+13) + SH_C3_5 * z * (xx - yy) * *(c+14) + SH_C3_6 * x * (xx - 3.0f*yy) * *(c+15);
    g += SH_C3_0 * y * (3.0f*xx - yy) * *(c+25) + SH_C3_1 * xy * z * *(c+26) + SH_C3_2 * y * (4.0f*zz - xx - yy) * *(c+27) + SH_C3_3 * z * (2.0f*zz - 3.0f*xx - 3.0f*yy) * *(c+28) + SH_C3_4 * x * (4.0f*zz - xx - yy) * *(c+29) + SH_C3_5 * z * (xx - yy) * *(c+30) + SH_C3_6 * x * (xx - 3.0f*yy) * *(c+31);
    b += SH_C3_0 * y * (3.0f*xx - yy) * *(c+41) + SH_C3_1 * xy * z * *(c+42) + SH_C3_2 * y * (4.0f*zz - xx - yy) * *(c+43) + SH_C3_3 * z * (2.0f*zz - 3.0f*xx - 3.0f*yy) * *(c+44) + SH_C3_4 * x * (4.0f*zz - xx - yy) * *(c+45) + SH_C3_5 * z * (xx - yy) * *(c+46) + SH_C3_6 * x * (xx - 3.0f*yy) * *(c+47);

    r = r + 0.5f; g = g + 0.5f; b = b + 0.5f;
    frame_luts[s_idx][lut_idx].r = (unsigned char)(fmaxf(0.0f, fminf(255.0f, r * 255.0f)));
    frame_luts[s_idx][lut_idx].g = (unsigned char)(fmaxf(0.0f, fminf(255.0f, g * 255.0f)));
    frame_luts[s_idx][lut_idx].b = (unsigned char)(fmaxf(0.0f, fminf(255.0f, b * 255.0f)));
}

static inline void radix_sort_splats_front_to_back(TSplat *in, TSplat *out_tmp, size_t count) {
    size_t i, shift, buckets[256];
    TSplat *src = in, *dst = out_tmp;
    for (shift = 0; shift < 32; shift += 8) {
        for (i = 0; i < 256; i++) buckets[i] = 0;
        for (i = 0; i < count; i++) {
            unsigned int bits = *(unsigned int *)&src[i].z;
            unsigned char byte = (bits >> shift) & 0xFF;
            buckets[byte]++;
        }
        size_t current_offset = 0;
        for (i = 0; i < 256; i++) {
            size_t b_size = buckets[i]; buckets[i] = current_offset; current_offset += b_size;
        }
        for (i = 0; i < count; i++) {
            unsigned int bits = *(unsigned int *)&src[i].z;
            unsigned char byte = (bits >> shift) & 0xFF;
            size_t dst_idx = buckets[byte]++; dst[dst_idx] = src[i];
        }
        TSplat *tmp = src; src = dst; dst = tmp;
    }
    if (src != in) { for (i = 0; i < count; i++) in[i] = src[i]; }
}


void* scene_worker(void *arg) {
    SceneContext *ctx = (SceneContext*)arg;
    size_t s_idx = ctx->scene_idx;
    ctx->v_cnt = 0;

    for (uint32_t i = 0; i < ctx->num_patterns; i++) {
        compute_sh_element_local(s_idx, i, ctx->view_dx, ctx->view_dy, ctx->view_dz);
    }

    for (size_t p = 0; p < ctx->n_s; p++) {
        SCached cur = s_data[s_idx][p];
        float tx = cur.x - ctx->px - global_offset_x;
        float ty = cur.y - ctx->py - global_offset_y;
        float tz = cur.z - ctx->ph - global_offset_z;

        float rx = tx * ctx->cy - tz * ctx->sy, rz = tx * ctx->sy + tz * ctx->cy;
        float fz = ty * ctx->sp + rz * ctx->cp, fx = rx, fy = ty * ctx->cp - rz * ctx->sp;

        if (fz <= 0.15f) continue;
        float inv_z = 1.0f / fz, p_x = fx * ctx->fov * inv_z + ctx->hw, p_y = fy * ctx->fov * inv_z + ctx->hh;
        if (p_x < -64.f || p_x > w + 64.f || p_y < -64.f || p_y > h + 64.f) continue;

        float J00 = ctx->fov * inv_z, J02 = -(fx * ctx->fov) * (inv_z * inv_z), J12 = -(fy * ctx->fov) * (inv_z * inv_z);

        float c11 = J00 * J00 * cur.v0 + 2.f * J00 * J02 * cur.v2 + J02 * J02 * cur.v5 + 0.3f;
        float c12 = J00 * J00 * cur.v1 + J00 * J12 * cur.v2 + J02 * J00 * cur.v2 + J02 * J12 * cur.v5;
        float c22 = J00 * J00 * cur.v3 + 2.f * J00 * J12 * cur.v4 + J12 * J12 * cur.v5 + 0.3f;

        if (c11 > 1e18f || c22 > 1e18f || c11 < 0.f || c22 < 0.f) continue;
        float det = c11 * c22 - c12 * c12; if (det <= 0.0001f || isnan(det) || isinf(det)) continue;
        float mid = 0.5f * (c11 + c22); float term = mid * mid - det; if (term < 0.f) continue;
        float rad = ceilf(3.0f * sqrtf(term)); if (rad > 64.f) rad = 64.f;

        if (p_x + rad < 0.f || p_x - rad >= w || p_y + rad < 0.f || p_y - rad >= h || ctx->v_cnt >= MAX_SPLATS_PER_SCENE) continue;

        RGB color = frame_luts[s_idx][cur.sh_idx];

        v_data[s_idx][ctx->v_cnt++] = (TSplat){p_x, p_y, fz, c11, c12, c22, color.r, color.g, color.b, cur.opacity};
    }

    radix_sort_splats_front_to_back(v_data[s_idx], tmp_v_data[s_idx], ctx->v_cnt);
    return NULL;
}



int main(void) {
    init_lut();

    if (getenv("SCREEN_WIDTH"))  w = atoi(getenv("SCREEN_WIDTH"));
    if (getenv("SCREEN_HEIGHT")) h = atoi(getenv("SCREEN_HEIGHT"));

    err_out("SCREEN: ", 8); err_outi(w); err_out("x", 1); err_outi(h); err_out("\n", 1);

    FILE *f4 = fdopen(4, "r"); if (!f4) return 1;
    setvbuf(f4, NULL, _IONBF, 0);

    size_t n_pix = (size_t)w * h, p;
    uint32_t current_file_patterns = 0;
    uint32_t current_file_splats = 0;

    err_out("Lese Szenen-Streams von stdin...\n", 31);

    while (fread(&current_file_patterns, sizeof(uint32_t), 1, stdin) == 1) {
        if (total_scenes_loaded >= MAX_SCENES) {
            err_out("WARNUNG: MAX_SCENES erreicht! Kappe Stream.\n", 44);
            break;
        }

        SceneContext *ctx = &scenes[total_scenes_loaded];
        ctx->scene_idx = total_scenes_loaded;
        ctx->num_patterns = (current_file_patterns > MAX_PATTERNS) ? MAX_PATTERNS : current_file_patterns;

        if (ctx->num_patterns > 0) {
            if (fread(codebooks[total_scenes_loaded], sizeof(SHPattern), ctx->num_patterns, stdin) != ctx->num_patterns) {
                err_out("FEHLER: Codebuch unvollstaendig.\n", 32); return 1;
            }
        }

        if (fread(&current_file_splats, sizeof(uint32_t), 1, stdin) != 1) {
            err_out("FEHLER: Splat-Anzahl fehlt.\n", 29); return 1;
        }

        FSplat temp_in;
        ctx->n_s = 0;
        for (uint32_t i = 0; i < current_file_splats; i++) {
            if (fread(&temp_in, sizeof(FSplat), 1, stdin) != 1) break;

            if (ctx->n_s >= MAX_SPLATS_PER_SCENE) {
                if (ctx->n_s == MAX_SPLATS_PER_SCENE) {
                    err_out("WARNUNG: Teilszene gekappt (MAX_SPLATS_PER_SCENE).\n", 50);
                }
                ctx->n_s++;
                continue;
            }

            float q0 = (temp_in.q0 - 128.f) / 128.f, q1 = (temp_in.q1 - 128.f) / 128.f;
            float q2 = (temp_in.q2 - 128.f) / 128.f, q3 = (temp_in.q3 - 128.f) / 128.f;
            float ql = sqrtf(q0*q0 + q1*q1 + q2*q2 + q3*q3);
            if (ql > 0.f) { q0 /= ql; q1 /= ql; q2 /= ql; q3 /= ql; }

            float r00 = 1.f - 2.f*(q2*q2 + q3*q3), r01 = 2.f*(q1*q2 - q0*q3), r02 = 2.f*(q1*q3 + q0*q2);
            float r10 = 2.f*(q1*q2 + q0*q3), r11 = 1.f - 2.f*(q1*q1 + q3*q3), r12 = 2.f*(q2*q3 - q0*q1);
            float r20 = 2.f*(q1*q3 - q0*q2), r21 = 2.f*(q2*q3 + q0*q1), r22 = 1.f - 2.f*(q1*q1 + q2*q2);
            float sx2 = temp_in.sx * temp_in.sx, sy2 = temp_in.sy * temp_in.sy, sz2 = temp_in.sz * temp_in.sz;

            s_data[total_scenes_loaded][ctx->n_s].x = temp_in.x; 
            s_data[total_scenes_loaded][ctx->n_s].y = temp_in.y; 
            s_data[total_scenes_loaded][ctx->n_s].z = temp_in.z;
            s_data[total_scenes_loaded][ctx->n_s].v0 = r00*r00*sx2 + r01*r01*sy2 + r02*r02*sz2;
            s_data[total_scenes_loaded][ctx->n_s].v1 = r00*r10*sx2 + r01*r11*sy2 + r02*r12*sz2;
            s_data[total_scenes_loaded][ctx->n_s].v2 = r00*r20*sx2 + r01*r21*sy2 + r02*r22*sz2;
            s_data[total_scenes_loaded][ctx->n_s].v3 = r10*r10*sx2 + r11*r11*sy2 + r12*r12*sz2;
            s_data[total_scenes_loaded][ctx->n_s].v4 = r10*r20*sx2 + r11*r11*sy2 + r12*r22*sz2;
            s_data[total_scenes_loaded][ctx->n_s].v5 = r20*r20*sx2 + r21*r21*sy2 + r22*r22*sz2;

            s_data[total_scenes_loaded][ctx->n_s].opacity = temp_in.opacity;
            s_data[total_scenes_loaded][ctx->n_s].sh_idx = temp_in.sh_idx;
            ctx->n_s++;
        }

        if (ctx->n_s > MAX_SPLATS_PER_SCENE) { ctx->n_s = MAX_SPLATS_PER_SCENE; }
        total_scenes_loaded++;
    }

    err_out("Streaming beendet. Szenen aktiv: ", 32);
    err_outi((int)total_scenes_loaded); err_out("\n", 1);

    char line_buffer[256];
    float px, py, ph, hw = w * 0.5f, hh = h * 0.5f;
    int yaw, pitch, zoom;
    pthread_t threads[MAX_SCENES];

    while (fgets(line_buffer, sizeof(line_buffer), f4)) {
        if (line_buffer[0] == 'O') {
            char *ptr = line_buffer + 1;
            global_offset_x = strtof(ptr, &ptr);
            global_offset_y = strtof(ptr, &ptr);
            global_offset_z = strtof(ptr, &ptr);
            continue;
        }

        if (line_buffer[0] != 'P') continue;
        char *ptr = line_buffer + 1;
        px = strtof(ptr, &ptr); py = strtof(ptr, &ptr); ph = strtof(ptr, &ptr);
        yaw = (int)strtol(ptr, &ptr, 10); pitch = (int)strtol(ptr, &ptr, 10); zoom = (int)strtol(ptr, &ptr, 10);

        if (zoom <= 0) zoom = 1;
        if (n_pix > MAX_RES_P_IDX) n_pix = MAX_RES_P_IDX;

        for (p = 0; p < n_pix; p++) {
            img[p].r = 0.f;
            img[p].g = 0.f;
            img[p].b = 0.f;
            rem_alpha[p] = 1.0f;
        }

        float cy = cosf(yaw * PI / 180.f), sy = sinf(yaw * PI / 180.f);
        float cp = cosf(pitch * PI / 180.f), sp = sinf(pitch * PI / 180.f);
        float fov = (w * 0.75f) * ((float)zoom / 100.f);

        for (size_t i = 0; i < total_scenes_loaded; i++) {
            scenes[i].px = px; scenes[i].py = py; scenes[i].ph = ph;
            scenes[i].view_dx = sy * cp; scenes[i].view_dy = sp; scenes[i].view_dz = cy * cp;
            scenes[i].fov = fov; scenes[i].hw = hw; scenes[i].hh = hh;
            scenes[i].cy = cy; scenes[i].sy = sy; scenes[i].cp = cp; scenes[i].sp = sp;
            pthread_create(&threads[i], NULL, scene_worker, &scenes[i]);
        }

        for (size_t i = 0; i < total_scenes_loaded; i++) {
            pthread_join(threads[i], NULL);
        }

        for (size_t i = 0; i < total_scenes_loaded; i++) {
            SceneContext *ctx = &scenes[i];
            size_t s_idx = ctx->scene_idx;
            for (p = 0; p < ctx->v_cnt; p++) {
                TSplat cs = v_data[s_idx][p];
                float det = cs.c11 * cs.c22 - cs.c12 * cs.c12;
                if (det <= 0.0001f) continue;
                float inv_det = 1.0f / det, ic11 = 0.5f * cs.c22 * inv_det, ic12 = 0.5f * -cs.c12 * inv_det, ic22 = 0.5f * cs.c11 * inv_det;
                float mid = 0.5f * (cs.c11 + cs.c22), term = mid * mid - det;
                int radius = (term < 0.f) ? 1 : (int)ceilf(3.f * sqrtf(term)); if (radius > 64) radius = 64;
                int px_i = (int)cs.x, py_i = (int)cs.y;
                int xs = fmaxf(0, px_i - radius), xe = fminf(w - 1, px_i + radius);
                int ys = fmaxf(0, py_i - radius), ye = fminf(h - 1, py_i + radius);
                if (xs > xe || ys > ye) continue;
                float ab = cs.a / 255.f, vr = cs.r, vg = cs.g, vb = cs.b, vs_x = cs.x, vs_y = cs.y;

                for (int y_p = ys; y_p <= ye; y_p++) {
                    float dy = (float)y_p - vs_y, dy_part = dy * dy * ic22, cross_part = 2.0f * dy * ic12;
                    size_t offset = (size_t)y_p * w;
                    for (int x_p = xs; x_p <= xe; x_p++) {
                        size_t idx = offset + x_p;
                        if (rem_alpha[idx] < 0.01f) continue;
                        float dx = (float)x_p - vs_x, exp_val = dx * dx * ic11 + dx * cross_part + dy_part;
                        float weight = f_exp(exp_val);
                        if (weight < 0.05f) continue;
                        float alpha = ab * weight;
                        float blended_alpha = alpha * rem_alpha[idx];
                        img[idx].r += blended_alpha * vr;
                        img[idx].g += blended_alpha * vg;
                        img[idx].b += blended_alpha * vb;
                        rem_alpha[idx] *= (1.0f - alpha);
                    }
                }
            }
        }

        size_t p_idx = 0;
        for (int y_p = 0; y_p < h; y_p++) {
            for (int x_p = 0; x_p < w; x_p++) {
                out('p'); out(' '); outi(x_p); out(' '); outi(y_p); out(' ');
                outi((int)(img[p_idx].r + 0.5f)); out(' ');
                outi((int)(img[p_idx].g + 0.5f)); out(' ');
                outi((int)(img[p_idx].b + 0.5f)); out('\n');
                p_idx++;
            }
        }
        if (op > 0) { write(1, ob, op); op = 0; }
        write(1, "s 33333\n", 8);
        err_out(".\n", 2);
    }

    fclose(f4);
    return 0;
}