#include "solid_angle.h"
#include <math.h>
#include <stdio.h>
#include "optics.h"
#include "muon.h"
#include "mt19937ar.h"
#include <gsl/gsl_sf_gamma.h>
#include "misc.h"
#include "sno_charge.h"
#include <gsl/gsl_integration.h>
#include "path.h"
#include "random.h"
#include "pdg.h"
#include <gsl/gsl_spline.h>

typedef int testFunction(char *err);

/* Table of some of the tabulated values of the refractive index of water as a
 * function of wavelength and temperature. In all cases, I just used the values
 * for standard atmospheric pressure and assume this corresponds approximately
 * to a density of 1000 kg/m^3.
 *
 * See Table 7 in https://aip.scitation.org/doi/pdf/10.1063/1.555859. */

struct refractive_index_results {
    double p;
    double T;
    double wavelength;
    double n;
} refractive_index_results[] = {
    {1.0, 0 ,  226.50, 1.39468},
    {1.0, 10,  226.50, 1.39439},
    {1.0, 20,  226.50, 1.39353},
    {1.0, 30,  226.50, 1.39224},
    {1.0, 0 ,  404.41, 1.34431},
    {1.0, 10,  404.41, 1.34404},
    {1.0, 20,  404.41, 1.34329},
    {1.0, 30,  404.41, 1.34218},
    {1.0, 0 ,  589.00, 1.33447},
    {1.0, 10,  589.00, 1.33422},
    {1.0, 20,  589.00, 1.33350},
    {1.0, 30,  589.00, 1.33243},
    {1.0, 0 ,  632.80, 1.33321},
    {1.0, 10,  632.80, 1.33296},
    {1.0, 20,  632.80, 1.33224},
    {1.0, 30,  632.80, 1.33118},
    {1.0, 0 , 1013.98, 1.32626},
    {1.0, 10, 1013.98, 1.32604},
    {1.0, 20, 1013.98, 1.32537},
    {1.0, 30, 1013.98, 1.32437},
    {1.0, 0 , 2325.42, 1.27663},
    {1.0, 10, 2325.42, 1.27663},
    {1.0, 20, 2325.42, 1.27627},
    {1.0, 30, 2325.42, 1.27563},
};

/* Table of the values of solid angle for various values of r0/r and L/r.
 *
 * See Table 1 in http://www.umich.edu/~ners312/CourseLibrary/SolidAngleOfADiskOffAxis.pdf. */

struct solid_angle_results {
    double L;
    double r0;
    double omega;
} solid_angle_results[] = {
    {0.5,0.0,3.4732594},
    {0.5,0.2,3.4184435},
    {0.5,0.4,3.2435434},
    {0.5,0.6,2.9185178},
    {0.5,0.8,2.4122535},
    {0.5,1.0,1.7687239},
    {0.5,1.2,1.1661307},
    {0.5,1.4,0.7428889},
    {0.5,1.6,0.4841273},
    {0.5,1.8,0.3287007},
    {0.5,2.0,0.2324189},
    {1.0,0.0,1.8403024},
    {1.0,0.2,1.8070933},
    {1.0,0.4,1.7089486},
    {1.0,0.6,1.5517370},
    {1.0,0.8,1.3488367},
    {1.0,1.0,1.1226876},
    {1.0,1.2,0.9003572},
    {1.0,1.4,0.7039130},
    {1.0,1.6,0.5436956},
    {1.0,1.8,0.4195415},
    {1.0,2.0,0.3257993},
    {1.5,0.0,1.0552591},
    {1.5,0.2,1.0405177},
    {1.5,0.4,0.9975504},
    {1.5,0.6,0.9301028},
    {1.5,0.8,0.8441578},
    {1.5,1.0,0.7472299},
    {1.5,1.2,0.6472056},
    {1.5,1.4,0.5509617},
    {1.5,1.6,0.4632819},
    {1.5,1.8,0.3866757},
    {1.5,2.0,0.3217142},
    {2.0,0.0,0.6633335},
    {2.0,0.2,0.6566352},
    {2.0,0.4,0.6370508},
    {2.0,0.6,0.6060694},
    {2.0,0.8,0.5659755},
    {2.0,1.0,0.5195359},
    {2.0,1.2,0.4696858},
    {2.0,1.4,0.4191714},
    {2.0,1.6,0.3702014},
    {2.0,1.8,0.3243908},
    {2.0,2.0,0.282707}
};

int test_muon_get_energy(char *err)
{
    /* A very simple test to make sure the energy as a function of distance
     * along the track makes sense. Should include more detailed tests later. */
    double T0, T, range;
    muon_energy *m;

    /* Assume initial kinetic energy is 1 GeV. */
    T0 = 1000.0;
    m = muon_init_energy(T0,1.0,10000);
    T = muon_get_energy(1e-9,m);

    /* At the beginning of the track we should have roughly the same energy. */
    if (!isclose(T, T0, 1e-5, 0)) {
        sprintf(err, "KE = %.5f, but expected %.5f", T, T0);
        return 1;
    }

    range = get_range(T0,1.0);
    T = muon_get_energy(range,m);

    /* At the end of the track the energy should be approximately 0. */
    if (!isclose(T, 0, 1e-5, 1e-5)) {
        sprintf(err, "KE = %.5f, but expected %.5f", T, 0.0);
        return 1;
    }

    return 0;
}

int test_muon_get_range(char *err)
{
    /* A very simple test to make sure we read in the PDG table correctly. */
    double value;

    value = get_range(1.0,1.0);

    if (!isclose(value, 2.071e-3, 1e-5, 0)) {
        sprintf(err, "range = %.5f, but expected %.5f", value, 1.863e-3);
        return 1;
    }

    return 0;
}

int test_muon_get_dEdx(char *err)
{
    /* A very simple test to make sure we read in the PDG table correctly. */
    double value;

    value = get_dEdx(1.0,1.0);

    if (!isclose(value, 5.471, 1e-5, 0)) {
        sprintf(err, "dE/dx = %.5f, but expected %.5f", value, 6.097);
        return 1;
    }

    return 0;
}

int test_refractive_index(char *err)
{
    /* Tests the get_index() function. */
    int i;
    double n;
    struct refractive_index_results result;

    for (i = 0; i < sizeof(refractive_index_results)/sizeof(struct refractive_index_results); i++) {
        result = refractive_index_results[i];
        n = get_index(result.p, result.wavelength, result.T);

        if (!isclose(n, result.n, 1e-2, 0)) {
            sprintf(err, "n = %.5f, but expected %.5f", n, result.n);
            return 1;
        }
    }

    return 0;
}

int test_solid_angle_approx(char *err)
{
    /* Tests the get_solid_angle_approx() function. */
    int i;
    double pmt[3] = {0,0,0};
    double pos[3] = {0,0,1};
    double n[3] = {0,0,1};
    double r = 1.0;
    double solid_angle;

    solid_angle = get_solid_angle_approx(pos,pmt,n,r);

    if (!isclose(solid_angle, 2*M_PI*(1-1/sqrt(2)), 1e-2, 0)) {
        sprintf(err, "solid angle = %.5f, but expected %.5f", solid_angle, 2*M_PI*(1-1/sqrt(2)));
        return 1;
    }

    for (i = 0; i < sizeof(solid_angle_results)/sizeof(struct solid_angle_results); i++) {
        pos[0] = solid_angle_results[i].r0*r;
        pos[2] = solid_angle_results[i].L*r;

        solid_angle = get_solid_angle_approx(pos,pmt,n,r);

        if (!isclose(solid_angle, solid_angle_results[i].omega, 1e-2, 0)) {
            sprintf(err, "solid angle = %.5f, but expected %.5f", solid_angle, solid_angle_results[i].omega);
            return 1;
        }
    }

    return 0;
}

int test_solid_angle(char *err)
{
    /* Tests the get_solid_angle() function. */
    int i;
    double pmt[3] = {0,0,0};
    double pos[3] = {0,0,1};
    double n[3] = {0,0,1};
    double r = 1.0;
    double solid_angle;

    solid_angle = get_solid_angle(pos,pmt,n,r);

    if (!isclose(solid_angle, 2*M_PI*(1-1/sqrt(2)), 1e-9, 0)) {
        sprintf(err, "solid angle = %.5f, but expected %.5f", solid_angle, 2*M_PI*(1-1/sqrt(2)));
        return 1;
    }

    for (i = 0; i < sizeof(solid_angle_results)/sizeof(struct solid_angle_results); i++) {
        pos[0] = solid_angle_results[i].r0*r;
        pos[2] = solid_angle_results[i].L*r;

        solid_angle = get_solid_angle(pos,pmt,n,r);

        if (!isclose(solid_angle, solid_angle_results[i].omega, 1e-4, 0)) {
            sprintf(err, "solid angle = %.5f, but expected %.5f", solid_angle, solid_angle_results[i].omega);
            return 1;
        }
    }

    return 0;
}

static double sno_charge(double q, void *params)
{
    int n = ((int *) params)[0];
    return pq(q,n);
}

int test_sno_charge_integral(char *err)
{
    /* Test that the single PE charge distribution integrates to one. */
    double result, error;
    gsl_function F;
    size_t nevals;
    gsl_integration_cquad_workspace *w;
    int params[1];

    w = gsl_integration_cquad_workspace_alloc(100);

    F.function = &sno_charge;
    params[0] = 1;
    F.params = params;

    init_charge();

    gsl_integration_cquad(&F, -10.0, 1000.0, 0, 1e-9, w, &result, &error, &nevals);

    if (!isclose(result, 1.0, 1e-9, 1e-3)) {
        sprintf(err, "integral of single PE charge distribution is %.2f", result);
        goto err;
    }

    gsl_integration_cquad_workspace_free(w);

    return 0;

err:
    gsl_integration_cquad_workspace_free(w);
    return 1;
}

int test_logsumexp(char *err)
{
    /* Tests the logsumexp() function. */
    int i, j;
    double logp[100], mu, result, expected;

    init_genrand(0);

    for (i = 0; i < 100; i++) {
        mu = genrand_real2();

        for (j = 0; j < sizeof(logp)/sizeof(double); j++) {
            logp[j] = -mu + j*log(mu) - gsl_sf_lnfact(j);
        }
        result = logsumexp(logp, sizeof(logp)/sizeof(double));

        expected = 0.0;

        if (!isclose(result, expected, 1e-9, 1e-9)) {
            sprintf(err, "logsumexp(logp) = %.5g, but expected %.5g", result, expected);
            return 1;
        }
    }

    return 0;
}

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

int test_path(char *err)
{
    /* Tests the logsumexp() function. */
    int i, j, k;
    double pos0[3], dir0[3], T0, range, m;
    double T, t;
    double pos_expected[3], t_expected;
    double pos[3], dir[3];
    double E, mom, beta;
    double s;
    double theta0;

    T0 = 1.0;
    range = 1.0;
    m = 1.0;

    init_genrand(0);

    for (i = 0; i < 100; i++) {
        pos0[0] = genrand_real2();
        pos0[1] = genrand_real2();
        pos0[2] = genrand_real2();

        rand_sphere(dir0);

        path *p = path_init(pos0,dir0,T0,range,0.1,getKineticEnergy,NULL,NULL,NULL,0,m);

        for (j = 0; j < 100; j++) {
            s = range*j/99;
            pos_expected[0] = pos0[0] + dir0[0]*s;
            pos_expected[1] = pos0[1] + dir0[1]*s;
            pos_expected[2] = pos0[2] + dir0[2]*s;

            /* Calculate total energy */
            E = T0 + m;
            mom = sqrt(E*E - m*m);
            beta = mom/E;

            t_expected = s/(beta*SPEED_OF_LIGHT);

            path_eval(p,s,pos,dir,&T,&t,&theta0);

            for (k = 0; k < 3; k++) {
                if (!isclose(pos[k], pos_expected[k], 1e-9, 1e-9)) {
                    sprintf(err, "path_eval(%.2g) returned pos[%i] = %.5g, but expected %.5g", s, k, pos[k], pos_expected[k]);
                    return 1;
                }
            }

            for (k = 0; k < 3; k++) {
                if (!isclose(dir[k], dir0[k], 1e-9, 1e-9)) {
                    sprintf(err, "path_eval(%.2g) returned dir[%i] = %.5g, but expected %.5g", s, k, dir[k], dir0[k]);
                    return 1;
                }
            }

            if (!isclose(T, 1.0, 1e-9, 1e-9)) {
                sprintf(err, "path_eval(%.2g) returned T = %.5g, but expected %.5g", s, T, 1.0);
                return 1;
            }

            if (!isclose(t, t_expected, 1e-9, 1e-9)) {
                sprintf(err, "path_eval(%.2g) returned t = %.5g, but expected %.5g", s, t, t_expected);
                return 1;
            }
        }

        path_free(p);
    }

    return 0;
}

int test_interp1d(char *err)
{
    /* Tests the interp1d() function. */
    size_t i;
    double xp[100], yp[100], range, x, y, expected;
    gsl_interp_accel *acc;
    gsl_spline *spline;

    range = 1.0;

    init_genrand(0);

    for (i = 0; i < sizeof(xp)/sizeof(xp[0]); i++) {
        xp[i] = range*i/(100-1);
        yp[i] = genrand_real2();
    }

    acc = gsl_interp_accel_alloc();
    spline = gsl_spline_alloc(gsl_interp_linear,100);
    gsl_spline_init(spline,xp,yp,100);

    for (i = 0; i < 1000; i++) {
        x = genrand_real2()*range;
        expected = gsl_spline_eval(spline,x,acc);
        y = interp1d(x,xp,yp,100);
        if (!isclose(y, expected, 1e-9, 1e-9)) {
            sprintf(err, "interp1d(%.2g) returned %.5g, but expected %.5g", x, y, expected);
            goto err;
        }
    }

    gsl_interp_accel_free(acc);
    gsl_spline_free(spline);

    return 0;

err:
    gsl_interp_accel_free(acc);
    gsl_spline_free(spline);

    return 1;
}

int test_kahan_sum(char *err)
{
    /* Tests the kahan_sum() function. */
    size_t i;
    double x[100], sum, expected;

    init_genrand(0);

    expected = 0.0;
    for (i = 0; i < sizeof(x)/sizeof(x[0]); i++) {
        x[i] = genrand_real2();
        expected += x[i];
    }

    sum = kahan_sum(x,sizeof(x)/sizeof(x[0]));

    if (!isclose(sum, expected, 1e-9, 1e-9)) {
        sprintf(err, "kahan_sum returned %.5g, but expected %.5g", sum, expected);
        goto err;
    }

    return 0;

err:
    return 1;
}

int test_lnfact(char *err)
{
    /* Tests the lnfact() function. */
    size_t i;
    double x, expected;

    for (i = 0; i < 1000; i++) {
        x = lnfact(i);
        expected = gsl_sf_lnfact(i);

        if (!isclose(x, expected, 1e-9, 1e-9)) {
            sprintf(err, "lnfact(%zu) returned %.5g, but expected %.5g", i, x, expected);
            goto err;
        }
    }

    return 0;

err:
    return 1;
}

int test_ln(char *err)
{
    /* Tests the lnfact() function. */
    size_t i;
    double x, expected;

    for (i = 1; i < 1000; i++) {
        x = ln(i);
        expected = log(i);

        if (!isclose(x, expected, 1e-9, 1e-9)) {
            sprintf(err, "ln(%zu) returned %.5g, but expected %.5g", i, x, expected);
            goto err;
        }
    }

    return 0;

err:
    return 1;
}

struct tests {
    testFunction *test;
    char *name;
} tests[] = {
    {test_solid_angle, "test_solid_angle"},
    {test_solid_angle_approx, "test_solid_angle_approx"},
    {test_refractive_index, "test_refractive_index"},
    {test_muon_get_dEdx, "test_muon_get_dEdx"},
    {test_muon_get_range, "test_muon_get_range"},
    {test_muon_get_energy, "test_muon_get_energy"},
    {test_logsumexp, "test_logsumexp"},
    {test_sno_charge_integral, "test_sno_charge_integral"},
    {test_path, "test_path"},
    {test_interp1d, "test_interp1d"},
    {test_kahan_sum, "test_kahan_sum"},
    {test_lnfact, "test_lnfact"},
    {test_ln, "test_ln"},
};

int main(int argc, char **argv)
{
    int i;
    char err[256];
    int retval = 0;
    struct tests test;

    for (i = 0; i < sizeof(tests)/sizeof(struct tests); i++) {
        test = tests[i];

        if (!test.test(err)) {
            printf("[\033[92mok\033[0m] %s\n", test.name);
        } else {
            printf("[\033[91mfail\033[0m] %s: %s\n", test.name, err);
            retval = 1;
        }
    }

    return retval;
}