#include <cstdio>
#include <numeric>
#include "black_holes.h"

/* BEGIN legacy inclusion */
// I'm not going to touch this C file
#define SQR(x) ((x)*(x))
double L[3]; // needed in pn_bh_spin.c
#include "pn_bh_spin.c"
#undef SQR
/* END legacy inclusion */

void two_body_gravity(const Particle_ref& j, const Particle_ref& i, const double eps, double& pot, double3& acc, double3& jrk)
{
    double3 dx = i.x - j.x;
    double3 dv = i.v - j.v;
    double r2 = dx.norm2() + eps*eps;
    double r  = sqrt(r2);
    double r3 = r2*r;
    double r4 = r2*r2;
    double RP = 3*(dx*dv)/r;

    pot = -j.m/r;
    acc = -j.m*dx/r3;
    jrk = -j.m*(dv/r3 - RP*dx/r4);
}

void Black_hole_physics::adjust_softening(const std::vector<Particle_ref>& particles)
{
    if (eps_new < 0) return;
    double pot_old, pot_new;
    double3 acc_old, acc_new, jrk_old, jrk_new;
    for (int j = 0; j < particles.size(); j++) {
        for (int i = 0; i < particles.size(); i++) {
            if (i == j) continue;
            two_body_gravity(particles[j], particles[i], eps_old, pot_old, acc_old, jrk_old);
            two_body_gravity(particles[j], particles[i], eps_new, pot_new, acc_new, jrk_new);
            particles[i].pot += pot_new - pot_old;
            particles[i].acc += acc_new - acc_old;
            particles[i].jrk += jrk_new - jrk_old;
        }
    }
}

#if 0
void Black_hole_physics::adjust_post_newtonian(
    const double dt_bh, // pn_usage should be const
    double3& acc1, double3& acc2,
    double3& jrk1, double3& jrk2)
{
    // calculate and "plus" the new BH <-> BH : PN1, PN2, PN2.5, PN3, PN3.5 : acc, jerk
    // TODO maybe have the PN terms as local variables here?
    int tmp;
    tmp = calc_force_pn_BH(masses[0], x1, v1, bbh_grav.spin1,
                           masses[1], x2, v2, bbh_grav.spin2,
                           c,  dt_bh, pn_usage,
                           (double(*)[3])bbh_grav.a_pn1, (double(*)[3])bbh_grav.adot_pn1,
                           (double(*)[3])bbh_grav.a_pn2, (double(*)[3])bbh_grav.adot_pn2, myRank, rootRank);
    if (tmp == 505) exit(-1); // Very ugly way to terminate

    // NOTE we have these _corr variables accumulating the corrections before
    // applying it. It's almost the same but different from applying each
    // correction term in a loop.
    double3 acc1_corr(0,0,0), acc2_corr(0,0,0), jrk1_corr(0,0,0), jrk2_corr(0,0,0);
    for (int i=1; i<7; i++) {
        acc1_corr += bbh_grav.a_pn1[i];
        acc2_corr += bbh_grav.a_pn2[i];
        jrk1_corr += bbh_grav.adot_pn1[i];
        jrk2_corr += bbh_grav.adot_pn2[i];
    }
    acc1 += acc1_corr;
    acc2 += acc2_corr;
    jrk1 += jrk1_corr;
    jrk2 += jrk2_corr;
}
#endif

void Black_hole_physics::write_bh_data(const double time_cur, const int count, const std::vector<double> &m, const std::vector<double3> &x, const std::vector<double3> &v, const std::vector<double> &pot, const std::vector<double3> &a, const std::vector<double3> &adot, const std::vector<double> &dt)
{
    auto out = fopen("bh.dat", "a");
    for (int i = 0; i < count; i++) {
        fprintf(out,"%.16E \t %.8E \t % .16E % .16E % .16E \t %.16E \t % .16E % .16E % .16E \t %.16E \t % .8E \t % .8E % .8E % .8E \t %.8E \t % .8E % .8E % .8E \t %.8E \t %.8E \n",
            time_cur, m[i],
            x[i][0], x[i][1], x[i][2], x[i].norm(),
            v[i][0], v[i][1], v[i][2], v[i].norm(),
            pot[i],
            a[i][0],    a[i][1],    a[i][2],    a[i].norm(),
            adot[i][0], adot[i][1], adot[i][2], adot[i].norm(),
            dt[i]);
    }
    fprintf(out, "\n");
    fclose(out);
}

void Write_bh_nb_data::operator()(double time_cur)
{
    for (int i_bh=0; i_bh < smbh_count; i_bh++) {
        for (int i=0; i<N; i++) var_sort[i] = (x[i]-x[i_bh]).norm();
        std::iota(ind_sort.begin(), ind_sort.end(), 0);
        std::partial_sort(ind_sort.begin(), ind_sort.begin() + nb, ind_sort.begin() + N, [&](int i, int j) {return var_sort[i] < var_sort[j];});

        fprintf(out,"%.16E \t %07d \t %.8E \t % .8E % .8E % .8E \t % .8E % .8E % .8E \t",
                time_cur,
                i_bh,
                m[i_bh],
                x[i_bh][0], x[i_bh][1], x[i_bh][2],
                v[i_bh][0], v[i_bh][1], v[i_bh][2]);

        for (int j=1; j < nb; j++) {
            int i = ind_sort[j];
            fprintf(out,"%02d %07d   %.8E   % .8E % .8E % .8E   %.8E   % .8E % .8E % .8E   %.8E \t",
                    j, i,
                    m[i],
                    x[i][0], x[i][1], x[i][2], (x[i]-x[i_bh]).norm(),
                    v[i][0], v[i][1], v[i][2], (v[i]-v[i_bh]).norm());
        }
        fprintf(out,"\n");
    }
    fprintf(out, "\n"); // this is redundant
    fflush(out);
}

void Binary_smbh_influence_sphere_output::operator()(const std::vector<int>& ind_act, int n_act, double timesteps, double time_cur)
{
    double m_bh1  = m[0];
    double m_bh2  = m[1];
    double3 x_bh1 = x[0];
    double3 x_bh2 = x[1];
    double3 v_bh1 = v[0];
    double3 v_bh2 = v[1];

    double3 x_bbhc = (m_bh1*x_bh1 + m_bh2*x_bh2)/(m_bh1 + m_bh2);
    double3 v_bbhc = (m_bh1*v_bh1 + m_bh2*v_bh2)/(m_bh1 + m_bh2);

    double DR2 = (x_bh1 - x_bh2).norm2();
    double DV2 = (v_bh1 - v_bh2).norm2();
    double EB = -(m_bh1 + m_bh2) / sqrt(DR2) + 0.5 * DV2;
    double SEMI_a = -0.5 * (m_bh1 + m_bh2)/EB;
    double SEMI_a2 = SEMI_a*SEMI_a;

    for (int i=0; i<n_act; i++) {
        int j_act = ind_act[i];
        if (j_act<2) continue;
        const double& pot_bh1 = pot[0];
        const double& pot_bh2 = pot[1];
        const double& m_act   = m[j_act];
        const double3& x_act  = x[j_act];
        const double3& v_act  = v[j_act];
        const double& dt_act  = dt[j_act];
        const double& pot_act = pot[j_act];
        double tmp_r2 = (x_act - x_bbhc).norm2();
        if (tmp_r2 < SEMI_a2*factor*factor) {
            if (inf_event[j_act] == 0) {
                fprintf(out,"INF1 %.6E %.16E %07d %07d   %.6E % .6E % .6E % .6E % .6E % .6E % .6E   %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E  %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E   %.6E   %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E  %.6E \n",
                    timesteps, time_cur, i, j_act,
                    sqrt(DR2), x_bbhc[0], x_bbhc[1], x_bbhc[2], v_bbhc[0], v_bbhc[1], v_bbhc[2],
                    m_bh1, x_bh1[0], x_bh1[1], x_bh1[2], v_bh1[0], v_bh1[1], v_bh1[2], pot_bh1,
                    m_bh2, x_bh2[0], x_bh2[1], x_bh2[2], v_bh2[0], v_bh2[1], v_bh2[2], pot_bh2,
                    sqrt(tmp_r2),
                    m_act, x_act[0], x_act[1], x_act[2], v_act[0], v_act[1], v_act[2], pot_act,
                    dt_act);

            inf_event[j_act] = 1;
            }
        } else {
            if (inf_event[j_act] == 1) {
                fprintf(out,"INF2 %.6E %.16E %07d %07d   %.6E % .6E % .6E % .6E % .6E % .6E % .6E   %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E  %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E   %.6E   %.6E % .6E % .6E % .6E % .6E % .6E % .6E % .6E  %.6E \n",
                    timesteps, time_cur, i, ind_act[i],
                    sqrt(DR2), x_bbhc[0], x_bbhc[1], x_bbhc[2], v_bbhc[0], v_bbhc[1], v_bbhc[2],
                    m_bh1, x_bh1[0], x_bh1[1], x_bh1[2], v_bh1[0], v_bh1[1], v_bh1[2], pot_bh1,
                    m_bh2, x_bh2[0], x_bh2[1], x_bh2[2], v_bh2[0], v_bh2[1], v_bh2[2], pot_bh2,
                    sqrt(tmp_r2),
                    m_act, x_act[0], x_act[1], x_act[2], v_act[0], v_act[1], v_act[2], pot_act,
                    dt_act);
            }
            inf_event[j_act] = 0;
        } /* if (tmp_r2 < DR2*R_INF2) */
    } /* i */
    fflush(out);
}