/*
    James William Fletcher (github.com/mrbid)
        April 2024
*/

#include <math.h>
#include <stdio.h>
#include <string.h>
#include <time.h>
#include <dirent.h>

#include <sys/file.h>
#include <unistd.h>

#include "rply.h"
#include "../icosphere.h"

#define uint unsigned int
#define INPUT_SIZE 2
#define OUTPUT_SIZE 6

#pragma GCC diagnostic ignored "-Wunused-result"

//*************************************
// utility functions
//*************************************
void timestamp(char* ts){const time_t tt = time(0); strftime(ts, 16, "%H:%M:%S", localtime(&tt));}
float urandf32()
{
    static const float RECIP_FLOAT_UINT32_MAX = 1.f/(float)((unsigned int)-1);
    int f = open("/dev/urandom", O_RDONLY | O_CLOEXEC);
    unsigned int s = 0;
    read(f, &s, sizeof(unsigned int));
    close(f);
    return ((float)s) * RECIP_FLOAT_UINT32_MAX;
}

//*************************************
// load models from file to gpu memory
//*************************************
// permenant 8 MB staging buffer
#define MAX_SIZE 2097152
float vbuff[MAX_SIZE];
uint vbl = 0; // vertex buffer length
uint nverts = 0; // number of vertices
static int vertex_cb(p_ply_argument argument)
{
    if(vbl > MAX_SIZE-1)
    {
        puts("terminated: model too big for internal staging vertex buffer");
        return 0;
    }
    static uint vc = 0;
    long eol;
    ply_get_argument_user_data(argument, NULL, &eol);
    vbuff[vbl] = ply_get_argument_value(argument);
    if(vc > 5){vbuff[vbl] *= 0.003921569f;}
    vbl++;
    vc++;
    if(eol){vc = 0;}
    return 1;
}
void processModels(const char* dir_path)
{
    const size_t INPUT_SIZE_BYTES = INPUT_SIZE*4;
    const size_t OUTPUT_SIZE_BYTES = OUTPUT_SIZE*4;
    float input[INPUT_SIZE];
    float output[OUTPUT_SIZE];
    uint total_processed = 0;
    DIR* d = opendir(dir_path);
    if(d != NULL)
    {
        struct dirent *dir;
        while((dir = readdir(d)) != NULL)
        {
            if(dir->d_name[0] != '.')
            {
                // reset buffers
                vbl = 0;
                const float model_seed = urandf32();

                // new file load path
                char fp[384];
                sprintf(fp, "%s/%s", dir_path, dir->d_name);
                //puts(fp);

                // open file
                p_ply ply = ply_open(fp, NULL, 0, NULL);
                if(!ply){continue;} // rply spits an error to console anyway if the open fails
                if(!ply_read_header(ply))
                {
                    ply_close(ply);
                    continue;
                }

                // read file setup
                nverts = ply_set_read_cb(ply, "vertex", "x", vertex_cb, NULL, 0);
                ply_set_read_cb(ply, "vertex", "y", vertex_cb, NULL, 0);
                ply_set_read_cb(ply, "vertex", "z", vertex_cb, NULL, 0);
                ply_set_read_cb(ply, "vertex", "nx", vertex_cb, NULL, 0);
                ply_set_read_cb(ply, "vertex", "ny", vertex_cb, NULL, 0);
                ply_set_read_cb(ply, "vertex", "nz", vertex_cb, NULL, 0);
                ply_set_read_cb(ply, "vertex", "red", vertex_cb, NULL, 0);
                ply_set_read_cb(ply, "vertex", "green", vertex_cb, NULL, 0);
                ply_set_read_cb(ply, "vertex", "blue", vertex_cb, NULL, 1);

                // read file
                if(!ply_read(ply))
                {
                    ply_close(ply);
                    continue;
                }

                // close file
                ply_close(ply);

                // process vertices (36 bytes / 9 floats per vertex)
                for(uint i = 0; i < nverts; i++)
                {
                    // printf("x: %f\n",  vbuff[9*i]);
                    // printf("y: %f\n",  vbuff[(9*i)+1]);
                    // printf("z: %f\n",  vbuff[(9*i)+2]);
                    // printf("nx: %f\n", vbuff[(9*i)+3]);
                    // printf("ny: %f\n", vbuff[(9*i)+4]);
                    // printf("nz: %f\n", vbuff[(9*i)+5]);
                    // printf("r: %f\n",  vbuff[(9*i)+6]);
                    // printf("g: %f\n",  vbuff[(9*i)+7]);
                    // printf("b: %f\n",  vbuff[(9*i)+8]);
                    // printf("---\n");
                    memcpy(&output[0], &vbuff[9*i], 12);
                    memcpy(&output[3], &vbuff[(9*i)+6], 12);
                    // printf("x: %f\n", ((float*)output)[0]);
                    // printf("y: %f\n", ((float*)output)[1]);
                    // printf("z: %f\n", ((float*)output)[2]);
                    // printf("r: %f\n", ((float*)output)[3]);
                    // printf("g: %f\n", ((float*)output)[4]);
                    // printf("b: %f\n", ((float*)output)[5]);
                    // printf("###\n");

                    // get vertex pos
                    const float vx = vbuff[9*i];
                    const float vy = vbuff[(9*i)+1];
                    const float vz = vbuff[(9*i)+2];
                    //printf("v: %f %f %f\n", vx, vy, vz);

                    // get vertex normal
                    const float vnx = vbuff[9*i+3];
                    const float vny = vbuff[(9*i)+4];
                    const float vnz = vbuff[(9*i)+5];
                    //printf("vn: %f %f %f\n", vnx, vny, vnz);

                    // get remainder distance of vertex from 0,0,0
                    const float pd = 1.f-sqrtf(vx*vx + vy*vy + vz*vz);
                    //printf("pd: %f\n", pd);

                    // scale the vertex normal by remainder
                    const float pdx = vnx*pd;
                    const float pdy = vny*pd;
                    const float pdz = vnz*pd;
                    //printf("vn-s: %f %f %f\n", pdx, pdy, pdz);

                    // project vertex positon to unit sphere
                    const float upx = vx+pdx;
                    const float upy = vy+pdy;
                    const float upz = vz+pdz;
                    //printf("v-p: %f %f %f\n", upx, upy, upz);

                    // find closest vertex on the ico sphere
                    int closest = -1;
                    float lcd = 9999.f;
                    for(uint i = 0; i < icosphere_size; i+=3)
                    {
                        const float xm = icosphere[i]   - upx;
                        const float ym = icosphere[i+1] - upy;
                        const float zm = icosphere[i+2] - upz;
                        const float d = sqrtf(xm*xm + ym*ym + zm*zm);
                        if(d < lcd){lcd = d; closest=i;}
                    }
                    // printf("a: %f %f %f\n", icosphere[closest], icosphere[closest+1], icosphere[closest+2]);
                    // printf("b: %f %f %f\n", upx, upy, upz);
                    // printf("lcd: %f\n", lcd);
                    if(closest == -1)
                    {
                        puts("found no closest unit sphere point on icosphere");
                        continue;
                    }

                    input[0] = model_seed; // random
                    input[1] = ((float)closest) / ((float)icosphere_size); // icosphere index
                    //if(input[1] > 1.f){printf("ind: %f\n", input[1]);}
                    // input[2] = -vnx; // inverted vertex normal
                    // input[3] = -vny;
                    // input[4] = -vnz;

                    // write processed vertex
                    FILE* fp = fopen("train_x.dat", "ab");
                    if(fp != NULL)
                    {
                        if(fwrite(&input, 1, INPUT_SIZE_BYTES, fp) != INPUT_SIZE_BYTES)
                            puts("train_x.dat: write error.");
                        fclose(fp);
                    }
                    fp = fopen("train_y.dat", "ab");
                    if(fp != NULL)
                    {
                        if(fwrite(&output, 1, OUTPUT_SIZE_BYTES, fp) != OUTPUT_SIZE_BYTES)
                            puts("train_y.dat: write error.");
                        fclose(fp);
                    }
                }

                // stats
                char strts[16];
                timestamp(&strts[0]);
                printf("[%s] [%u] %s\n", strts, total_processed, fp);
                total_processed++;
            }
        }
    }
}

//*************************************
// process entry point
//*************************************
int main(int argc, char** argv)
{
    const time_t st = time(0);
    processModels("../ply");
    printf("Time Taken: %lu seconds\n", time(0)-st);
    return 0;
}
