#!icc -lm -pthread -std=c99

#define _GNU_SOURCE

#include <getopt.h>
#include <math.h>
#include <pthread.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#define DPRINT printf("ok at %d\n", __LINE__);
#define zeros(n, type) (type*)calloc((n), sizeof(type))

struct global_data {
    float x_span;
    float y_span;
    float z_span;
    int ngrid;
    int nbins;
    int nupc; // muons per cell
    int nthreads;
    struct thread_data* threads;
    float l_min;
    float l_max;
    float l_scale;
    float th_max;
    float th_scale;
};

struct thread_data {
    pthread_t thread_id;
    int thread_num;
    int i_x1;
    int* r_l_bins;
    int* r_th_bins;
    struct global_data* gd;
};

inline float random_one() {
    return (float)random() / RAND_MAX - (float)random() / RAND_MAX;
}

void distance_distrib(struct thread_data* td) {
    struct global_data* gd = td->gd;
    int i_x1 = td->i_x1;
    float x_span = gd->x_span;
    float y_span = gd->y_span;
    float z_span = gd->z_span;
    int ngrid = gd->ngrid;
    int nbins = gd->nbins;
    float l_min = gd->l_min;
    float l_scale = gd->l_scale;
    float th_scale = gd->th_scale;
    for (int i_y1 = 0; i_y1 < ngrid; i_y1++) {
        for (int i_x2 = 0; i_x2 < ngrid; i_x2++) {
            float dx = x_span * (i_x1 - i_x2) / ngrid;
            for (int i_y2 = 0; i_y2 < ngrid; i_y2++) {
                float dy = y_span * (i_y1 - i_y2) / ngrid;
                float l = sqrt(dx * dx + dy * dy + z_span * z_span);
                float nmu = gd->nupc * (z_span * z_span) / (l * l);
                for (int i = 0; i < nmu; i++) {
                    float dx_r = dx + random_one();
                    float dy_r = dy + random_one();
                    float l_r = sqrt(dx_r * dx_r + dy_r * dy_r + z_span * z_span);
                    float th_r = acos(z_span / l_r);
                    td->r_l_bins[(int)((l_r - l_min) * l_scale)] += 1;
                    td->r_th_bins[(int)(th_r * th_scale)] += 1;
                }
            }
        }
    }
}

void* thread_run(void* arg) {
    struct thread_data* td = (struct thread_data*)arg;
    struct global_data* gd = td->gd;
    for (td->i_x1 = td->thread_num; td->i_x1 < gd->ngrid; td->i_x1 += gd->nthreads) {
        distance_distrib(td);
    }
    return NULL;
}

struct option opts[] = {
    {"bins", 1, 0, 'n'},
    {"flux", 1, 0, 'f'},
    {"resolution", 1, 0, 'r'},
    {0, 0, 0, 0}
};

static struct global_data* glob_data = NULL;

void print_state(int signum) {
    int nthreads = glob_data->nthreads;
    for (int i = 0; i < nthreads; i++) {
        fprintf(stderr, "thread %d: i_x1 = %d\n", i, glob_data->threads[i].i_x1);
    }
}

int main(int argc, char** argv) {
    // Declare values
    int nbins = 10000;
    int ngrid = 100;
    int nupc = 4;
    int nthreads = 1;
    //float l_min = 0.0;
    //float l_max = 0.0;
    //float th_max = 0.0;
    // Parse options
    while (1) {
        int c = getopt_long(argc, argv, "f:j:n:r:", opts, NULL);
        if (c == -1) {
            break;
        }
        switch (c) {
            case 'f':
                nupc = atoi(optarg);
                break;
            case 'j':
                nthreads = atoi(optarg);
                break;
            case 'n':
                nbins = atoi(optarg);
                break;
            case 'r':
                ngrid = atoi(optarg);
                break;
            default:
                return 2;
        }
    }
    if (argc - optind < 3) {
        printf("Usage: %s [-f|--flux flux] [-n|--bins nbins] [-r|--resolution res] xdim ydim zdim\n", argv[0]);
        return 2;
    }
    float xdim = (float)strtof(argv[optind++], NULL);
    float ydim = (float)strtof(argv[optind++], NULL);
    float zdim = (float)strtof(argv[optind++], NULL);
    if (xdim == 0.0 || ydim == 0.0 || zdim == 0.0) {
        printf("Invalid number\nUsage: %s [-f|--flux flux] [-n|--bins nbins] [-r|--resolution res] xdim ydim zdim\n", argv[0]);
        return 2;
    }
    int* l_bins = zeros(nbins, int);
    if (!l_bins) {
        return 1;
    }
    int* th_bins = zeros(nbins, int);
    if (!th_bins) {
        return 1;
    }
    // set up signal
    struct sigaction sa;
    sa.sa_handler = &print_state;
    sigaction(SIGUSR1, &sa, NULL);
    // initialize some values
    struct global_data gd;
    glob_data = &gd;
    gd.x_span = xdim;
    gd.y_span = ydim;
    gd.z_span = zdim;
    gd.ngrid = ngrid;
    gd.nbins = nbins;
    gd.nupc = nupc;
    gd.nthreads = nthreads;
    gd.l_min = gd.z_span;
    gd.l_max = sqrt(gd.x_span * gd.x_span + gd.y_span * gd.y_span + gd.z_span * gd.z_span);
    gd.th_max = acos(gd.z_span / gd.l_max);
    gd.l_scale = gd.nbins / (gd.l_max - gd.l_min);
    gd.th_scale = gd.nbins / gd.th_max;
    // create threads
    struct thread_data* threads = calloc(nthreads, sizeof(struct thread_data));
    if (!threads) {
        exit(3); // error
    }
    gd.threads = threads;
    for (int t = 0; t < nthreads; t++) {
        threads[t].i_x1 = 0;
        threads[t].thread_num = t;
        threads[t].gd = &gd;
        threads[t].r_l_bins = zeros(nbins, int);
        if (!threads[t].r_l_bins) {
            exit(3);
        }
        threads[t].r_th_bins = zeros(nbins, int);
        if (!threads[t].r_th_bins) {
            exit(3);
        }
        int s = pthread_create(&threads[t].thread_id, NULL, &thread_run, &threads[t]);
        if (s != 0) {
            exit(3);
        }
    }
    for (int t = 0; t < nthreads; t++) {
        // wait for termination
        int s = pthread_join(threads[t].thread_id, NULL);
        if (s != 0) {
            exit(3);
        }
        // and merge the results
        for (int i = 0; i < nbins; i++) {
            l_bins[i] += threads[t].r_l_bins[i];
            th_bins[i] += threads[t].r_th_bins[i];
        }
        free(threads[t].r_l_bins);
        free(threads[t].r_th_bins);
    }
    free(threads);

    printf("%5s\t%7s\t%10s\t%7s\t%10s\n", "bin", "l", "l_count", "theta", "th_count");
    for (int i = 0; i < nbins; i++) {
        printf("%5d\t%7.3f\t%10d\t%7.5f\t%10d\n", i, gd.l_min + (i * (float)(gd.l_max - gd.l_min)) / nbins, l_bins[i], i * (float)gd.th_max / nbins, th_bins[i]);
    }
    printf("%5d\t%7.3f\t%10d\t%7.5f\t%10d\n", nbins, gd.l_max, 0, gd.th_max, 0);

    free(l_bins);
    free(th_bins);
    return 0;
}