#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <sys/stat.h>
#include <errno.h>
#include <string.h>
#include <math.h>

#define MAX_PATH 256
#define DEFAULT_SIZE 1024

// Check if number is power of 2
int is_power_of_2(int n) {
    return n > 0 && (n & (n - 1)) == 0;
}

// Create directory if it doesn't exist
void ensure_directory(const char *path) {
    char tmp[MAX_PATH];
    char *p = NULL;
    size_t len;

    snprintf(tmp, sizeof(tmp), "%s", path);
    len = strlen(tmp);
    if (tmp[len - 1] == '/')
        tmp[len - 1] = 0;
    for (p = tmp + 1; *p; p++) {
        if (*p == '/') {
            *p = 0;
            mkdir(tmp, S_IRWXU);
            *p = '/';
        }
    }
    mkdir(tmp, S_IRWXU);
}

void print_usage(const char *program_name) {
    printf("Usage: %s [size]\n", program_name);
    printf("  size: Number of elements (must be a power of 2)\n");
    printf("  Default size: %d\n", DEFAULT_SIZE);
    printf("  Example sizes: 1024, 2048, 4096, 8192, 16384, etc.\n");
}

int main(int argc, char *argv[]) {
    int N = DEFAULT_SIZE;
    
    // Parse command line arguments
    if (argc > 1) {
        if (strcmp(argv[1], "-h") == 0 || strcmp(argv[1], "--help") == 0) {
            print_usage(argv[0]);
            return 0;
        }
        
        N = atoi(argv[1]);
        if (N <= 0) {
            fprintf(stderr, "Error: Size must be positive\n");
            print_usage(argv[0]);
            return 1;
        }
        if (!is_power_of_2(N)) {
            fprintf(stderr, "Error: Size must be a power of 2\n");
            print_usage(argv[0]);
            return 1;
        }
    }

    char dir_path[MAX_PATH];
    char file_path[MAX_PATH];

    // Create paths
    snprintf(dir_path, sizeof(dir_path), "dataset/%d", N);
    snprintf(file_path, sizeof(file_path), "%s/input.bin", dir_path);

    // Ensure directory exists
    ensure_directory(dir_path);

    FILE *file = fopen(file_path, "wb");
    if (!file) {
        perror("Cannot create file");
        return 1;
    }

    // Allocate memory
    float *data = (float *)malloc(N * sizeof(float));
    if (!data) {
        perror("Memory allocation failed");
        fclose(file);
        return 1;
    }

    // Initialize random seed
    srand(time(NULL));

    // Generate random floating point numbers
    for (int i = 0; i < N; i++) {
        data[i] = (float)rand() / RAND_MAX * 1000.0f;  // Random numbers between 0 and 1000
    }

    // Write to file
    size_t written = fwrite(data, sizeof(float), N, file);
    if (written != N) {
        perror("Failed to write all data");
        free(data);
        fclose(file);
        return 1;
    }

    // Cleanup
    free(data);
    fclose(file);

    // Print success message with size in human-readable format
    const char *size_suffix = "";
    float size_value = N;
    if (N >= 1048576) {
        size_suffix = "M";
        size_value = N / 1048576.0f;
    } else if (N >= 1024) {
        size_suffix = "K";
        size_value = N / 1024.0f;
    }

    printf("Successfully generated %.2f%s elements and saved to %s\n", 
           size_value, size_suffix, file_path);
    printf("Data size: %.2f MB\n", (N * sizeof(float)) / (1024.0f * 1024.0f));

    return 0;
}