#include <algorithm>
#include <array>
#include <boost/numeric/odeint.hpp>
#include <Eigen/Geometry> 
#include <fstream>
#include <gsl/gsl_spline.h>
#include <iostream>
#include <numeric>
#include <string>
#include <stdexcept>
#include <vector>

#include "loadtxt.h"

using double3 = Eigen::Vector3d;

// Our units are {kiloparsec, solar mass, gigayear}
constexpr double G = 4.498317481097514e-06;

class Interp {
    // This is a wrapper around GSL spline interpolation. I tried to use
    // boost::math::interpolators but as of version 1.72 none were suitable:
    // barycentric_rational is the one suitalbe for non-uniform sampling but it
    // is very slow. I also tried to resample the data uniformly using
    // barycentric rational interpolation and then using cardinal_cubic_b_spline
    // on the result, but was still slower than GSL.
public:
    Interp(std::vector<double>& x, std::vector<double>& y)
    {
        acc    = gsl_interp_accel_alloc();
        spline = gsl_spline_alloc(gsl_interp_cspline, x.size());
        gsl_spline_init(spline, x.data(), y.data(), x.size());
    }
    Interp() {}
    inline double operator()(double x) const
    {
        return gsl_spline_eval(spline, x, acc);
    }
private:
    gsl_interp_accel *acc;
    gsl_spline *spline;
};

double3 plummer(const double M, const double b, const double3& pos)
{
    double r2 = pos.squaredNorm() + b*b;
    double r  = sqrt(r2);
    double r3_inv = 1/(r*r2);
    return -G*M*pos*r3_inv;
}

double3 nfw(const double rho_0, const double b, const double3& pos)
{
    double r2 = pos.squaredNorm();
    double r  = sqrt(r2);
    double tmp = -4*M_PI*G*rho_0*b*b*b*(log1p(r/b) - r/(b+r))/(r2*r);
    return pos*tmp;
}

double3 miyamoto_nagai(const double M, const double a, const double b, const double phi, const double theta, const double3& pos)
{
    // Construct new z-axis from the angles
    double cos_theta = cos(theta);
    double sin_theta = sqrt(1-cos_theta*cos_theta);
    double cos_phi   = cos(phi);
    double sin_phi   = sqrt(1-cos_phi*cos_phi);
    Eigen::Vector3d old_z_axis = {cos_phi*sin_theta, sin_phi*sin_theta, cos_theta};

    // Construct rotation
    auto rot = Eigen::Quaternion<double>::FromTwoVectors(old_z_axis, Eigen::Vector3d::UnitZ());

    // Rotate the position vector
    auto new_pos = rot * pos;

    // Calculate acceleration in new frame
    Eigen::Vector3d acc_in_rotated_frame;
    auto z_tmp  = sqrt(new_pos[2]*new_pos[2] + b*b);
    auto r2_tmp = new_pos[0]*new_pos[0] + new_pos[1]*new_pos[1] + (z_tmp + a)*(z_tmp + a);
    auto r_tmp  = sqrt(r2_tmp);
    auto tmp = G*M / (r2_tmp*r_tmp);
    acc_in_rotated_frame[0] = - tmp * new_pos[0];
    acc_in_rotated_frame[1] = - tmp * new_pos[1];
    acc_in_rotated_frame[2] = - tmp * new_pos[2] * (z_tmp + a)/z_tmp;

    // Return to original frame
    return rot.inverse() * acc_in_rotated_frame;
}

class Galaxy {
public:
    Galaxy(std::string file_name)
    {
        auto data = Loadtxt(file_name, {1, 2, 3, 4, 5, 6, 7, 8}).get_cols();
        auto& t_data = data[0];
        std::transform(begin(t_data), end(t_data), begin(t_data), [t0=t_data[0]](const double& x){ return x-t0; });
        for (int i=0; i<7; i++) interp[i] = Interp(t_data, data[i+1]);
    }
    void func(const std::array<double, 6> &y, std::array<double, 6> &f, const double t)
    {
        f[0] = y[3]; // vx -> x'
        f[1] = y[4]; // vy -> y'
        f[2] = y[5]; // vz -> z'
        double phi         = interp[0](t);
        double theta       = interp[1](t);
        double m_disk      = interp[2](t);
        double a_disk      = interp[3](t);
        double b_disk      = interp[4](t);
        double rho_0_nfw   = interp[5](t);
        double b_nfw       = interp[6](t);

        double3 pos(y.data());
        double3 acc_disk = miyamoto_nagai(m_disk, a_disk, b_disk, phi, theta, pos);
        double3 acc_halo = nfw(rho_0_nfw, b_nfw, pos);
        for (int i=0; i<3; i++) f[3+i] = acc_disk[i] + acc_halo[i];
    }
private:
    Interp interp[10];
} galaxy("subhalo_411321_parameters_processed.dat");

void integrate(const std::array<double,6> y0, const double t_max, const double stride_size, double y[])
// NOTE we're not responsible to the memory of the output parameter y. Caller should allocate
{
    using namespace boost::numeric::odeint;
    using Coordinates = std::array<double, 6>;
    auto stepper = bulirsch_stoer<Coordinates>(1E-7, 0);
    auto function_wrapper = [](const Coordinates &x, Coordinates &dxdt, const double t) { return galaxy.func(x, dxdt, t); };
    const int stride_count = t_max / stride_size;
    Coordinates y_current;
    std::copy(begin(y0), end(y0), begin(y_current));
    std::copy(begin(y0), end(y0), y);
    double t = 0;
    const double h = std::min(1./4096., stride_size);
    for (int i=0; i<stride_count; i++) {
        integrate_adaptive(stepper, function_wrapper, y_current, t, t+stride_size, h);
        // NOTE h here is just a recommended initial step size for the stepper,
        // the actual step is adapted as needed. Since the integration is
        // interrupted here in order for the data to be saved, the result
        // somewhat depends on stride_size.
        std::copy(begin(y_current), end(y_current), y+(i+1)*6);
        t += stride_size;
    }
}

extern "C"
void integrate(const double y0[], const double t_max, const double stride_size, double y[])
{
    std::array<double,6> y0_array;
    std::copy(y0, y0+6, begin(y0_array));
    integrate(y0_array, t_max, stride_size, y);
}