/* Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
 *
 * This program is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option)
 * any later version.

 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.

 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include "path.h"
#include <stddef.h> /* for size_t */
#include "mt19937ar.h"
#include <math.h> /* for sin(), cos(), etc. */
#include "vector.h"
#include <errno.h> /* for errno */
#include <string.h> /* for strerror() */
#include "random.h"
#include <gsl/gsl_histogram.h>
#include "misc.h"

typedef struct sim {
    double *pos;
    double *dir;
    path *p;
    double *z1;
    double *z2;
} sim;

static double getKineticEnergy(double x, void *params)
{
    return 1.0;
}

sim *simulate_path(double *pos0, double *dir0, double range, double theta0, size_t N, size_t n)
{
    size_t i;
    double *s, *theta1, *theta2, *x, *y, *z, tmp[3], tmp2[3], k[3], normal[3];
    double dx, dy, dz, theta, phi;

    sim *simvalue = malloc(sizeof(sim));

    s = calloc(N,sizeof(double));
    theta1 = calloc(N,sizeof(double));
    theta2 = calloc(N,sizeof(double));
    x = calloc(N,sizeof(double));
    y = calloc(N,sizeof(double));
    z = calloc(N,sizeof(double));
    simvalue->pos = calloc(N*3,sizeof(double));
    simvalue->dir = calloc(N*3,sizeof(double));
    simvalue->z1 = calloc(n,sizeof(double));
    simvalue->z2 = calloc(n,sizeof(double));

    simvalue->pos[0] = pos0[0];
    simvalue->pos[1] = pos0[1];
    simvalue->pos[2] = pos0[2];

    for (i = 1; i < N; i++) {
        s[i] = range*i/(N-1);

        theta1[i] = theta1[i-1] + randn()*theta0*sqrt(s[i]-s[i-1]);
        theta2[i] = theta2[i-1] + randn()*theta0*sqrt(s[i]-s[i-1]);

        theta = sqrt(theta1[i]*theta1[i] + theta2[i]*theta2[i]);
        phi = atan2(theta2[i],theta1[i]);

        dx = (s[i]-s[i-1])*sin(theta)*cos(phi);
        dy = (s[i]-s[i-1])*sin(theta)*sin(phi);
        dz = (s[i]-s[i-1])*cos(theta);

        x[i] = x[i-1] + dx;
        y[i] = y[i-1] + dy;
        z[i] = z[i-1] + dz;

        tmp[0] = x[i];
        tmp[1] = y[i];
        tmp[2] = z[i];

        k[0] = 0.0;
        k[1] = 0.0;
        k[2] = 1.0;

        CROSS(normal,k,dir0);

        normalize(normal);

        phi = acos(DOT(k,dir0));

        rotate(tmp2,tmp,normal,phi);

        simvalue->pos[i*3] = pos0[0] + tmp2[0];
        simvalue->pos[i*3+1] = pos0[1] + tmp2[1];
        simvalue->pos[i*3+2] = pos0[2] + tmp2[2];

        tmp[0] = dx;
        tmp[1] = dy;
        tmp[2] = dz;

        normalize(tmp);

        rotate(tmp2,tmp,normal,phi);

        simvalue->dir[i*3] = tmp2[0];
        simvalue->dir[i*3+1] = tmp2[1];
        simvalue->dir[i*3+2] = tmp2[2];
    }

    for (i = 0; i < n; i++) {
        simvalue->z1[i] = path_get_coefficient(i+1,s,theta1,theta0,N);
        simvalue->z2[i] = path_get_coefficient(i+1,s,theta2,theta0,N);
    }

    simvalue->p = path_init(pos0,dir0,1.0,range,1000,theta0,getKineticEnergy,NULL,simvalue->z1,simvalue->z2,n,1.0);

    free(s);
    free(theta1);
    free(theta2);
    free(x);
    free(y);
    free(z);

    return simvalue;
}

void sim_free(sim *simvalue)
{
    free(simvalue->pos);
    free(simvalue->dir);
    free(simvalue->z1);
    free(simvalue->z2);
    path_free(simvalue->p);
    free(simvalue);
}

void usage(void)
{
    fprintf(stderr,"Usage: ./test-path [options]\n");
    fprintf(stderr,"  -N      number of points along track\n");
    fprintf(stderr,"  -z      number of terms in KL expansion\n");
    fprintf(stderr,"  -n      number of simulations\n");
    fprintf(stderr,"  -t      standard deviation of angular distribution\n");
    fprintf(stderr,"  --range range of track\n");
    fprintf(stderr,"  -b      number of bins\n");
    fprintf(stderr,"  -h      display this help message\n");
    exit(1);
}

int main(int argc, char **argv)
{
    size_t i, j, n, N, z, bins;
    double pos0[3], dir0[3], theta0, range, pos2[3], dir[3], beta, t, tmp;
    double *theta1, *theta2;

    z = 2;
    n = 1000;
    N = 1000;
    theta0 = 0.01;
    range = 100.0;
    bins = 100;

    for (i = 1; i < argc; i++) {
        if (!strncmp(argv[i], "--", 2)) {
            if (!strcmp(argv[i]+2,"range")) {
                range = strtod(argv[++i],NULL);
                continue;
            }
        } else if (argv[i][0] == '-') {
            switch (argv[i][1]) {
            case 'N':
                N = atoi(argv[++i]);
                break;
            case 'n':
                n = atoi(argv[++i]);
                break;
            case 'z':
                z = atoi(argv[++i]);
                break;
            case 't':
                theta0 = strtod(argv[++i],NULL);
                break;
            case 'b':
                bins = atoi(argv[++i]);
                break;
            case 'h':
                usage();
            default:
                fprintf(stderr, "unrecognized option '%s'\n", argv[i]);
                exit(1);
            }
        }
    }

    pos0[0] = 0.0;
    pos0[1] = 0.0;
    pos0[2] = 0.0;

    dir0[0] = 1.0;
    dir0[1] = 0.0;
    dir0[2] = 0.0;

    theta1 = calloc(N*n,sizeof(double));
    theta2 = calloc(N*n,sizeof(double));

    gsl_histogram **h1 = malloc(sizeof(gsl_histogram *)*z);
    gsl_histogram **h2 = malloc(sizeof(gsl_histogram *)*z);

    for (i = 0; i < z; i++) {
        h1[i] = gsl_histogram_alloc(bins);
        gsl_histogram_set_ranges_uniform(h1[i],-10.0,10.0);
        h2[i] = gsl_histogram_alloc(bins);
        gsl_histogram_set_ranges_uniform(h2[i],-10.0,10.0);
    }

    init_genrand(0);

    for (i = 0; i < n; i++) {
        sim *simvalue = simulate_path(pos0,dir0,range,theta0,N,z);
        for (j = 0; j < z; j++) {
            gsl_histogram_increment(h1[j], simvalue->z1[j]);
            gsl_histogram_increment(h2[j], simvalue->z2[j]);
        }
        for (j = 0; j < N; j++) {
            path_eval(simvalue->p,range*j/(N-1),pos2,dir,&beta,&t,&tmp);
            theta1[j*n+i] = simvalue->dir[j*3+1] - dir[1];
            theta2[j*n+i] = simvalue->dir[j*3+2] - dir[2];
        }
        sim_free(simvalue);
    }

    FILE *pipe = popen("gnuplot -p", "w");

    if (!pipe) {
        fprintf(stderr, "error running gnuplot command: %s\n", strerror(errno));
        exit(1);
    }

    fprintf(pipe, "set title 'Path Coefficients'\n");
    fprintf(pipe, "plot ");
    for (i = 0; i < z; i++) {
        fprintf(pipe, "'-' title 'z1[%zu]' with steps, \\\n",i);
        fprintf(pipe, "'-' title 'z2[%zu]' with steps",i);
        if (i != z-1) fprintf(pipe, ", \\");
        fprintf(pipe, "\n");
    }

    for (i = 0; i < z; i++) {
        for (j = 0; j < h1[i]->n; j++) {
            fprintf(pipe, "%g %g\n", h1[i]->range[j], h1[i]->bin[j]);
        }
        fprintf(pipe, "e\n");
        for (j = 0; j < h2[i]->n; j++) {
            fprintf(pipe, "%g %g\n", h2[i]->range[j], h2[i]->bin[j]);
        }
        if (i != z-1) fprintf(pipe, "e\n");
    }

    for (i = 0; i < z; i++) {
        gsl_histogram_free(h1[i]);
        gsl_histogram_free(h2[i]);
    }

    free(h1);
    free(h2);

    if (pclose(pipe)) {
        fprintf(stderr, "error closing gnuplot command: %s\n", strerror(errno));
        exit(1);
    }

    pipe = popen("gnuplot -p", "w");

    if (!pipe) {
        fprintf(stderr, "error running gnuplot command: %s\n", strerror(errno));
        exit(1);
    }

    fprintf(pipe, "set title 'Standard Deviation of the Path Direction along the Track'\n");
    fprintf(pipe, "set xlabel 'Distance (cm)'\n");
    fprintf(pipe, "set ylabel '{/Symbol s}({/Symbol q})'\n");
    fprintf(pipe, "plot '-' title 'x' with lines, \\\n");
    fprintf(pipe, "     '-' title 'y' with lines, \\\n");
    fprintf(pipe, "     '-' title 'predicted' with lines\n");
    fprintf(pipe, "\n");
    for (i = 0; i < N; i++) {
        fprintf(pipe, "%g %g\n", i*range/(N-1), std(theta1+n*i,n));
    }
    fprintf(pipe, "e\n");
    for (i = 0; i < N; i++) {
        fprintf(pipe, "%g %g\n", i*range/(N-1), std(theta2+n*i,n));
    }
    fprintf(pipe, "e\n");
    sim *simvalue = simulate_path(pos0,dir0,range,theta0,N,z);
    for (i = 0; i < N; i++) {
        path_eval(simvalue->p,range*i/(N-1),pos2,dir,&beta,&t,&tmp);
        fprintf(pipe, "%g %g\n", i*range/(N-1), tmp);
    }

    if (pclose(pipe)) {
        fprintf(stderr, "error closing gnuplot command: %s\n", strerror(errno));
        exit(1);
    }

    pipe = popen("gnuplot -p", "w");

    if (!pipe) {
        fprintf(stderr, "error running gnuplot command: %s\n", strerror(errno));
        exit(1);
    }

    fprintf(pipe, "set title 'Simulated Path'\n");
    fprintf(pipe, "splot '-' u 1:2:3 title 'simulation' with lines, \\\n");
    fprintf(pipe, "      '-' u 1:2:3 title 'path' with lines\n");
    for (i = 0; i < N; i++) {
        fprintf(pipe,"%.10g %.10g %.10g\n", simvalue->pos[i*3], simvalue->pos[i*3+1], simvalue->pos[i*3+2]);
    }
    fprintf(pipe,"e\n");

    for (i = 0; i < N; i++) {
        path_eval(simvalue->p,i*range/(N-1),pos2,dir,&beta,&t,&theta0);
        fprintf(pipe,"%.10g %.10g %.10g\n", pos2[0], pos2[1], pos2[2]);
    }
    fprintf(pipe,"e\n");

    /* Pause so that you can rotate the 3D graph. */
    fprintf(pipe,"pause mouse keypress\n");

    sim_free(simvalue);

    if (pclose(pipe)) {
        fprintf(stderr, "error closing gnuplot command: %s\n", strerror(errno));
        exit(1);
    }

    return 0;
}