#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. */
#include "misc.h"

#define N 1000

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;
    double normal[3], k[3], tmp[3], tmp2[3];

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

    p->theta0 = theta0;

    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];
        }
    }

    /* Now, we rotate and translate the path so that it starts at `pos` and
     * points in the direction `dir`. */

    /* We need to compute an arbitrary vector which is orthogonal to the
     * direction vector. To do this, all we need is another vector not parallel
     * to the direction. */
    k[0] = 0.0;
    k[1] = 0.0;
    k[2] = 1.0;

    /* If k is approximately equal to the unit vector in the z direction, we
     * switch to the unit vector in the x direction. */
    if (allclose(k,dir,3,1e-9,1e-9)) {
        k[0] = 1.0;
        k[1] = 0.0;
        k[2] = 0.0;
    }

    /* Take the cross product between `k` and `dir` to get a vector orthogonal
     * to `dir`. */
    CROSS(normal,k,dir);

    /* Make sure it's normalized. */
    normalize(normal);

    /* Compute the angle required to rotate the unit vector to `dir`. */
    phi = acos(DOT(k,dir));

    /* Now we rotate and translate all the positions and rotate all the
     * directions. */
    for (i = 0; i < N; i++) {
        tmp[0] = x[i];
        tmp[1] = y[i];
        tmp[2] = z[i];

        rotate(tmp2,tmp,normal,phi);

        ADD(tmp2,tmp2,pos);

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

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

        normalize(tmp);

        rotate(tmp2,tmp,normal,phi);

        dx[i] = tmp2[0];
        dy[i] = tmp2[1];
        dz[i] = tmp2[2];
    }

    p->s = s;
    p->x = x;
    p->y = y;
    p->z = z;
    p->T = T;
    p->t = t;
    p->dx = dx;
    p->dy = dy;
    p->dz = dz;

    free(theta1);
    free(theta2);

    return p;
}

void path_eval(path *p, double s, double *pos, double *dir, double *T, double *t, double *theta0)
{
    pos[0] = interp1d(s,p->s,p->x,N);
    pos[1] = interp1d(s,p->s,p->y,N);
    pos[2] = interp1d(s,p->s,p->z,N);

    *T = interp1d(s,p->s,p->T,N);
    *t = interp1d(s,p->s,p->t,N);

    dir[0] = interp1d(s,p->s,p->dx,N);
    dir[1] = interp1d(s,p->s,p->dy,N);
    dir[2] = interp1d(s,p->s,p->dz,N);

    /* FIXME: This should be the *residual* scattering RMS. */
    *theta0 = fmax(MIN_THETA0, p->theta0*sqrt(s));
}

void path_free(path *p)
{
    free(p->s);
    free(p->x);
    free(p->y);
    free(p->z);
    free(p->T);
    free(p->t);
    free(p->dx);
    free(p->dy);
    free(p->dz);

    free(p);
}