/* 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 "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"
#include <gsl/gsl_sf_psi.h> /* for gsl_sf_psi_n() */
#include "scattering.h"

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, size_t len, double theta0, getKineticEnergyFunc *fun, void *params, double *z1, double *z2, size_t n, double m)
{
    size_t i, j;
    double T, E, mom, theta, phi;
    double normal[3], k[3], tmp[3], tmp2[3];

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

    p->n = n;
    p->len = len;

    double *s = calloc(len,sizeof(double));
    double *theta1 = calloc(len,sizeof(double));
    double *theta2 = calloc(len,sizeof(double));
    double *x = calloc(len,sizeof(double));
    double *y = calloc(len,sizeof(double));
    double *z = calloc(len,sizeof(double));
    double *beta = calloc(len,sizeof(double));
    double *t = calloc(len,sizeof(double));
    double *dx = calloc(len,sizeof(double));
    double *dy = calloc(len,sizeof(double));
    double *dz = calloc(len,sizeof(double));

    dz[0] = 1.0;
    for (i = 0; i < len; i++) {
        s[i] = range*i/(len-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 = fun(s[i],params);
        E = T + m;
        mom = sqrt(E*E - m*m);
        beta[i] = mom/E;

        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 */
            if (beta[i] > 0)
                t[i] = t[i-1] + (s[i]-s[i-1])/(beta[i]*SPEED_OF_LIGHT);
            else
                t[i] = t[i-1];

            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 (allclose(k,dir,3,0,0)) {
        /* If the direction vector is already in the z direction, we don't need
         * to rotate it, so the normal direction is arbitrary since `phi` will
         * be 0. */
        normal[0] = 1.0;
        normal[1] = 0.0;
        normal[2] = 0.0;
    } else {
        /* Take the cross product between `k` and `dir` to get the axis of
         * rotation. */
        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 < len; 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];
    }

    /* Calculate the approximate residual scattering RMS.
     *
     * For a given truncated KL expansion the residual error in the
     * approximation is given by:
     *
     * \sum_{n+1}^\infty sqrt(lambda_k)*sqrt(2/range)*sin(pi*s*(k-0.5)/range)
     *
     * Since there is no nice closed form solution for this, we estimate it by
     * evaluating the result at the middle of the track. */
    if (p->n > 0) {
        p->theta0 = theta0*sqrt(range)*sqrt(gsl_sf_psi_n(1,p->n+0.5))/M_PI;
        p->theta0 = fmax(MIN_THETA0, p->theta0);
        p->theta0 = fmin(MAX_THETA0, p->theta0);
    } else {
        p->theta0 = theta0;
    }

    p->s = s;
    p->x = x;
    p->y = y;
    p->z = z;
    p->beta = beta;
    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 *beta, double *t, double *theta0)
{
    pos[0] = interp1d(s,p->s,p->x,p->len);
    pos[1] = interp1d(s,p->s,p->y,p->len);
    pos[2] = interp1d(s,p->s,p->z,p->len);

    *beta = interp1d(s,p->s,p->beta,p->len);
    *t = interp1d(s,p->s,p->t,p->len);

    /* Technically we should be interpolating between the direction vectors
     * using an algorithm like Slerp (see https://en.wikipedia.org/wiki/Slerp),
     * but since we expect the direction vectors to be pretty close to each
     * other, we just linearly interpolate and then normalize it. */

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

    normalize(dir);

    if (p->n > 0) {
        *theta0 = p->theta0;
    } else {
        *theta0 = fmax(MIN_THETA0,p->theta0*sqrt(s));
        *theta0 = fmin(MAX_THETA0,*theta0);
    }
}

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

    free(p);
}