#include "path.h"
#include <math.h>
#include <gsl/gsl_spline.h>
#include <stddef.h> /* for size_t */
#include <gsl/gsl_integration.h>
#include "pdg.h"
#include "vector.h"
#include <stdlib.h> /* for malloc(), calloc(), etc. */

#define N 100

static double foo(double s, double range, double theta0, int k)
{
    return sqrt(2*pow(theta0,2)*range)*sin(M_PI*s*(k-0.5)/range)/(M_PI*(k-0.5));
}

double path_get_coefficient(unsigned int k, double *s, double *x, double theta0, size_t n)
{
    size_t i;
    double sum, range;

    range = s[n-1];

    sum = 0.0;
    for (i = 1; i < n; i++) {
        sum += (foo(s[i],range,theta0,k)*x[i] + foo(s[i-1],range,theta0,k)*x[i-1])*(s[i]-s[i-1])/2.0;
    }

    return sum*pow(k-0.5,2)*pow(M_PI,2)/pow(theta0*range,2);
}

path *path_init(double *pos, double *dir, double T0, double range, double theta0, getKineticEnergyFunc *fun, double *z1, double *z2, size_t n, double m)
{
    size_t i, j;
    double E, mom, beta, theta, phi;

    path *p = malloc(sizeof(path));

    p->pos[0] = pos[0];
    p->pos[1] = pos[1];
    p->pos[2] = pos[2];

    p->dir[0] = dir[0];
    p->dir[1] = dir[1];
    p->dir[2] = dir[2];

    double *s = malloc(sizeof(double)*N);
    double *theta1 = calloc(N,sizeof(double));
    double *theta2 = calloc(N,sizeof(double));
    double *x = calloc(N,sizeof(double));
    double *y = calloc(N,sizeof(double));
    double *z = calloc(N,sizeof(double));
    double *T = calloc(N,sizeof(double));
    double *t = calloc(N,sizeof(double));
    double *dx = calloc(N,sizeof(double));
    double *dy = calloc(N,sizeof(double));
    double *dz = calloc(N,sizeof(double));

    dz[0] = 1.0;
    for (i = 0; i < N; i++) {
        s[i] = range*i/(N-1);
        for (j = 0; j < n; j++) {
            theta1[i] += z1[j]*foo(s[i],range,theta0,j+1);
            theta2[i] += z2[j]*foo(s[i],range,theta0,j+1);
        }
        T[i] = fun(s[i],T0);
        if (i > 0) {
            theta = sqrt(theta1[i]*theta1[i] + theta2[i]*theta2[i]);
            phi = atan2(theta2[i],theta1[i]);

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

            /* Calculate total energy */
            E = T[i] + m;
            mom = sqrt(E*E - m*m);
            beta = mom/E;

            t[i] = t[i-1] + (s[i]-s[i-1])/(beta*SPEED_OF_LIGHT);

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

    for (i = 0; i < 8; i++) {
        p->acc[i] = gsl_interp_accel_alloc();
        p->spline[i] = gsl_spline_alloc(gsl_interp_linear,N);
    }

    gsl_spline_init(p->spline[0],s,x,N);
    gsl_spline_init(p->spline[1],s,y,N);
    gsl_spline_init(p->spline[2],s,z,N);
    gsl_spline_init(p->spline[3],s,T,N);
    gsl_spline_init(p->spline[4],s,t,N);
    gsl_spline_init(p->spline[5],s,dx,N);
    gsl_spline_init(p->spline[6],s,dy,N);
    gsl_spline_init(p->spline[7],s,dz,N);

    free(s);
    free(theta1);
    free(theta2);
    free(x);
    free(y);
    free(z);
    free(T);
    free(t);
    free(dx);
    free(dy);
    free(dz);

    return p;
}

void path_eval(path *p, double s, double *pos, double *dir, double *T, double *t)
{
    double normal[3], k[3], tmp[3], phi;

    tmp[0] = gsl_spline_eval(p->spline[0],s,p->acc[0]);
    tmp[1] = gsl_spline_eval(p->spline[1],s,p->acc[1]);
    tmp[2] = gsl_spline_eval(p->spline[2],s,p->acc[2]);

    *T = gsl_spline_eval(p->spline[3],s,p->acc[3]);
    *t = gsl_spline_eval(p->spline[4],s,p->acc[4]);

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

    CROSS(normal,k,p->dir);

    normalize(normal);

    phi = acos(DOT(k,p->dir));

    rotate(pos,tmp,normal,phi);

    ADD(pos,pos,p->pos);

    tmp[0] = gsl_spline_eval(p->spline[5],s,p->acc[5]);
    tmp[1] = gsl_spline_eval(p->spline[6],s,p->acc[6]);
    tmp[2] = gsl_spline_eval(p->spline[7],s,p->acc[7]);

    normalize(tmp);

    rotate(dir,tmp,normal,phi);
}

void path_free(path *p)
{
    size_t i;

    for (i = 0; i < 5; i++) {
        gsl_interp_accel_free(p->acc[i]);
        gsl_spline_free(p->spline[i]);
    }

    free(p);
}